xref: /aosp_15_r20/external/federated-compute/fcp/secagg/shared/aes_ctr_prng_test.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 <cstdint>
20 #include <vector>
21 
22 #include "gmock/gmock.h"
23 #include "gtest/gtest.h"
24 #include "fcp/secagg/shared/aes_ctr_prng_factory.h"
25 
26 namespace fcp {
27 namespace secagg {
28 namespace {
29 
30 using ::testing::Contains;
31 using ::testing::Eq;
32 using ::testing::Ne;
33 using ::testing::Not;
34 using ::testing::Pointwise;
35 
TEST(AesCtrPrngTest,Rand8ReturnsSameValuesGivenSameInputs)36 TEST(AesCtrPrngTest, Rand8ReturnsSameValuesGivenSameInputs) {
37   uint8_t seed_data[32];
38   memset(seed_data, '1', 32);
39   AesKey seed(seed_data);
40 
41   AesCtrPrngFactory factory;
42   std::unique_ptr<SecurePrng> prng0 = factory.MakePrng(seed);
43   std::unique_ptr<SecurePrng> prng1 = factory.MakePrng(seed);
44   std::vector<uint8_t> output0;
45   std::vector<uint8_t> output1;
46   for (int i = 0; i < 16; ++i) {
47     output0.push_back(prng0->Rand8());
48     output1.push_back(prng1->Rand8());
49   }
50   EXPECT_THAT(output0, Eq(output1));
51 }
52 
TEST(AesCtrPrngTest,Rand64ReturnsSameValuesGivenSameInputs)53 TEST(AesCtrPrngTest, Rand64ReturnsSameValuesGivenSameInputs) {
54   uint8_t seed_data[32];
55   memset(seed_data, '1', 32);
56   AesKey seed(seed_data);
57 
58   AesCtrPrngFactory factory;
59   std::unique_ptr<SecurePrng> prng0 = factory.MakePrng(seed);
60   std::unique_ptr<SecurePrng> prng1 = factory.MakePrng(seed);
61   std::vector<uint64_t> output0;
62   std::vector<uint64_t> output1;
63   for (int i = 0; i < 16; ++i) {
64     output0.push_back(prng0->Rand64());
65     output1.push_back(prng1->Rand64());
66   }
67   EXPECT_THAT(output0, Eq(output1));
68 }
69 
TEST(AesCtrPrngTest,MixedRandCallsReturnSameValuesGivenSameInputs)70 TEST(AesCtrPrngTest, MixedRandCallsReturnSameValuesGivenSameInputs) {
71   uint8_t seed_data[32];
72   memset(seed_data, '1', 32);
73   AesKey seed(seed_data);
74 
75   AesCtrPrngFactory factory;
76   std::unique_ptr<SecurePrng> prng0 = factory.MakePrng(seed);
77   std::unique_ptr<SecurePrng> prng1 = factory.MakePrng(seed);
78   std::vector<uint64_t> output0;
79   std::vector<uint64_t> output1;
80   for (int i = 0; i < 5; ++i) {
81     output0.push_back(prng0->Rand8());
82     output1.push_back(prng1->Rand8());
83   }
84   for (int i = 0; i < 5; ++i) {
85     output0.push_back(prng0->Rand64());
86     output1.push_back(prng1->Rand64());
87   }
88   for (int i = 0; i < 10; ++i) {
89     output0.push_back(prng0->Rand8());
90     output1.push_back(prng1->Rand8());
91   }
92   EXPECT_THAT(output0, Eq(output1));
93 }
94 
95 // While for random seeds or IVs there would be a very small chance of
96 // duplication, these tests are not flaky because this PRNG is deterministic.
TEST(AesCtrPrngTest,DifferentSeedsGenerateDifferentValues)97 TEST(AesCtrPrngTest, DifferentSeedsGenerateDifferentValues) {
98   uint8_t seed_data[32];
99   memset(seed_data, '1', 32);
100   AesKey seed1(seed_data);
101   memset(seed_data, '3', 32);
102   AesKey seed2(seed_data);
103 
104   AesCtrPrngFactory factory;
105   std::unique_ptr<SecurePrng> prng0 = factory.MakePrng(seed1);
106   std::unique_ptr<SecurePrng> prng1 = factory.MakePrng(seed2);
107   std::vector<uint64_t> output0;
108   std::vector<uint64_t> output1;
109   for (int i = 0; i < 16; ++i) {
110     output0.push_back(prng0->Rand64());
111     output1.push_back(prng1->Rand64());
112   }
113   // output0 differs from output1 at every point
114   EXPECT_THAT(output0, Pointwise(Ne(), output1));
115 }
116 
TEST(AesCtrPrngTest,DoesntGenerateRepeatedValues)117 TEST(AesCtrPrngTest, DoesntGenerateRepeatedValues) {
118   uint8_t seed_data[32];
119   memset(seed_data, '1', 32);
120   AesKey seed(seed_data);
121 
122   AesCtrPrngFactory factory;
123   std::unique_ptr<SecurePrng> prng = factory.MakePrng(seed);
124   std::vector<uint64_t> output;
125   uint64_t val;
126   for (int i = 0; i < 16; ++i) {
127     val = prng->Rand64();
128     EXPECT_THAT(output, Not(Contains(val)));
129     output.push_back(val);
130   }
131 }
132 
TEST(AesCtrPrngTest,GeneratesExpectedValues)133 TEST(AesCtrPrngTest, GeneratesExpectedValues) {
134   uint8_t iv[16];
135   memset(iv, 0, sizeof(iv));
136 
137   uint8_t seed_data[32];
138   memset(seed_data, '1', sizeof(seed_data));
139   AesKey seed(seed_data);
140 
141   EVP_CIPHER_CTX* ctx;
142   ctx = EVP_CIPHER_CTX_new();
143   ASSERT_THAT(ctx, Ne(nullptr));
144 
145   ASSERT_THAT(
146       EVP_EncryptInit_ex(ctx, EVP_aes_256_ctr(), nullptr, seed_data, iv),
147       Eq(1));
148 
149   const int kBlockSize = 16 * 32;
150 
151   static constexpr uint8_t zeroes[kBlockSize] = {0};
152 
153   // These are processed separately in the class
154   uint8_t expected_uint8_t[kBlockSize];
155   uint8_t expected_uint64_t[kBlockSize];
156 
157   // Obtain the ciphertext incrementally to verify identical output of versions
158   // using a different block size.
159   int len;
160   for (auto i = 0; i < kBlockSize; i += 16) {
161     ASSERT_THAT(EVP_EncryptUpdate(ctx, &expected_uint8_t[i], &len, zeroes, 16),
162                 Ne(0));
163     ASSERT_THAT(len, 16);
164   }
165   for (auto i = 0; i < kBlockSize; i += 16) {
166     ASSERT_THAT(EVP_EncryptUpdate(ctx, &expected_uint64_t[i], &len, zeroes, 16),
167                 Ne(0));
168     ASSERT_THAT(len, 16);
169   }
170 
171   AesCtrPrngFactory factory;
172   std::unique_ptr<SecurePrng> prng = factory.MakePrng(seed);
173 
174   for (int i = 0; i < sizeof(expected_uint8_t); i++) {
175     EXPECT_THAT(prng->Rand8(), Eq(expected_uint8_t[i]));
176   }
177   for (int i = 0; i < sizeof(expected_uint64_t) / sizeof(uint64_t); i++) {
178     uint64_t value = 0;
179     for (int j = 0; j < sizeof(uint64_t); j++) {
180       value |=
181           static_cast<uint64_t>(expected_uint64_t[i * sizeof(uint64_t) + j])
182           << (8 * j);
183     }
184     EXPECT_THAT(prng->Rand64(), Eq(value));
185   }
186   EVP_CIPHER_CTX_free(ctx);
187 }
188 
TEST(AesCtrPrngTest,RandBufferIsConsistentWithRand8)189 TEST(AesCtrPrngTest, RandBufferIsConsistentWithRand8) {
190   uint8_t seed_data[32];
191   memset(seed_data, '1', 32);
192   AesKey seed(seed_data);
193 
194   AesCtrPrngFactory factory1;
195   AesCtrPrngFactory factory2;
196   std::unique_ptr<SecurePrng> prng1 = factory1.MakePrng(seed);
197   std::unique_ptr<SecurePrng> prng2 = factory2.MakePrng(seed);
198   auto batch_prng = static_cast<SecureBatchPrng*>(prng2.get());
199 
200   constexpr int kSize = 16000;
201   std::vector<uint8_t> output1(kSize);
202   std::vector<uint8_t> output2(kSize);
203 
204   // Fill output1 using Rand8
205   for (int i = 0; i < kSize; ++i) {
206     output1[i] = prng1->Rand8();
207   }
208 
209   // Fill output2 using RandBuffer
210   int bytes_received = 0;
211   while (bytes_received < kSize) {
212     bytes_received += batch_prng->RandBuffer(output2.data() + bytes_received,
213                                              kSize - bytes_received);
214   }
215 
216   // output1 and output2 should be the same.
217   EXPECT_THAT(output1, Eq(output2));
218 }
219 
220 }  // namespace
221 }  // namespace secagg
222 }  // namespace fcp
223