1 // Copyright 2020 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 #include "tink/internal/keyset_wrapper_impl.h"
17
18 #include <memory>
19 #include <string>
20 #include <utility>
21 #include <vector>
22
23 #include "gmock/gmock.h"
24 #include "gtest/gtest.h"
25 #include "absl/container/flat_hash_map.h"
26 #include "absl/memory/memory.h"
27 #include "absl/status/status.h"
28 #include "absl/strings/match.h"
29 #include "absl/strings/string_view.h"
30 #include "tink/primitive_set.h"
31 #include "tink/primitive_wrapper.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/tink.pb.h"
37
38 namespace crypto {
39 namespace tink {
40 namespace internal {
41
42 namespace {
43
44 using ::crypto::tink::test::AddKeyData;
45 using ::crypto::tink::test::IsOk;
46 using ::crypto::tink::test::IsOkAndHolds;
47 using ::google::crypto::tink::Keyset;
48 using ::testing::HasSubstr;
49 using ::testing::Not;
50 using ::testing::Pair;
51 using ::testing::Pointee;
52 using ::testing::UnorderedElementsAre;
53
54 using InputPrimitive = std::string;
55 using OutputPrimitive = std::vector<std::pair<int, std::string>>;
56
57 // This "Wrapper" wraps primitives of type std::string into primitives of type
58 // std::vector<int, std::string> simply by returning pairs {key_id, string}.
59 // It appends " (primary)" to the string for the primary id.
60 class Wrapper : public PrimitiveWrapper<InputPrimitive, OutputPrimitive> {
61 public:
Wrap(std::unique_ptr<PrimitiveSet<InputPrimitive>> primitive_set) const62 crypto::tink::util::StatusOr<std::unique_ptr<OutputPrimitive>> Wrap(
63 std::unique_ptr<PrimitiveSet<InputPrimitive>> primitive_set)
64 const override {
65 auto result = absl::make_unique<OutputPrimitive>();
66 for (const auto* entry : primitive_set->get_all()) {
67 (*result).push_back(
68 std::make_pair(entry->get_key_id(), entry->get_primitive()));
69 if (entry->get_key_id() == primitive_set->get_primary()->get_key_id()) {
70 result->back().second.append(" (primary)");
71 }
72 }
73 return std::move(result);
74 }
75 };
76
CreateIn(const google::crypto::tink::KeyData & key_data)77 crypto::tink::util::StatusOr<std::unique_ptr<InputPrimitive>> CreateIn(
78 const google::crypto::tink::KeyData& key_data) {
79 if (absl::StartsWith(key_data.type_url(), "error:")) {
80 return crypto::tink::util::Status(absl::StatusCode::kInvalidArgument,
81 key_data.type_url());
82 } else {
83 return absl::make_unique<InputPrimitive>(key_data.type_url());
84 }
85 }
86
OnlyTypeUrlKeyData(absl::string_view type_url)87 google::crypto::tink::KeyData OnlyTypeUrlKeyData(absl::string_view type_url) {
88 google::crypto::tink::KeyData result;
89 result.set_type_url(std::string(type_url));
90 return result;
91 }
92
CreateKeyset(const std::vector<std::pair<int,std::string>> & keydata)93 google::crypto::tink::Keyset CreateKeyset(
94 const std::vector<std::pair<int, std::string>>& keydata) {
95 google::crypto::tink::Keyset keyset;
96 for (const auto& pair : keydata) {
97 AddKeyData(OnlyTypeUrlKeyData(pair.second), pair.first,
98 google::crypto::tink::OutputPrefixType::TINK,
99 google::crypto::tink::KeyStatusType::ENABLED, &keyset);
100 }
101 return keyset;
102 }
103
TEST(KeysetWrapperImplTest,Basic)104 TEST(KeysetWrapperImplTest, Basic) {
105 Wrapper wrapper;
106 auto wrapper_or =
107 absl::make_unique<KeysetWrapperImpl<InputPrimitive, OutputPrimitive>>(
108 &wrapper, &CreateIn);
109 std::vector<std::pair<int, std::string>> keydata = {
110 {111, "one"}, {222, "two"}, {333, "three"}};
111 google::crypto::tink::Keyset keyset = CreateKeyset(keydata);
112 keyset.set_primary_key_id(222);
113
114 util::StatusOr<std::unique_ptr<OutputPrimitive>> wrapped =
115 wrapper_or->Wrap(keyset, /*annotations=*/{});
116
117 ASSERT_THAT(wrapped, IsOk());
118 ASSERT_THAT(*wrapped.value(),
119 UnorderedElementsAre(Pair(111, "one"), Pair(222, "two (primary)"),
120 Pair(333, "three")));
121 }
122
TEST(KeysetWrapperImplTest,FailingGetPrimitive)123 TEST(KeysetWrapperImplTest, FailingGetPrimitive) {
124 Wrapper wrapper;
125 auto wrapper_or =
126 absl::make_unique<KeysetWrapperImpl<InputPrimitive, OutputPrimitive>>(
127 &wrapper, &CreateIn);
128 std::vector<std::pair<int, std::string>> keydata = {{1, "ok:one"},
129 {2, "error:two"}};
130 google::crypto::tink::Keyset keyset = CreateKeyset(keydata);
131 keyset.set_primary_key_id(1);
132
133 util::StatusOr<std::unique_ptr<OutputPrimitive>> wrapped =
134 wrapper_or->Wrap(keyset, /*annotations=*/{});
135
136 ASSERT_THAT(wrapped, Not(IsOk()));
137 ASSERT_THAT(std::string(wrapped.status().message()), HasSubstr("error:two"));
138 }
139
140 // This test checks that validate keyset is called. We simply pass an empty
141 // keyset.
TEST(KeysetWrapperImplTest,ValidatesKeyset)142 TEST(KeysetWrapperImplTest, ValidatesKeyset) {
143 Wrapper wrapper;
144 auto wrapper_or =
145 absl::make_unique<KeysetWrapperImpl<InputPrimitive, OutputPrimitive>>(
146 &wrapper, &CreateIn);
147 util::StatusOr<std::unique_ptr<OutputPrimitive>> wrapped =
148 wrapper_or->Wrap(google::crypto::tink::Keyset(), /*annotations=*/{});
149
150 ASSERT_THAT(wrapped, Not(IsOk()));
151 }
152
153 // This test checks that only enabled keys are used to create the primitive set.
TEST(KeysetWrapperImplTest,OnlyEnabled)154 TEST(KeysetWrapperImplTest, OnlyEnabled) {
155 Wrapper wrapper;
156 auto wrapper_or =
157 absl::make_unique<KeysetWrapperImpl<InputPrimitive, OutputPrimitive>>(
158 &wrapper, &CreateIn);
159 std::vector<std::pair<int, std::string>> keydata = {
160 {111, "one"}, {222, "two"}, {333, "three"}, {444, "four"}};
161 google::crypto::tink::Keyset keyset = CreateKeyset(keydata);
162 keyset.set_primary_key_id(222);
163 // KeyId 333 is index 2.
164 keyset.mutable_key(2)->set_status(google::crypto::tink::DISABLED);
165 util::StatusOr<std::unique_ptr<OutputPrimitive>> wrapped =
166 wrapper_or->Wrap(keyset, /*annotations=*/{});
167
168 ASSERT_THAT(wrapped, IsOk());
169 ASSERT_THAT(*wrapped.value(),
170 UnorderedElementsAre(Pair(111, "one"), Pair(222, "two (primary)"),
171 Pair(444, "four")));
172 }
173
174 // Mock PrimitiveWrapper with input primitive I and output primitive O.
175 template <class I, class O>
176 class MockWrapper : public PrimitiveWrapper<I, O> {
177 public:
178 MOCK_METHOD(util::StatusOr<std::unique_ptr<O>>, Wrap,
179 (std::unique_ptr<PrimitiveSet<I>> primitive_set),
180 (const, override));
181 };
182
183 // Returns a valid output primitive.
GetOutputPrimitiveForTesting()184 std::unique_ptr<OutputPrimitive> GetOutputPrimitiveForTesting() {
185 auto output_primitive = absl::make_unique<OutputPrimitive>();
186 output_primitive->push_back({111, "one"});
187 output_primitive->push_back({222, "two (primary)"});
188 output_primitive->push_back({333, "three"});
189 output_primitive->push_back({444, "four"});
190 return output_primitive;
191 }
192
193 // Tests that annotations are correctly passed on to the generated PrimitiveSet.
TEST(KeysetWrapperImplTest,WrapWithAnnotationCorrectlyWrittenToPrimitiveSet)194 TEST(KeysetWrapperImplTest, WrapWithAnnotationCorrectlyWrittenToPrimitiveSet) {
195 MockWrapper<InputPrimitive, OutputPrimitive> wrapper;
196 auto keyset_wrapper =
197 absl::make_unique<KeysetWrapperImpl<InputPrimitive, OutputPrimitive>>(
198 &wrapper, CreateIn);
199 Keyset keyset = CreateKeyset(
200 /*keydata=*/{{111, "one"}, {222, "two"}, {333, "three"}, {444, "four"}});
201 keyset.set_primary_key_id(222);
202 const absl::flat_hash_map<std::string, std::string> kExpectedAnnotations = {
203 {"key1", "value1"}, {"key2", "value2"}};
204
205 absl::flat_hash_map<std::string, std::string> generated_annotations;
206 EXPECT_CALL(wrapper, Wrap(testing::_))
207 .WillOnce(
208 [&generated_annotations](std::unique_ptr<PrimitiveSet<InputPrimitive>>
209 generated_primitive_set) {
210 // We are interested to check if the annotations are what we
211 // expected, so we copy them to `generated_annotations`.
212 generated_annotations = generated_primitive_set->get_annotations();
213 // Return a valid output primitive.
214 return GetOutputPrimitiveForTesting();
215 });
216
217 util::StatusOr<std::unique_ptr<OutputPrimitive>> wrapped_primitive =
218 keyset_wrapper->Wrap(keyset, kExpectedAnnotations);
219
220 EXPECT_EQ(generated_annotations, kExpectedAnnotations);
221 EXPECT_THAT(wrapped_primitive,
222 IsOkAndHolds(Pointee(UnorderedElementsAre(
223 Pair(111, "one"), Pair(222, "two (primary)"),
224 Pair(333, "three"), Pair(444, "four")))));
225 }
226
227 } // namespace
228
229 } // namespace internal
230 } // namespace tink
231 } // namespace crypto
232