xref: /aosp_15_r20/external/federated-compute/fcp/secagg/server/secagg_server_r1_share_keys_state.cc (revision 14675a029014e728ec732f129a32e299b2da0601)
1 /*
2  * Copyright 2019 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 #include "fcp/secagg/server/secagg_server_r1_share_keys_state.h"
18 
19 #include <algorithm>
20 #include <memory>
21 #include <string>
22 #include <utility>
23 
24 #include "fcp/base/monitoring.h"
25 #include "fcp/secagg/server/secagg_server_r2_masked_input_coll_state.h"
26 
27 namespace fcp {
28 namespace secagg {
29 
SecAggServerR1ShareKeysState(std::unique_ptr<SecAggServerProtocolImpl> impl,int number_of_clients_failed_after_sending_masked_input,int number_of_clients_failed_before_sending_masked_input,int number_of_clients_terminated_without_unmasking)30 SecAggServerR1ShareKeysState::SecAggServerR1ShareKeysState(
31     std::unique_ptr<SecAggServerProtocolImpl> impl,
32     int number_of_clients_failed_after_sending_masked_input,
33     int number_of_clients_failed_before_sending_masked_input,
34     int number_of_clients_terminated_without_unmasking)
35     : SecAggServerState(number_of_clients_failed_after_sending_masked_input,
36                         number_of_clients_failed_before_sending_masked_input,
37                         number_of_clients_terminated_without_unmasking,
38                         SecAggServerStateKind::R1_SHARE_KEYS, std::move(impl)) {
39 }
40 
~SecAggServerR1ShareKeysState()41 SecAggServerR1ShareKeysState::~SecAggServerR1ShareKeysState() {}
42 
HandleMessage(uint32_t client_id,const ClientToServerWrapperMessage & message)43 Status SecAggServerR1ShareKeysState::HandleMessage(
44     uint32_t client_id, const ClientToServerWrapperMessage& message) {
45   if (message.has_abort()) {
46     MessageReceived(message, false);
47     AbortClient(client_id, "", ClientDropReason::SENT_ABORT_MESSAGE,
48                 /*notify=*/false);
49     return FCP_STATUS(OK);
50   }
51   // If the client has aborted or sent a message already, ignore its messages.
52   if (client_status(client_id) != ClientStatus::ADVERTISE_KEYS_RECEIVED) {
53     MessageReceived(message, false);
54     AbortClient(client_id,
55                 "Not expecting an ShareKeysResponse from this "
56                 "client - either the client already aborted or one such "
57                 "message was already received.",
58                 ClientDropReason::SHARE_KEYS_UNEXPECTED);
59     return FCP_STATUS(OK);
60   }
61   if (!message.has_share_keys_response()) {
62     MessageReceived(message, false);
63     AbortClient(client_id,
64                 "Message type received is different from what was expected.",
65                 ClientDropReason::UNEXPECTED_MESSAGE_TYPE);
66     return FCP_STATUS(OK);
67   }
68   MessageReceived(message, true);
69 
70   Status status =
71       impl()->HandleShareKeysResponse(client_id, message.share_keys_response());
72   if (!status.ok()) {
73     AbortClient(client_id, std::string(status.message()),
74                 ClientDropReason::INVALID_SHARE_KEYS_RESPONSE);
75     return FCP_STATUS(OK);
76   }
77 
78   set_client_status(client_id, ClientStatus::SHARE_KEYS_RECEIVED);
79   number_of_messages_received_in_this_round_++;
80   number_of_clients_ready_for_next_round_++;
81   return FCP_STATUS(OK);
82 }
83 
IsNumberOfIncludedInputsCommitted() const84 bool SecAggServerR1ShareKeysState::IsNumberOfIncludedInputsCommitted() const {
85   return false;
86 }
87 
MinimumMessagesNeededForNextRound() const88 int SecAggServerR1ShareKeysState::MinimumMessagesNeededForNextRound() const {
89   return std::max(0, minimum_number_of_clients_to_proceed() -
90                          number_of_clients_ready_for_next_round_);
91 }
92 
NumberOfPendingClients() const93 int SecAggServerR1ShareKeysState::NumberOfPendingClients() const {
94   return NumberOfAliveClients() - number_of_clients_ready_for_next_round_;
95 }
96 
HandleAbortClient(uint32_t client_id,ClientDropReason reason_code)97 void SecAggServerR1ShareKeysState::HandleAbortClient(
98     uint32_t client_id, ClientDropReason reason_code) {
99   number_of_clients_failed_before_sending_masked_input_++;
100   if (client_status(client_id) == ClientStatus::SHARE_KEYS_RECEIVED) {
101     number_of_clients_ready_for_next_round_--;
102     // Remove that client's shared keys as if they were never sent. This will
103     // avoid wasted computation on both client and server ends.
104     impl()->EraseShareKeysForClient(client_id);
105   }
106   set_client_status(client_id,
107                     ClientStatus::DEAD_AFTER_ADVERTISE_KEYS_RECEIVED);
108   if (NumberOfAliveClients() < minimum_number_of_clients_to_proceed()) {
109     needs_to_abort_ = true;
110   }
111 }
112 
113 StatusOr<std::unique_ptr<SecAggServerState>>
ProceedToNextRound()114 SecAggServerR1ShareKeysState::ProceedToNextRound() {
115   if (!ReadyForNextRound()) {
116     return FCP_STATUS(UNAVAILABLE);
117   }
118   if (needs_to_abort_) {
119     std::string error_string = "Too many clients aborted.";
120     ServerToClientWrapperMessage message;
121     message.mutable_abort()->set_diagnostic_info(error_string);
122     message.mutable_abort()->set_early_success(false);
123     SendBroadcast(message);
124 
125     return AbortState(error_string,
126                       SecAggServerOutcome::NOT_ENOUGH_CLIENTS_REMAINING);
127   }
128 
129   // Abort all clients that haven't yet sent a message, and send a message to
130   // all clients that are still alive.
131   for (int i = 0; i < total_number_of_clients(); ++i) {
132     if (!IsClientDead(i) &&
133         client_status(i) != ClientStatus::SHARE_KEYS_RECEIVED) {
134       AbortClient(
135           i, "Client did not send ShareKeysResponse before round transition.",
136           ClientDropReason::NO_SHARE_KEYS);
137     } else if (client_status(i) == ClientStatus::SHARE_KEYS_RECEIVED) {
138       ServerToClientWrapperMessage message;
139       impl()->PrepareMaskedInputCollectionRequestForClient(
140           i, message.mutable_masked_input_request());
141       Send(i, message);
142     }
143   }
144 
145   // Encrypted shares are no longer needed beyond this point as the server has
146   // already forwarded them to the clients.
147   impl()->ClearShareKeys();
148 
149   return std::make_unique<SecAggServerR2MaskedInputCollState>(
150       ExitState(StateTransition::kSuccess),
151       number_of_clients_failed_after_sending_masked_input_,
152       number_of_clients_failed_before_sending_masked_input_,
153       number_of_clients_terminated_without_unmasking_);
154 }
155 
ReadyForNextRound() const156 bool SecAggServerR1ShareKeysState::ReadyForNextRound() const {
157   return (number_of_clients_ready_for_next_round_ >=
158           minimum_number_of_clients_to_proceed()) ||
159          (needs_to_abort_);
160 }
161 
162 }  // namespace secagg
163 }  // namespace fcp
164