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