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/prf/prf_set_wrapper.h"
17
18 #include <cstdint>
19 #include <map>
20 #include <memory>
21 #include <string>
22 #include <utility>
23
24 #include "gmock/gmock.h"
25 #include "gtest/gtest.h"
26 #include "absl/memory/memory.h"
27 #include "absl/status/status.h"
28 #include "absl/strings/string_view.h"
29 #include "tink/monitoring/monitoring_client_mocks.h"
30 #include "tink/prf/prf_set.h"
31 #include "tink/primitive_set.h"
32 #include "tink/registry.h"
33 #include "tink/util/status.h"
34 #include "tink/util/statusor.h"
35 #include "tink/util/test_matchers.h"
36 #include "proto/tink.pb.h"
37
38 namespace crypto {
39 namespace tink {
40 namespace {
41
42 using ::crypto::tink::test::IsOk;
43 using ::crypto::tink::test::IsOkAndHolds;
44 using ::google::crypto::tink::KeysetInfo;
45 using ::google::crypto::tink::KeyStatusType;
46 using ::testing::_;
47 using ::testing::ByMove;
48 using ::testing::Key;
49 using ::testing::NiceMock;
50 using ::testing::Not;
51 using ::testing::Return;
52 using ::testing::StrEq;
53 using ::testing::Test;
54 using ::testing::UnorderedElementsAre;
55
MakeKey(uint32_t id)56 KeysetInfo::KeyInfo MakeKey(uint32_t id) {
57 KeysetInfo::KeyInfo key;
58 key.set_output_prefix_type(google::crypto::tink::OutputPrefixType::RAW);
59 key.set_key_id(id);
60 key.set_status(KeyStatusType::ENABLED);
61 return key;
62 }
63
64 class FakePrf : public Prf {
65 public:
FakePrf(const std::string & output)66 explicit FakePrf(const std::string& output) : output_(output) {}
Compute(absl::string_view input,size_t output_length) const67 util::StatusOr<std::string> Compute(absl::string_view input,
68 size_t output_length) const override {
69 return output_;
70 }
71
72 private:
73 std::string output_;
74 };
75
76 class PrfSetWrapperTest : public ::testing::Test {
77 protected:
SetUp()78 void SetUp() override { prf_set_ = absl::make_unique<PrimitiveSet<Prf>>(); }
79
AddPrf(const std::string & output,const KeysetInfo::KeyInfo & key_info)80 util::StatusOr<PrimitiveSet<Prf>::Entry<Prf>*> AddPrf(
81 const std::string& output, const KeysetInfo::KeyInfo& key_info) {
82 auto prf = absl::make_unique<FakePrf>(output);
83 return prf_set_->AddPrimitive(std::move(prf), key_info);
84 }
85
PrfSet()86 std::unique_ptr<PrimitiveSet<Prf>>& PrfSet() { return prf_set_; }
87
88 private:
89 std::unique_ptr<PrimitiveSet<Prf>> prf_set_;
90 };
91
TEST_F(PrfSetWrapperTest,NullPrfSet)92 TEST_F(PrfSetWrapperTest, NullPrfSet) {
93 PrfSetWrapper wrapper;
94 EXPECT_THAT(wrapper.Wrap(nullptr), Not(IsOk()));
95 }
96
TEST_F(PrfSetWrapperTest,EmptyPrfSet)97 TEST_F(PrfSetWrapperTest, EmptyPrfSet) {
98 PrfSetWrapper wrapper;
99 EXPECT_THAT(wrapper.Wrap(absl::make_unique<PrimitiveSet<Prf>>()).status(),
100 Not(IsOk()));
101 }
102
TEST_F(PrfSetWrapperTest,NonRawKeyType)103 TEST_F(PrfSetWrapperTest, NonRawKeyType) {
104 KeysetInfo::KeyInfo key_info = MakeKey(1);
105 key_info.set_output_prefix_type(google::crypto::tink::OutputPrefixType::TINK);
106 auto entry = AddPrf("output", key_info);
107 ASSERT_THAT(entry, IsOk());
108 ASSERT_THAT(PrfSet()->set_primary(entry.value()), IsOk());
109 PrfSetWrapper wrapper;
110 EXPECT_THAT(wrapper.Wrap(std::move(PrfSet())), Not(IsOk()));
111 }
112
TEST_F(PrfSetWrapperTest,WrapOkay)113 TEST_F(PrfSetWrapperTest, WrapOkay) {
114 auto entry = AddPrf("output", MakeKey(1));
115 ASSERT_THAT(entry, IsOk());
116 ASSERT_THAT(PrfSet()->set_primary(entry.value()), IsOk());
117 PrfSetWrapper wrapper;
118 auto wrapped = wrapper.Wrap(std::move(PrfSet()));
119 ASSERT_THAT(wrapped, IsOk());
120 EXPECT_THAT(wrapped.value()->ComputePrimary("input", 6),
121 IsOkAndHolds(StrEq("output")));
122 }
123
TEST_F(PrfSetWrapperTest,WrapTwo)124 TEST_F(PrfSetWrapperTest, WrapTwo) {
125 std::string primary_output("output");
126 auto entry = AddPrf(primary_output, MakeKey(1));
127 ASSERT_THAT(entry, IsOk());
128 ASSERT_THAT(PrfSet()->set_primary(entry.value()), IsOk());
129
130 ASSERT_THAT(AddPrf(primary_output, MakeKey(1)), IsOk());
131 std::string secondary_output("different");
132 ASSERT_THAT(AddPrf(secondary_output, MakeKey(2)), IsOk());
133 PrfSetWrapper wrapper;
134 auto wrapped_or = wrapper.Wrap(std::move(PrfSet()));
135 ASSERT_THAT(wrapped_or, IsOk());
136 auto wrapped = std::move(wrapped_or.value());
137 EXPECT_THAT(wrapped->ComputePrimary("input", 6),
138 IsOkAndHolds(StrEq("output")));
139 const auto& prf_map = wrapped->GetPrfs();
140 ASSERT_THAT(prf_map, UnorderedElementsAre(Key(1), Key(2)));
141 EXPECT_THAT(prf_map.find(1)->second->Compute("input", 6),
142 IsOkAndHolds(StrEq("output")));
143 EXPECT_THAT(prf_map.find(2)->second->Compute("input", 6),
144 IsOkAndHolds(StrEq("different")));
145 }
146
147 // Tests for the monitoring behavior.
148 class PrfSetWrapperWithMonitoringTest : public Test {
149 protected:
150 // Reset the global registry.
SetUp()151 void SetUp() override {
152 Registry::Reset();
153 // Setup mocks for catching Monitoring calls.
154 auto monitoring_client_factory =
155 absl::make_unique<MockMonitoringClientFactory>();
156 auto monitoring_client =
157 absl::make_unique<NiceMock<MockMonitoringClient>>();
158 monitoring_client_ref_ = monitoring_client.get();
159 // Monitoring tests expect that the client factory will create the
160 // corresponding MockMonitoringClients.
161 EXPECT_CALL(*monitoring_client_factory, New(_))
162 .WillOnce(
163 Return(ByMove(util::StatusOr<std::unique_ptr<MonitoringClient>>(
164 std::move(monitoring_client)))));
165
166 ASSERT_THAT(internal::RegistryImpl::GlobalInstance()
167 .RegisterMonitoringClientFactory(
168 std::move(monitoring_client_factory)),
169 IsOk());
170 ASSERT_THAT(
171 internal::RegistryImpl::GlobalInstance().GetMonitoringClientFactory(),
172 Not(testing::IsNull()));
173 }
174
175 // Cleanup the registry to avoid mock leaks.
~PrfSetWrapperWithMonitoringTest()176 ~PrfSetWrapperWithMonitoringTest() override { Registry::Reset(); }
177
178 NiceMock<MockMonitoringClient>* monitoring_client_ref_;
179 };
180
181 class AlwaysFailingPrf : public Prf {
182 public:
183 AlwaysFailingPrf() = default;
184
Compute(absl::string_view input,size_t output_length) const185 util::StatusOr<std::string> Compute(absl::string_view input,
186 size_t output_length) const override {
187 return util::Status(absl::StatusCode::kOutOfRange, "AlwaysFailingPrf");
188 }
189 };
190
TEST_F(PrfSetWrapperWithMonitoringTest,WrapKeysetWithMonitoringFailure)191 TEST_F(PrfSetWrapperWithMonitoringTest, WrapKeysetWithMonitoringFailure) {
192 const absl::flat_hash_map<std::string, std::string> annotations = {
193 {"key1", "value1"}, {"key2", "value2"}, {"key3", "value3"}};
194 auto primitive_set = absl::make_unique<PrimitiveSet<Prf>>(annotations);
195 util::StatusOr<PrimitiveSet<Prf>::Entry<Prf>*> entry =
196 primitive_set->AddPrimitive(absl::make_unique<AlwaysFailingPrf>(),
197 MakeKey(/*id=*/1));
198 ASSERT_THAT(entry, IsOk());
199 ASSERT_THAT(primitive_set->set_primary(entry.value()), IsOk());
200 ASSERT_THAT(primitive_set
201 ->AddPrimitive(absl::make_unique<FakePrf>("output"),
202 MakeKey(/*id=*/1))
203 .status(),
204 IsOk());
205 util::StatusOr<std::unique_ptr<PrfSet>> prf_set =
206 PrfSetWrapper().Wrap(std::move(primitive_set));
207 ASSERT_THAT(prf_set, IsOk());
208 EXPECT_CALL(*monitoring_client_ref_, LogFailure());
209 EXPECT_THAT((*prf_set)->ComputePrimary("input", /*output_length=*/16),
210 Not(IsOk()));
211 }
212
TEST_F(PrfSetWrapperWithMonitoringTest,WrapKeysetWithMonitoringVerifySuccess)213 TEST_F(PrfSetWrapperWithMonitoringTest, WrapKeysetWithMonitoringVerifySuccess) {
214 const absl::flat_hash_map<std::string, std::string> annotations = {
215 {"key1", "value1"}, {"key2", "value2"}, {"key3", "value3"}};
216 auto primitive_set = absl::make_unique<PrimitiveSet<Prf>>(annotations);
217
218 util::StatusOr<PrimitiveSet<Prf>::Entry<Prf>*> entry =
219 primitive_set->AddPrimitive(absl::make_unique<FakePrf>("output"),
220 MakeKey(/*id=*/1));
221 ASSERT_THAT(entry, IsOk());
222 ASSERT_THAT(primitive_set->set_primary(entry.value()), IsOk());
223 ASSERT_THAT(primitive_set
224 ->AddPrimitive(absl::make_unique<FakePrf>("output"),
225 MakeKey(/*id=*/1))
226 .status(),
227 IsOk());
228
229 util::StatusOr<std::unique_ptr<PrfSet>> prf_set =
230 PrfSetWrapper().Wrap(std::move(primitive_set));
231 ASSERT_THAT(prf_set, IsOk());
232 std::map<uint32_t, Prf*> prf_map = (*prf_set)->GetPrfs();
233 std::string input = "input";
234 for (const auto& entry : prf_map) {
235 EXPECT_CALL(*monitoring_client_ref_, Log(entry.first, input.size()));
236 EXPECT_THAT((entry.second)->Compute(input, /*output_length=*/16).status(),
237 IsOk());
238 }
239 input = "hello_world";
240 EXPECT_CALL(*monitoring_client_ref_,
241 Log((*prf_set)->GetPrimaryId(), input.size()));
242 EXPECT_THAT((*prf_set)->ComputePrimary(input, /*output_length=*/16), IsOk());
243 }
244
245 } // namespace
246 } // namespace tink
247 } // namespace crypto
248