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_r0_advertise_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_r1_share_keys_state.h"
26 
27 namespace fcp {
28 namespace secagg {
29 
SecAggServerR0AdvertiseKeysState(std::unique_ptr<SecAggServerProtocolImpl> impl)30 SecAggServerR0AdvertiseKeysState::SecAggServerR0AdvertiseKeysState(
31     std::unique_ptr<SecAggServerProtocolImpl> impl)
32     : SecAggServerState(0, 0, 0, SecAggServerStateKind::R0_ADVERTISE_KEYS,
33                         std::move(impl)) {
34   if (metrics()) {
35     // This is the initial state, so we count the start of the protocol from the
36     // moment it is constructed.
37     metrics()->ProtocolStarts(this->impl()->server_variant());
38   }
39 }
40 
~SecAggServerR0AdvertiseKeysState()41 SecAggServerR0AdvertiseKeysState::~SecAggServerR0AdvertiseKeysState() {}
42 
HandleMessage(uint32_t client_id,const ClientToServerWrapperMessage & message)43 Status SecAggServerR0AdvertiseKeysState::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::READY_TO_START) {
53     MessageReceived(message, false);
54     AbortClient(
55         client_id,
56         "Not expecting an AdvertiseKeys message from this client - either the "
57         "client already aborted or one such message was already received.",
58         ClientDropReason::ADVERTISE_KEYS_UNEXPECTED);
59     return FCP_STATUS(OK);
60   }
61   if (!message.has_advertise_keys()) {
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()->HandleAdvertiseKeys(client_id, message.advertise_keys());
72   if (!status.ok()) {
73     AbortClient(client_id, std::string(status.message()),
74                 ClientDropReason::INVALID_PUBLIC_KEY);
75     return FCP_STATUS(OK);
76   }
77 
78   set_client_status(client_id, ClientStatus::ADVERTISE_KEYS_RECEIVED);
79   number_of_clients_ready_for_next_round_++;
80   number_of_messages_received_in_this_round_++;
81 
82   return FCP_STATUS(OK);
83 }
84 
IsNumberOfIncludedInputsCommitted() const85 bool SecAggServerR0AdvertiseKeysState::IsNumberOfIncludedInputsCommitted()
86     const {
87   return false;
88 }
89 
MinimumMessagesNeededForNextRound() const90 int SecAggServerR0AdvertiseKeysState::MinimumMessagesNeededForNextRound()
91     const {
92   return std::max(0, minimum_number_of_clients_to_proceed() -
93                          number_of_clients_ready_for_next_round_);
94 }
95 
NumberOfPendingClients() const96 int SecAggServerR0AdvertiseKeysState::NumberOfPendingClients() const {
97   return NumberOfAliveClients() - number_of_clients_ready_for_next_round_;
98 }
99 
HandleAbortClient(uint32_t client_id,ClientDropReason reason_code)100 void SecAggServerR0AdvertiseKeysState::HandleAbortClient(
101     uint32_t client_id, ClientDropReason reason_code) {
102   number_of_clients_failed_before_sending_masked_input_++;
103   if (client_status(client_id) == ClientStatus::ADVERTISE_KEYS_RECEIVED) {
104     number_of_clients_ready_for_next_round_--;
105     // Remove that client's public keys as if they were never sent. This will
106     // avoid wasted computation and bandwidth.
107     impl()->ErasePublicKeysForClient(client_id);
108   }
109   set_client_status(client_id, ClientStatus::DEAD_BEFORE_SENDING_ANYTHING);
110   if (NumberOfAliveClients() < minimum_number_of_clients_to_proceed()) {
111     needs_to_abort_ = true;
112   }
113 }
114 
115 StatusOr<std::unique_ptr<SecAggServerState>>
ProceedToNextRound()116 SecAggServerR0AdvertiseKeysState::ProceedToNextRound() {
117   if (!ReadyForNextRound()) {
118     return FCP_STATUS(UNAVAILABLE);
119   }
120   if (needs_to_abort_) {
121     std::string error_string = "Too many clients aborted.";
122     ServerToClientWrapperMessage abort_message;
123     abort_message.mutable_abort()->set_diagnostic_info(error_string);
124     abort_message.mutable_abort()->set_early_success(false);
125     SendBroadcast(abort_message);
126 
127     return AbortState(error_string,
128                       SecAggServerOutcome::NOT_ENOUGH_CLIENTS_REMAINING);
129   }
130 
131   // Abort all clients that haven't yet sent a message.
132   for (int i = 0; i < total_number_of_clients(); ++i) {
133     if (!IsClientDead(i) &&
134         client_status(i) != ClientStatus::ADVERTISE_KEYS_RECEIVED) {
135       AbortClient(
136           i,
137           "Client did not send AdvertiseKeys message before round transition.",
138           ClientDropReason::NO_ADVERTISE_KEYS);
139     }
140   }
141 
142   impl()->ComputeSessionId();
143 
144   ServerToClientWrapperMessage message_to_client;
145   message_to_client.mutable_share_keys_request()->set_session_id(
146       impl()->session_id().data);
147   FCP_RETURN_IF_ERROR(impl()->InitializeShareKeysRequest(
148       message_to_client.mutable_share_keys_request()));
149 
150   for (int i = 0; i < total_number_of_clients(); ++i) {
151     //  Reuse the common parts of the ShareKeysRequest message and update the
152     //  client-specific parts.
153     if (!IsClientDead(i)) {
154       impl()->PrepareShareKeysRequestForClient(
155           i, message_to_client.mutable_share_keys_request());
156       Send(i, message_to_client);
157     }
158   }
159 
160   // Pairs of public keys are no longer needed beyond this point as the server
161   // has already forwarded them to the clients.
162   impl()->ClearPairsOfPublicKeys();
163 
164   return {std::make_unique<SecAggServerR1ShareKeysState>(
165       ExitState(StateTransition::kSuccess),
166       number_of_clients_failed_after_sending_masked_input_,
167       number_of_clients_failed_before_sending_masked_input_,
168       number_of_clients_terminated_without_unmasking_)};
169 }
170 
ReadyForNextRound() const171 bool SecAggServerR0AdvertiseKeysState::ReadyForNextRound() const {
172   return (number_of_clients_ready_for_next_round_ >=
173           minimum_number_of_clients_to_proceed()) ||
174          (needs_to_abort_);
175 }
176 
177 }  // namespace secagg
178 }  // namespace fcp
179