1 // Copyright 2023 Google LLC
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 //
15 ///////////////////////////////////////////////////////////////////////////////
16
17 #include "tink/internal/keyset_wrapper_store.h"
18
19 #include <functional>
20 #include <memory>
21 #include <string>
22 #include <utility>
23
24 #include "gmock/gmock.h"
25 #include "gtest/gtest.h"
26 #include "absl/status/status.h"
27 #include "tink/internal/registry_impl.h"
28 #include "tink/mac/mac_wrapper.h"
29 #include "tink/primitive_set.h"
30 #include "tink/primitive_wrapper.h"
31 #include "tink/subtle/random.h"
32 #include "tink/util/status.h"
33 #include "tink/util/statusor.h"
34 #include "tink/util/test_matchers.h"
35 #include "tink/util/test_util.h"
36 #include "proto/aes_gcm.pb.h"
37
38 namespace crypto {
39 namespace tink {
40 namespace internal {
41 namespace {
42
43 using ::crypto::tink::test::IsOk;
44 using ::crypto::tink::test::StatusIs;
45 using ::google::crypto::tink::AesGcmKey;
46 using ::google::crypto::tink::AesGcmKeyFormat;
47 using ::google::crypto::tink::KeyData;
48 using ::google::crypto::tink::Keyset;
49 using ::google::crypto::tink::KeysetInfo;
50 using ::google::crypto::tink::KeyStatusType;
51 using ::google::crypto::tink::OutputPrefixType;
52 using ::testing::Eq;
53
54 class FakePrimitive {
55 public:
FakePrimitive(std::string s)56 explicit FakePrimitive(std::string s) : s_(s) {}
get()57 std::string get() { return s_; }
58
59 private:
60 std::string s_;
61 };
62
63 class FakeKeyTypeManager
64 : public KeyTypeManager<AesGcmKey, AesGcmKeyFormat, List<FakePrimitive>> {
65 public:
66 class FakePrimitiveFactory : public PrimitiveFactory<FakePrimitive> {
67 public:
Create(const AesGcmKey & key) const68 util::StatusOr<std::unique_ptr<FakePrimitive>> Create(
69 const AesGcmKey& key) const override {
70 return absl::make_unique<FakePrimitive>(key.key_value());
71 }
72 };
73
FakeKeyTypeManager()74 FakeKeyTypeManager()
75 : KeyTypeManager(absl::make_unique<FakePrimitiveFactory>()) {}
76
key_material_type() const77 KeyData::KeyMaterialType key_material_type() const override {
78 return KeyData::SYMMETRIC;
79 }
80
get_version() const81 uint32_t get_version() const override { return 0; }
82
get_key_type() const83 const std::string& get_key_type() const override { return key_type_; }
84
ValidateKey(const AesGcmKey & key) const85 util::Status ValidateKey(const AesGcmKey& key) const override {
86 return util::OkStatus();
87 }
88
ValidateKeyFormat(const AesGcmKeyFormat & key_format) const89 util::Status ValidateKeyFormat(
90 const AesGcmKeyFormat& key_format) const override {
91 return util::OkStatus();
92 }
93
CreateKey(const AesGcmKeyFormat & key_format) const94 util::StatusOr<AesGcmKey> CreateKey(
95 const AesGcmKeyFormat& key_format) const override {
96 return AesGcmKey();
97 }
98
DeriveKey(const AesGcmKeyFormat & key_format,InputStream * input_stream) const99 util::StatusOr<AesGcmKey> DeriveKey(
100 const AesGcmKeyFormat& key_format,
101 InputStream* input_stream) const override {
102 return AesGcmKey();
103 }
104
105 private:
106 const std::string key_type_ =
107 "type.googleapis.com/google.crypto.tink.AesGcmKey";
108 };
109
110 class FakePrimitiveWrapper
111 : public PrimitiveWrapper<FakePrimitive, FakePrimitive> {
112 public:
Wrap(std::unique_ptr<PrimitiveSet<FakePrimitive>> primitive_set) const113 util::StatusOr<std::unique_ptr<FakePrimitive>> Wrap(
114 std::unique_ptr<PrimitiveSet<FakePrimitive>> primitive_set)
115 const override {
116 return absl::make_unique<FakePrimitive>(
117 primitive_set->get_primary()->get_primitive().get());
118 }
119 };
120
121 class FakePrimitiveWrapper2
122 : public PrimitiveWrapper<FakePrimitive, FakePrimitive> {
123 public:
Wrap(std::unique_ptr<PrimitiveSet<FakePrimitive>> primitive_set) const124 util::StatusOr<std::unique_ptr<FakePrimitive>> Wrap(
125 std::unique_ptr<PrimitiveSet<FakePrimitive>> primitive_set)
126 const override {
127 return absl::make_unique<FakePrimitive>(
128 primitive_set->get_primary()->get_primitive().get());
129 }
130 };
131
AddAesGcmKeyToKeyset(Keyset & keyset,uint32_t key_id,OutputPrefixType output_prefix_type,KeyStatusType key_status_type)132 std::string AddAesGcmKeyToKeyset(Keyset& keyset, uint32_t key_id,
133 OutputPrefixType output_prefix_type,
134 KeyStatusType key_status_type) {
135 AesGcmKey key;
136 key.set_version(0);
137 key.set_key_value(subtle::Random::GetRandomBytes(16));
138 KeyData key_data;
139 key_data.set_value(key.SerializeAsString());
140 key_data.set_type_url("type.googleapis.com/google.crypto.tink.AesGcmKey");
141 test::AddKeyData(key_data, key_id, output_prefix_type, key_status_type,
142 &keyset);
143 return key.key_value();
144 }
145
146 // Returns the function that relies on `registry` to transform `key_data` into
147 // FakePrimitive.
148 util::StatusOr<std::function<
149 util::StatusOr<std::unique_ptr<FakePrimitive>>(const KeyData& key_data)>>
PrimitiveGetter(RegistryImpl & registry)150 PrimitiveGetter(RegistryImpl& registry) {
151 util::Status status =
152 registry.RegisterKeyTypeManager<AesGcmKey, AesGcmKeyFormat,
153 List<FakePrimitive>>(
154 absl::make_unique<FakeKeyTypeManager>(),
155 /*new_key_allowed=*/true);
156 if (!status.ok()) {
157 return status;
158 }
159 return [®istry](const KeyData& key_data) {
160 return registry.GetPrimitive<FakePrimitive>(key_data);
161 };
162 }
163
TEST(KeysetWrapperStoreTest,Add)164 TEST(KeysetWrapperStoreTest, Add) {
165 RegistryImpl registry;
166 util::StatusOr<std::function<util::StatusOr<std::unique_ptr<FakePrimitive>>(
167 const KeyData& key_data)>>
168 primitive_getter = PrimitiveGetter(registry);
169 ASSERT_THAT(primitive_getter, IsOk());
170
171 KeysetWrapperStore store;
172 EXPECT_THAT(
173 (store.Add<FakePrimitive, FakePrimitive>(
174 absl::make_unique<FakePrimitiveWrapper>(), *primitive_getter)),
175 IsOk());
176 }
177
TEST(KeysetWrapperStoreTest,AddNull)178 TEST(KeysetWrapperStoreTest, AddNull) {
179 RegistryImpl registry;
180 util::StatusOr<std::function<util::StatusOr<std::unique_ptr<FakePrimitive>>(
181 const KeyData& key_data)>>
182 primitive_getter = PrimitiveGetter(registry);
183 ASSERT_THAT(primitive_getter, IsOk());
184
185 KeysetWrapperStore store;
186 EXPECT_THAT((store.Add<FakePrimitive, FakePrimitive>(/*wrapper=*/nullptr,
187 *primitive_getter)),
188 StatusIs(absl::StatusCode::kInvalidArgument));
189 }
190
TEST(KeysetWrapperStoreTest,AddWrappersForDifferentPrimitivesSucceeds)191 TEST(KeysetWrapperStoreTest, AddWrappersForDifferentPrimitivesSucceeds) {
192 RegistryImpl registry;
193 util::StatusOr<std::function<util::StatusOr<std::unique_ptr<FakePrimitive>>(
194 const KeyData& key_data)>>
195 primitive_getter = PrimitiveGetter(registry);
196 ASSERT_THAT(primitive_getter, IsOk());
197
198 KeysetWrapperStore store;
199 ASSERT_THAT(
200 (store.Add<FakePrimitive, FakePrimitive>(
201 absl::make_unique<FakePrimitiveWrapper>(), *primitive_getter)),
202 IsOk());
203
204 std::function<util::StatusOr<std::unique_ptr<Mac>>(const KeyData& key_data)>
205 primitive_getter_mac = [®istry](const KeyData& key_data) {
206 return registry.GetPrimitive<Mac>(key_data);
207 };
208 EXPECT_THAT((store.Add<Mac, Mac>(absl::make_unique<MacWrapper>(),
209 primitive_getter_mac)),
210 IsOk());
211 }
212
TEST(KeysetWrapperStoreTest,AddSameWrapperTwiceSucceeds)213 TEST(KeysetWrapperStoreTest, AddSameWrapperTwiceSucceeds) {
214 RegistryImpl registry;
215 util::StatusOr<std::function<util::StatusOr<std::unique_ptr<FakePrimitive>>(
216 const KeyData& key_data)>>
217 primitive_getter = PrimitiveGetter(registry);
218 ASSERT_THAT(primitive_getter, IsOk());
219
220 KeysetWrapperStore store;
221 ASSERT_THAT(
222 (store.Add<FakePrimitive, FakePrimitive>(
223 absl::make_unique<FakePrimitiveWrapper>(), *primitive_getter)),
224 IsOk());
225 EXPECT_THAT(
226 (store.Add<FakePrimitive, FakePrimitive>(
227 absl::make_unique<FakePrimitiveWrapper>(), *primitive_getter)),
228 IsOk());
229 }
230
TEST(KeysetWrapperStoreTest,AddDifferentWrappersForSamePrimitiveFails)231 TEST(KeysetWrapperStoreTest, AddDifferentWrappersForSamePrimitiveFails) {
232 RegistryImpl registry;
233 util::StatusOr<std::function<util::StatusOr<std::unique_ptr<FakePrimitive>>(
234 const KeyData& key_data)>>
235 primitive_getter = PrimitiveGetter(registry);
236 ASSERT_THAT(primitive_getter, IsOk());
237
238 KeysetWrapperStore store;
239 ASSERT_THAT(
240 (store.Add<FakePrimitive, FakePrimitive>(
241 absl::make_unique<FakePrimitiveWrapper>(), *primitive_getter)),
242 IsOk());
243 EXPECT_THAT(
244 (store.Add<FakePrimitive, FakePrimitive>(
245 absl::make_unique<FakePrimitiveWrapper2>(), *primitive_getter)),
246 StatusIs(absl::StatusCode::kAlreadyExists));
247 }
248
TEST(KeysetWrapperStoreTest,GetPrimitiveWrapper)249 TEST(KeysetWrapperStoreTest, GetPrimitiveWrapper) {
250 RegistryImpl registry;
251 util::StatusOr<std::function<util::StatusOr<std::unique_ptr<FakePrimitive>>(
252 const KeyData& key_data)>>
253 primitive_getter = PrimitiveGetter(registry);
254 ASSERT_THAT(primitive_getter, IsOk());
255
256 KeysetWrapperStore store;
257 ASSERT_THAT(
258 (store.Add<FakePrimitive, FakePrimitive>(
259 absl::make_unique<FakePrimitiveWrapper>(), *primitive_getter)),
260 IsOk());
261
262 util::StatusOr<const PrimitiveWrapper<FakePrimitive, FakePrimitive>*>
263 legacy_wrapper = store.GetPrimitiveWrapper<FakePrimitive>();
264 ASSERT_THAT(legacy_wrapper, IsOk());
265
266 Keyset keyset;
267 std::string raw_key = AddAesGcmKeyToKeyset(keyset, 13, OutputPrefixType::TINK,
268 KeyStatusType::ENABLED);
269 KeysetInfo keyset_info;
270 keyset_info.add_key_info();
271 keyset_info.mutable_key_info(0)->set_output_prefix_type(
272 OutputPrefixType::TINK);
273 keyset_info.mutable_key_info(0)->set_key_id(1234543);
274 keyset_info.mutable_key_info(0)->set_status(KeyStatusType::ENABLED);
275 keyset_info.set_primary_key_id(1234543);
276 std::unique_ptr<PrimitiveSet<FakePrimitive>> primitive_set(
277 new PrimitiveSet<FakePrimitive>());
278 auto entry = primitive_set->AddPrimitive(
279 absl::make_unique<FakePrimitive>(raw_key), keyset_info.key_info(0));
280 ASSERT_THAT(entry, IsOk());
281 ASSERT_THAT(primitive_set->set_primary(*entry), IsOk());
282
283 util::StatusOr<std::unique_ptr<FakePrimitive>> legacy_aead =
284 (*legacy_wrapper)->Wrap(std::move(primitive_set));
285 ASSERT_THAT(legacy_aead, IsOk());
286 EXPECT_THAT((*legacy_aead)->get(), Eq(raw_key));
287 }
288
TEST(KeysetWrapperStoreTest,GetPrimitiveWrapperNonexistentWrapperFails)289 TEST(KeysetWrapperStoreTest, GetPrimitiveWrapperNonexistentWrapperFails) {
290 RegistryImpl registry;
291 util::StatusOr<std::function<util::StatusOr<std::unique_ptr<FakePrimitive>>(
292 const KeyData& key_data)>>
293 primitive_getter = PrimitiveGetter(registry);
294 ASSERT_THAT(primitive_getter, IsOk());
295
296 KeysetWrapperStore store;
297 ASSERT_THAT(
298 (store.Add<FakePrimitive, FakePrimitive>(
299 absl::make_unique<FakePrimitiveWrapper>(), *primitive_getter)),
300 IsOk());
301
302 EXPECT_THAT(store.GetPrimitiveWrapper<Mac>().status(),
303 StatusIs(absl::StatusCode::kNotFound));
304 }
305
TEST(KeysetWrapperStoreTest,Get)306 TEST(KeysetWrapperStoreTest, Get) {
307 RegistryImpl registry;
308 util::StatusOr<std::function<util::StatusOr<std::unique_ptr<FakePrimitive>>(
309 const KeyData& key_data)>>
310 primitive_getter = PrimitiveGetter(registry);
311 ASSERT_THAT(primitive_getter, IsOk());
312
313 KeysetWrapperStore store;
314 ASSERT_THAT(
315 (store.Add<FakePrimitive, FakePrimitive>(
316 absl::make_unique<FakePrimitiveWrapper>(), *primitive_getter)),
317 IsOk());
318
319 util::StatusOr<const KeysetWrapper<FakePrimitive>*> wrapper =
320 store.Get<FakePrimitive>();
321 ASSERT_THAT(wrapper, IsOk());
322
323 Keyset keyset;
324 std::string raw_key = AddAesGcmKeyToKeyset(keyset, 13, OutputPrefixType::TINK,
325 KeyStatusType::ENABLED);
326 keyset.set_primary_key_id(13);
327
328 util::StatusOr<std::unique_ptr<FakePrimitive>> aead =
329 (*wrapper)->Wrap(keyset, /*annotations=*/{});
330 ASSERT_THAT(aead, IsOk());
331 EXPECT_THAT((*aead)->get(), Eq(raw_key));
332 }
333
TEST(KeysetWrapperStoreTest,GetNonexistentWrapperFails)334 TEST(KeysetWrapperStoreTest, GetNonexistentWrapperFails) {
335 RegistryImpl registry;
336 util::StatusOr<std::function<util::StatusOr<std::unique_ptr<FakePrimitive>>(
337 const KeyData& key_data)>>
338 primitive_getter = PrimitiveGetter(registry);
339 ASSERT_THAT(primitive_getter, IsOk());
340
341 KeysetWrapperStore store;
342 ASSERT_THAT(
343 (store.Add<FakePrimitive, FakePrimitive>(
344 absl::make_unique<FakePrimitiveWrapper>(), *primitive_getter)),
345 IsOk());
346
347 EXPECT_THAT(store.Get<Mac>().status(), StatusIs(absl::StatusCode::kNotFound));
348 }
349
TEST(KeysetWrapperStoreTest,IsEmpty)350 TEST(KeysetWrapperStoreTest, IsEmpty) {
351 KeysetWrapperStore store;
352 EXPECT_EQ(store.IsEmpty(), true);
353
354 RegistryImpl registry;
355 util::StatusOr<std::function<util::StatusOr<std::unique_ptr<FakePrimitive>>(
356 const KeyData& key_data)>>
357 primitive_getter = PrimitiveGetter(registry);
358 ASSERT_THAT(primitive_getter, IsOk());
359 ASSERT_THAT(
360 (store.Add<FakePrimitive, FakePrimitive>(
361 absl::make_unique<FakePrimitiveWrapper>(), *primitive_getter)),
362 IsOk());
363 EXPECT_THAT(store.IsEmpty(), false);
364 }
365
TEST(KeysetWrapperStoreTest,Move)366 TEST(KeysetWrapperStoreTest, Move) {
367 RegistryImpl registry;
368 util::StatusOr<std::function<util::StatusOr<std::unique_ptr<FakePrimitive>>(
369 const KeyData& key_data)>>
370 primitive_getter = PrimitiveGetter(registry);
371 ASSERT_THAT(primitive_getter, IsOk());
372
373 KeysetWrapperStore store;
374 ASSERT_THAT(
375 (store.Add<FakePrimitive, FakePrimitive>(
376 absl::make_unique<FakePrimitiveWrapper>(), *primitive_getter)),
377 IsOk());
378
379 util::StatusOr<const KeysetWrapper<FakePrimitive>*> wrapper =
380 store.Get<FakePrimitive>();
381 ASSERT_THAT(wrapper, IsOk());
382
383 KeysetWrapperStore new_store = std::move(store);
384 wrapper = new_store.Get<FakePrimitive>();
385 ASSERT_THAT(wrapper, IsOk());
386
387 Keyset keyset;
388 std::string raw_key = AddAesGcmKeyToKeyset(keyset, 13, OutputPrefixType::TINK,
389 KeyStatusType::ENABLED);
390 keyset.set_primary_key_id(13);
391
392 util::StatusOr<std::unique_ptr<FakePrimitive>> aead =
393 (*wrapper)->Wrap(keyset, /*annotations=*/{});
394 ASSERT_THAT(aead, IsOk());
395 EXPECT_THAT((*aead)->get(), Eq(raw_key));
396 }
397
398 } // namespace
399 } // namespace internal
400 } // namespace tink
401 } // namespace crypto
402