1 // Copyright 2021 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/aead/internal/zero_copy_aead_wrapper.h"
18
19 #include <cstring>
20 #include <memory>
21 #include <string>
22 #include <utility>
23 #include <vector>
24
25 #include "gtest/gtest.h"
26 #include "tink/aead/internal/mock_zero_copy_aead.h"
27 #include "tink/subtle/subtle_util.h"
28 #include "tink/util/test_matchers.h"
29
30 namespace crypto {
31 namespace tink {
32 namespace internal {
33 namespace {
34
35 using ::crypto::tink::PrimitiveSet;
36 using ::crypto::tink::test::IsOk;
37 using ::crypto::tink::test::StatusIs;
38 using ::crypto::tink::util::StatusOr;
39 using ::google::crypto::tink::KeysetInfo;
40 using ::google::crypto::tink::KeyStatusType;
41 using ::google::crypto::tink::OutputPrefixType;
42 using ::testing::_;
43 using ::testing::HasSubstr;
44 using ::testing::Invoke;
45 using ::testing::Return;
46 using ::testing::Unused;
47
48 constexpr absl::string_view kPlaintext = "Some data to encrypt.";
49 constexpr absl::string_view kAad = "Some data to authenticate.";
50 constexpr absl::string_view kCiphertext = "iv:Some data to encrypt.:tag";
51
52 using ZeroCopyAeadEntry =
53 crypto::tink::PrimitiveSet<ZeroCopyAead>::Entry<ZeroCopyAead>;
54
TEST(ZeroCopyAeadWrapperEmptyTest,Nullptr)55 TEST(ZeroCopyAeadWrapperEmptyTest, Nullptr) {
56 ZeroCopyAeadWrapper wrapper;
57 StatusOr<std::unique_ptr<Aead>> aead_set = wrapper.Wrap(nullptr);
58 EXPECT_THAT(aead_set.status(),
59 StatusIs(absl::StatusCode::kInternal, HasSubstr("non-NULL")));
60 }
61
TEST(ZeroCopyAeadWrapperEmptyTest,Empty)62 TEST(ZeroCopyAeadWrapperEmptyTest, Empty) {
63 ZeroCopyAeadWrapper wrapper;
64 StatusOr<std::unique_ptr<Aead>> aead_set =
65 wrapper.Wrap(absl::make_unique<PrimitiveSet<ZeroCopyAead>>());
66 EXPECT_THAT(aead_set.status(), StatusIs(absl::StatusCode::kInvalidArgument,
67 HasSubstr("no primary")));
68 }
69
70 class ZeroCopyAeadWrapperTest : public testing::Test {
71 protected:
SetUp()72 void SetUp() override {
73 // Defines a Tink-type key.
74 KeysetInfo::KeyInfo key_info;
75 key_info.set_output_prefix_type(OutputPrefixType::TINK);
76 key_info.set_key_id(1234543);
77 key_info.set_status(KeyStatusType::ENABLED);
78
79 // Creates a new AEAD set, adds a mock AEAD corresponding to the above key,
80 // and stores the set as aead_set_.
81 std::unique_ptr<PrimitiveSet<ZeroCopyAead>> aead_set(
82 new PrimitiveSet<ZeroCopyAead>());
83 auto entry = aead_set->AddPrimitive(SetUpMockZeroCopyAead(), key_info);
84 ASSERT_THAT(entry, IsOk());
85 ASSERT_THAT(aead_set->set_primary(*entry), IsOk());
86 aead_set_ = std::move(aead_set);
87 }
88
89 // Returns an AEAD with expected return values for all its functions set via
90 // EXPECT_CALL. All values are derived from constants kPlaintext, kAad, and
91 // kCiphertext.
SetUpMockZeroCopyAead()92 std::unique_ptr<MockZeroCopyAead> SetUpMockZeroCopyAead() {
93 auto aead = absl::make_unique<MockZeroCopyAead>();
94
95 EXPECT_CALL(*aead, MaxEncryptionSize(kPlaintext.size()))
96 .WillRepeatedly(Return(kCiphertext.size()));
97 EXPECT_CALL(*aead, Encrypt(kPlaintext, kAad, _))
98 .WillRepeatedly(Invoke([&](Unused, Unused, absl::Span<char> buffer) {
99 memcpy(buffer.data(), kCiphertext.data(), kCiphertext.size());
100 return kCiphertext.size();
101 }));
102 EXPECT_CALL(*aead, MaxDecryptionSize(kCiphertext.size()))
103 .WillRepeatedly(Return(kPlaintext.size()));
104 EXPECT_CALL(*aead, Decrypt(kCiphertext, kAad, _))
105 .WillRepeatedly(Invoke([&](Unused, Unused, absl::Span<char> buffer) {
106 std::memcpy(buffer.data(), kPlaintext.data(), kPlaintext.size());
107 return kPlaintext.size();
108 }));
109
110 return aead;
111 }
112
113 std::unique_ptr<PrimitiveSet<ZeroCopyAead>> aead_set_;
114 };
115
TEST_F(ZeroCopyAeadWrapperTest,EncryptDecrypt)116 TEST_F(ZeroCopyAeadWrapperTest, EncryptDecrypt) {
117 ZeroCopyAeadWrapper wrapper;
118 StatusOr<std::unique_ptr<Aead>> aead_set = wrapper.Wrap(std::move(aead_set_));
119 ASSERT_THAT(aead_set, IsOk());
120
121 StatusOr<std::string> ciphertext = (*aead_set)->Encrypt(kPlaintext, kAad);
122 ASSERT_THAT(ciphertext, IsOk());
123 StatusOr<std::string> plaintext = (*aead_set)->Decrypt(*ciphertext, kAad);
124 ASSERT_THAT(plaintext, IsOk());
125 EXPECT_EQ(*plaintext, kPlaintext);
126 }
127
TEST_F(ZeroCopyAeadWrapperTest,EncryptMultipleKeys)128 TEST_F(ZeroCopyAeadWrapperTest, EncryptMultipleKeys) {
129 // Manually encrypt with the primary key.
130 ZeroCopyAead& aead = aead_set_->get_primary()->get_primitive();
131 std::string ciphertext;
132 subtle::ResizeStringUninitialized(
133 &ciphertext, CryptoFormat::kNonRawPrefixSize +
134 aead.MaxEncryptionSize(kPlaintext.size()));
135 StatusOr<int64_t> ciphertext_size = aead.Encrypt(
136 kPlaintext, kAad,
137 absl::MakeSpan(ciphertext)
138 .subspan(CryptoFormat::kNonRawPrefixSize, ciphertext.size()));
139 ASSERT_THAT(ciphertext_size, IsOk());
140 const std::string& key_id = aead_set_->get_primary()->get_identifier();
141 std::memcpy(&ciphertext[0], key_id.data(), key_id.size());
142 ciphertext.resize(key_id.size() + *ciphertext_size);
143
144 // Add a second key.
145 KeysetInfo::KeyInfo key_info;
146 key_info.set_output_prefix_type(OutputPrefixType::TINK);
147 key_info.set_key_id(42);
148 key_info.set_status(KeyStatusType::ENABLED);
149 std::unique_ptr<ZeroCopyAead> aead1 = absl::make_unique<MockZeroCopyAead>();
150 ASSERT_THAT(aead_set_->AddPrimitive(std::move(aead1), key_info).status(),
151 IsOk());
152 ZeroCopyAeadWrapper wrapper;
153 StatusOr<std::unique_ptr<Aead>> aead_set = wrapper.Wrap(std::move(aead_set_));
154 ASSERT_THAT(aead_set, IsOk());
155
156 // Encrypt with the wrapped AEAD and check that the result is equal to
157 // encrypting directly with the primary key.
158 StatusOr<std::string> wrap_ciphertext =
159 (*aead_set)->Encrypt(kPlaintext, kAad);
160 ASSERT_THAT(wrap_ciphertext, IsOk());
161 EXPECT_EQ(*wrap_ciphertext, ciphertext);
162 }
163
TEST_F(ZeroCopyAeadWrapperTest,EncryptDecryptRawKey)164 TEST_F(ZeroCopyAeadWrapperTest, EncryptDecryptRawKey) {
165 // Add raw key to AEAD set.
166 KeysetInfo::KeyInfo key_info;
167 key_info.set_output_prefix_type(OutputPrefixType::RAW);
168 key_info.set_key_id(1234);
169 key_info.set_status(KeyStatusType::ENABLED);
170 auto entry = aead_set_->AddPrimitive(SetUpMockZeroCopyAead(), key_info);
171 ASSERT_THAT(entry, IsOk());
172 ASSERT_THAT(aead_set_->set_primary(*entry), IsOk());
173
174 // Manually encrypt with the raw key.
175 util::StatusOr<const std::vector<std::unique_ptr<ZeroCopyAeadEntry>>*>
176 raw_primitives = aead_set_->get_raw_primitives();
177 ASSERT_THAT(raw_primitives, IsOk());
178 EXPECT_EQ((*raw_primitives)->size(), 1);
179 ZeroCopyAead& aead = (*raw_primitives)->front()->get_primitive();
180 std::string ciphertext;
181 subtle::ResizeStringUninitialized(&ciphertext,
182 aead.MaxEncryptionSize(kPlaintext.size()));
183 util::StatusOr<int64_t> ciphertext_size =
184 aead.Encrypt(kPlaintext, kAad, absl::MakeSpan(ciphertext));
185 ASSERT_THAT(ciphertext_size, IsOk());
186 ciphertext.resize(*ciphertext_size);
187
188 // Encrypt with the wrapped AEAD and check that the result is equal to
189 // encrypting directly with the raw key.
190 ZeroCopyAeadWrapper wrapper;
191 StatusOr<std::unique_ptr<Aead>> aead_set = wrapper.Wrap(std::move(aead_set_));
192 ASSERT_THAT(aead_set, IsOk());
193 StatusOr<std::string> wrap_ciphertext =
194 (*aead_set)->Encrypt(kPlaintext, kAad);
195 ASSERT_THAT(wrap_ciphertext, IsOk());
196 EXPECT_EQ(*wrap_ciphertext, ciphertext);
197
198 // Manually decrypt with the raw key.
199 std::string plaintext;
200 subtle::ResizeStringUninitialized(&plaintext,
201 aead.MaxDecryptionSize(ciphertext.size()));
202 util::StatusOr<int64_t> plaintext_size =
203 aead.Decrypt(ciphertext, kAad, absl::MakeSpan(plaintext));
204 ASSERT_THAT(plaintext_size, IsOk());
205 plaintext.resize(*plaintext_size);
206 EXPECT_EQ(plaintext, kPlaintext);
207
208 // Decrypt with the wrapped AEAD.
209 StatusOr<std::string> wrap_plaintext = (*aead_set)->Decrypt(ciphertext, kAad);
210 ASSERT_THAT(wrap_plaintext, IsOk());
211 EXPECT_EQ(*wrap_plaintext, kPlaintext);
212 }
213
TEST_F(ZeroCopyAeadWrapperTest,EncryptBadDecrypt)214 TEST_F(ZeroCopyAeadWrapperTest, EncryptBadDecrypt) {
215 ZeroCopyAeadWrapper wrapper;
216 StatusOr<std::unique_ptr<Aead>> aead_set = wrapper.Wrap(std::move(aead_set_));
217 ASSERT_THAT(aead_set, IsOk());
218
219 StatusOr<std::string> plaintext =
220 (*aead_set)->Decrypt("some bad ciphertext", kAad);
221 EXPECT_EQ(plaintext.status().code(), absl::StatusCode::kInvalidArgument);
222 }
223
224 } // namespace
225 } // namespace internal
226 } // namespace tink
227 } // namespace crypto
228