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_SERVER_SECAGG_SERVER_PRNG_RUNNING_STATE_H_ 18 #define FCP_SECAGG_SERVER_SECAGG_SERVER_PRNG_RUNNING_STATE_H_ 19 20 #include <functional> 21 #include <memory> 22 #include <optional> 23 24 #include "absl/base/thread_annotations.h" 25 #include "absl/synchronization/mutex.h" 26 #include "absl/time/time.h" 27 #include "fcp/secagg/server/secagg_server_state.h" 28 29 namespace fcp { 30 namespace secagg { 31 32 // This class is the State for the SecAggServer when it has collected all secret 33 // shares from the clients and is ready to compute its final output. The 34 // protocol is essentially done, but this is a separate state from 35 // SecAggClientCompletedState because there the server still needs to run the 36 // potentially expensive step of using the PRNG to stretch client keys into 37 // masking vectors. 38 39 class SecAggServerPrngRunningState final : public SecAggServerState { 40 public: 41 SecAggServerPrngRunningState( 42 std::unique_ptr<SecAggServerProtocolImpl> impl, 43 int number_of_clients_failed_after_sending_masked_input, 44 int number_of_clients_failed_before_sending_masked_input, 45 int number_of_clients_terminated_without_unmasking); 46 47 ~SecAggServerPrngRunningState() override; 48 49 void EnterState() override; 50 51 // Handles abort message from a client. Any other type of message is 52 // unexpected and results in the client being aborted. 53 Status HandleMessage(uint32_t client_id, 54 const ClientToServerWrapperMessage& message) override; 55 56 bool IsNumberOfIncludedInputsCommitted() const override; 57 58 int NumberOfIncludedInputs() const override; 59 60 StatusOr<std::unique_ptr<SecAggServerState> > ProceedToNextRound() override; 61 62 bool ReadyForNextRound() const override; 63 64 bool SetAsyncCallback(std::function<void()> async_callback) override; 65 66 private: 67 void HandleAbort() override; 68 69 void HandleAbortClient(uint32_t client_id, 70 ClientDropReason reason_code) override; 71 72 // Called to perform the initial synchronous part of PRNG state. 73 StatusOr<SecAggServerProtocolImpl::PrngWorkItems> Initialize(); 74 75 // This is called when all computations are finished. 76 // final_status indicates whether PRNG computation has finished successfully. 77 void PrngRunnerFinished(Status final_status); 78 79 // The status is assigned when the state completes either successfully or 80 // unsuccessfully. 81 std::optional<Status> completion_status_ ABSL_GUARDED_BY(mutex_); 82 83 absl::Time prng_started_time_; 84 CancellationToken cancellation_token_; 85 86 std::function<void()> prng_done_callback_ ABSL_GUARDED_BY(mutex_); 87 88 // Protects this object from being destroyed while StartPrng call is still 89 // in progress. Also protects completion_status_ and prng_done_callback_. 90 mutable absl::Mutex mutex_; 91 }; 92 93 } // namespace secagg 94 } // namespace fcp 95 96 #endif // FCP_SECAGG_SERVER_SECAGG_SERVER_PRNG_RUNNING_STATE_H_ 97