xref: /aosp_15_r20/external/federated-compute/fcp/secagg/client/secagg_client.h (revision 14675a029014e728ec732f129a32e299b2da0601)
1 /*
2  * Copyright 2018 Google LLC
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #ifndef FCP_SECAGG_CLIENT_SECAGG_CLIENT_H_
18 #define FCP_SECAGG_CLIENT_SECAGG_CLIENT_H_
19 
20 #include <memory>
21 #include <string>
22 #include <vector>
23 
24 #include "absl/synchronization/mutex.h"
25 #include "fcp/base/monitoring.h"
26 #include "fcp/secagg/client/secagg_client_state.h"
27 #include "fcp/secagg/client/send_to_server_interface.h"
28 #include "fcp/secagg/client/state_transition_listener_interface.h"
29 #include "fcp/secagg/shared/aes_prng_factory.h"
30 #include "fcp/secagg/shared/async_abort.h"
31 #include "fcp/secagg/shared/input_vector_specification.h"
32 #include "fcp/secagg/shared/prng.h"
33 #include "fcp/secagg/shared/secagg_messages.pb.h"
34 #include "fcp/secagg/shared/secagg_vector.h"
35 
36 namespace fcp {
37 namespace secagg {
38 
39 // Represents a client for the secure aggregation protocol. Each instance of
40 // this class performs just *one* session of the protocol.
41 //
42 // To create a new instance, use the public constructor. The Start method can be
43 // used to produce the first message of the protocol, the SetInput method sets
44 // the input for this client and the ReceiveMessage method is used to process
45 // incoming messages from the server.
46 //
47 // The class is thread-safe, but will deadlock if accessed reentrantly from
48 // the SendToServerInterface callback.
49 //
50 // Functions are marked virtual for mockability.  Additional virtual attributes
51 // should be added as needed by tests.
52 
53 class SecAggClient {
54  public:
55   // Creates a new instance of the client.
56   //
57   // max_neighbors_expected is the upper bound on the total number of neighbors
58   // this client may interact with. If the server tries to start a protocol
59   // session with more than this many neighbors, this client will abort.
60   //
61   // minimum_surviving_neighbors_for_reconstruction is the threshold lower bound
62   // on the number of neighbors participating. If there are ever fewer than this
63   // number of remaining neighbors in the protocol, this client will abort.
64   //
65   // input_vector_specs must contain one InputVectorSpecification for each input
66   // vector which the protocol will aggregate.  This may optionally be moved
67   // from using std::move(caller_input_vector_specs).
68   //
69   // prng should always be an instance of CryptoRandPrng, except as needed for
70   // testing purposes. The client will consume prng, taking ownership of it.
71   //
72   // sender is used by the client to send messages to the server. The client
73   // will consume sender, taking ownership of it.
74   //
75   // transition_listener is used to trigger state transition events, used for
76   // logging.
77   //
78   // prng_factory is a pointer to an instance of a subclass of AesPrngFactory.
79   // The type of prng_factory must be consistent with the one used on the
80   // server.
81   //
82   // async_abort_for_test, optionally, allows the caller to reset the abort
83   // signal. This is used to exhaustively test all abort paths, and should not
84   // be used in production; specifically, if this paramter is not nullptr,
85   // Abort() will no longer abort a state-in-progress; it will only abort across
86   // state transitions.
87   SecAggClient(
88       int max_neighbors_expected,
89       int minimum_surviving_neighbors_for_reconstruction,
90       std::vector<InputVectorSpecification> input_vector_specs,
91       std::unique_ptr<SecurePrng> prng,
92       std::unique_ptr<SendToServerInterface> sender,
93       std::unique_ptr<StateTransitionListenerInterface> transition_listener,
94       std::unique_ptr<AesPrngFactory> prng_factory,
95       std::atomic<std::string*>* abort_signal_for_test = nullptr);
96   virtual ~SecAggClient() = default;
97 
98   // Disallow copy and move.
99   SecAggClient(const SecAggClient&) = delete;
100   SecAggClient& operator=(const SecAggClient&) = delete;
101 
102   // Initiates the protocol by computing its first message and sending it to
103   // the server. This method should only be called once. The output will be OK
104   // unless it is called more than once.
105   virtual Status Start();
106 
107   // Makes this client abort the protocol and sends a message to notify the
108   // server. All the state is erased. A new instance of SecAggClient will have
109   // to be created to restart the protocol.
110   //
111   // The status will be OK unless the protocol was already completed or aborted.
112   Status Abort();
113 
114   // Makes this client abort the protocol and sends a message to notify the
115   // server. All the state is erased. A new instance of SecAggClient will have
116   // to be created to restart the protocol.
117   //
118   // The specified reason for aborting will be sent to the server and logged.
119   //
120   // The status will be OK unless the protocol was already completed or aborted.
121   Status Abort(const std::string& reason);
122 
123   // Sets the input of this client for this protocol session. This method should
124   // only be called once.
125   //
126   // If the input does not match the format laid out in input_vector_specs,
127   // this will return INVALID_ARGUMENT. If SetInput has already been called or
128   // if the client is in an aborted or completed state, this will return
129   // FAILED_PRECONDITION. Otherwise returns OK.
130   Status SetInput(std::unique_ptr<SecAggVectorMap> input_map);
131 
132   // Returns a string uniquely describing the current state of the client's FSM.
133   ABSL_MUST_USE_RESULT std::string State() const;
134 
135   // Returns true if the client has aborted the protocol, false else.
136   ABSL_MUST_USE_RESULT bool IsAborted() const;
137 
138   // Returns true if the client has successfully completed the protocol,
139   // false else.
140   ABSL_MUST_USE_RESULT bool IsCompletedSuccessfully() const;
141 
142   // Returns a string describing the reason that the client aborted.
143   // If the client has not actually aborted, returns an error Status with code
144   // PRECONDITION_FAILED.
145   ABSL_MUST_USE_RESULT StatusOr<std::string> ErrorMessage() const;
146 
147   // Used to process an incoming message from the server. This method uses the
148   // SendToServerInterface passed to the constructor to send the response
149   // directly to the server.
150   //
151   // The output will be true if the client is still active, or false if the
152   // client is now in a terminal state. The output will be a failure status if
153   // the client did not process the message because it was in a terminal state,
154   // or because the message was the wrong type.
155   StatusOr<bool> ReceiveMessage(const ServerToClientWrapperMessage& incoming);
156 
157  private:
158   mutable absl::Mutex mu_;
159 
160   std::atomic<std::string*> abort_signal_;
161   AsyncAbort async_abort_;
162 
163   // The internal State object, containing details about this client's current
164   // state.
165   std::unique_ptr<SecAggClientState> state_ ABSL_GUARDED_BY(mu_);
166 };
167 
168 }  // namespace secagg
169 }  // namespace fcp
170 
171 #endif  // FCP_SECAGG_CLIENT_SECAGG_CLIENT_H_
172