xref: /aosp_15_r20/external/federated-compute/fcp/secagg/shared/aes_ctr_prng.cc (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  *     https://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/shared/aes_ctr_prng.h"
18 
19 #include <algorithm>
20 #include <cstdint>
21 #include <string>
22 
23 #include "fcp/base/monitoring.h"
24 #include "fcp/secagg/shared/aes_key.h"
25 #include "fcp/secagg/shared/prng.h"
26 #include "openssl/cipher.h"
27 #include "openssl/evp.h"
28 
29 namespace fcp {
30 namespace secagg {
31 
AesCtrPrng(const AesKey & seed)32 AesCtrPrng::AesCtrPrng(const AesKey& seed) {
33   uint8_t iv[kIvSize];
34   memset(iv, 0, kIvSize);
35   FCP_CHECK(ctx_ = EVP_CIPHER_CTX_new());
36 
37   FCP_CHECK(1 == EVP_EncryptInit_ex(ctx_, EVP_aes_256_ctr(), nullptr,
38                                     seed.data(), iv));
39 
40   // Initializing these to one past the end, in order to force a call to
41   // GenerateBytes on the first attempt to use each cache.
42   next_byte_pos_ = kCacheSize;
43   blocks_generated_ = 0;
44 }
45 
~AesCtrPrng()46 AesCtrPrng::~AesCtrPrng() { EVP_CIPHER_CTX_free(ctx_); }
47 
GenerateBytes(uint8_t * cache,int cache_size)48 void AesCtrPrng::GenerateBytes(uint8_t* cache, int cache_size) {
49   FCP_CHECK((cache_size % kBlockSize) == 0)
50       << "Number of bytes generated by AesCtrPrng must be a multiple of "
51       << kBlockSize;
52   FCP_CHECK(cache_size <= kCacheSize)
53       << "Requested number of bytes " << cache_size
54       << " exceeds maximum cache size " << kCacheSize;
55   FCP_CHECK(blocks_generated_ <= kMaxBlocks)
56       << "AesCtrPrng generated " << kMaxBlocks
57       << " blocks and needs a new seed.";
58   int bytes_written;
59   FCP_CHECK(
60       EVP_EncryptUpdate(ctx_, cache, &bytes_written, kAllZeroes, cache_size));
61   FCP_CHECK(bytes_written == cache_size);
62   blocks_generated_ += static_cast<size_t>(cache_size) / kBlockSize;
63 }
64 
Rand8()65 uint8_t AesCtrPrng::Rand8() {
66   if (next_byte_pos_ >= kCacheSize) {
67     GenerateBytes(cache_, kCacheSize);
68     next_byte_pos_ = 0;
69   }
70   // Return the next byte and then increment the position.
71   return cache_[next_byte_pos_++];
72 }
73 
Rand64()74 uint64_t AesCtrPrng::Rand64() {
75   uint64_t output = 0;
76   for (size_t i = 0; i < sizeof(uint64_t); ++i) {
77     output |= static_cast<uint64_t>(Rand8()) << 8 * i;
78   }
79   return output;
80 }
81 
RandBuffer(uint8_t * buffer,int buffer_size)82 int AesCtrPrng::RandBuffer(uint8_t* buffer, int buffer_size) {
83   buffer_size = std::min(buffer_size, kCacheSize);
84   GenerateBytes(buffer, buffer_size);
85   return buffer_size;
86 }
87 
88 }  // namespace secagg
89 }  // namespace fcp
90