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 17package registry_test 18 19import ( 20 "bytes" 21 "errors" 22 "fmt" 23 "testing" 24 25 wrapperspb "google.golang.org/protobuf/types/known/wrapperspb" 26 "google.golang.org/protobuf/proto" 27 "github.com/google/tink/go/aead" 28 "github.com/google/tink/go/aead/subtle" 29 "github.com/google/tink/go/core/registry" 30 "github.com/google/tink/go/insecurecleartextkeyset" 31 "github.com/google/tink/go/internal/tinkerror" 32 "github.com/google/tink/go/keyset" 33 "github.com/google/tink/go/subtle/random" 34 tinkpb "github.com/google/tink/go/proto/tink_go_proto" 35) 36 37const ( 38 customTypeURL = "type.googleapis.com/google.crypto.tink.CustomAesGcmKey" 39) 40 41// customKeyManager is a custom implementation of registry.KeyManager for AES GCM 128. 42type customKeyManager struct{} 43 44// Assert that customKeyManager implements the KeyManager interface. 45var _ registry.KeyManager = (*customKeyManager)(nil) 46 47func (km *customKeyManager) Primitive(serializedKey []byte) (interface{}, error) { 48 key := new(wrapperspb.BytesValue) 49 if err := proto.Unmarshal(serializedKey, key); err != nil { 50 return nil, fmt.Errorf("invalid key") 51 } 52 if len(key.GetValue()) != 16 { 53 return nil, fmt.Errorf("invalid key") 54 } 55 return subtle.NewAESGCM(key.GetValue()) 56} 57 58// NewKey is only used by registry.NewKey, and that function is only used by KMSEnvelopeAEAD. 59// So there is no need to implement it. 60func (km *customKeyManager) NewKey(serializedKeyFormat []byte) (proto.Message, error) { 61 return nil, errors.New("not implemented") 62} 63 64func (km *customKeyManager) NewKeyData(serializedKeyFormat []byte) (*tinkpb.KeyData, error) { 65 keyFormat := new(wrapperspb.StringValue) 66 if err := proto.Unmarshal(serializedKeyFormat, keyFormat); err != nil { 67 return nil, fmt.Errorf("invalid key format") 68 } 69 if keyFormat.GetValue() != "AEAD_AES_GCM_128" { 70 return nil, fmt.Errorf("invalid key format") 71 } 72 keyValue := random.GetRandomBytes(16) 73 key := &wrapperspb.BytesValue{ 74 Value: keyValue, 75 } 76 serializedKey, err := proto.Marshal(key) 77 if err != nil { 78 return nil, err 79 } 80 return &tinkpb.KeyData{ 81 TypeUrl: customTypeURL, 82 Value: serializedKey, 83 KeyMaterialType: km.KeyMaterialType(), 84 }, nil 85} 86 87func (km *customKeyManager) DoesSupport(typeURL string) bool { 88 return typeURL == customTypeURL 89} 90 91func (km *customKeyManager) TypeURL() string { 92 return customTypeURL 93} 94 95func (km *customKeyManager) KeyMaterialType() tinkpb.KeyData_KeyMaterialType { 96 return tinkpb.KeyData_SYMMETRIC 97} 98 99// aesGCM128KeyTemplate creates a AES GCM 128 KeyTemplate for customKeyManager. 100func aesGCM128KeyTemplate() *tinkpb.KeyTemplate { 101 format := &wrapperspb.StringValue{ 102 Value: "AEAD_AES_GCM_128", 103 } 104 serializedFormat, err := proto.Marshal(format) 105 if err != nil { 106 tinkerror.Fail(fmt.Sprintf("failed to marshal key format: %s", err)) 107 } 108 return &tinkpb.KeyTemplate{ 109 TypeUrl: customTypeURL, 110 Value: serializedFormat, 111 OutputPrefixType: tinkpb.OutputPrefixType_RAW, 112 } 113} 114 115// aesGCM128KeyToKeysetHandle creates a keyset.Handle with one custom AES GCM 128 key. 116func aesGCM128KeyToKeysetHandle(rawAESKey []byte, keyID uint32, prefixType tinkpb.OutputPrefixType) (*keyset.Handle, error) { 117 if len(rawAESKey) != 16 { 118 return nil, fmt.Errorf("invalid key length") 119 } 120 key := &wrapperspb.BytesValue{Value: rawAESKey} 121 serializedKey, err := proto.Marshal(key) 122 if err != nil { 123 return nil, err 124 } 125 keyData := &tinkpb.KeyData{ 126 TypeUrl: customTypeURL, 127 Value: serializedKey, 128 KeyMaterialType: tinkpb.KeyData_SYMMETRIC, 129 } 130 ks := &tinkpb.Keyset{ 131 PrimaryKeyId: keyID, 132 Key: []*tinkpb.Keyset_Key{ 133 &tinkpb.Keyset_Key{ 134 KeyData: keyData, 135 Status: tinkpb.KeyStatusType_ENABLED, 136 KeyId: keyID, 137 OutputPrefixType: prefixType, 138 }, 139 }, 140 } 141 serializedKeyset, err := proto.Marshal(ks) 142 if err != nil { 143 return nil, err 144 } 145 return insecurecleartextkeyset.Read(keyset.NewBinaryReader(bytes.NewBuffer(serializedKeyset))) 146} 147 148func TestCreateEncryptDecrypt(t *testing.T) { 149 handle, err := keyset.NewHandle(aesGCM128KeyTemplate()) 150 if err != nil { 151 t.Fatalf("keyset.NewHandle(aesGCM128KeyTemplate()) err = %v, want nil", err) 152 } 153 primitive, err := aead.New(handle) 154 if err != nil { 155 t.Fatalf("aead.New(handle) err = %v, want nil", err) 156 } 157 158 plaintext := []byte("plaintext") 159 associatedData := []byte("associatedData") 160 161 ciphertext, err := primitive.Encrypt(plaintext, associatedData) 162 if err != nil { 163 t.Fatalf("primitive.Encrypt(plaintext, associatedData) err = %v, want nil", err) 164 } 165 decrypted, err := primitive.Decrypt(ciphertext, associatedData) 166 if err != nil { 167 t.Fatalf("primitive.Decrypt(ciphertext, associatedData) err = %v, want nil", err) 168 } 169 if !bytes.Equal(plaintext, decrypted) { 170 t.Errorf("primitive.Decrypt(ciphertext, associatedData) = %q, want: %q", decrypted, plaintext) 171 } 172} 173 174func TestImportExistingKeyDecryptsExistingCiphertext(t *testing.T) { 175 rawAesKey := random.GetRandomBytes(16) 176 plaintext := []byte("plaintext") 177 associatedData := []byte("associatedData") 178 179 // Create a AES GCM 128 ciphertext using rawAesKey. 180 aesGCMForRawAesKey, err := subtle.NewAESGCM(rawAesKey) 181 if err != nil { 182 t.Fatalf("subtle.NewAESGCM(rawAesKey) err = %v, want nil", err) 183 } 184 ciphertext, err := aesGCMForRawAesKey.Encrypt(plaintext, associatedData) 185 if err != nil { 186 t.Fatalf("aesGCMForRawAesKey.Encrypt(plaintext, associatedData) err = %v, want nil", err) 187 } 188 189 // Import rawAesKey into a Tink keyset.Handle, and decrypt the ciphertext. 190 handle, err := aesGCM128KeyToKeysetHandle(rawAesKey, 123, tinkpb.OutputPrefixType_RAW) 191 if err != nil { 192 t.Fatalf("aesGCM128KeyToKeysetHandle() err = %v, want nil", err) 193 } 194 primitive, err := aead.New(handle) 195 if err != nil { 196 t.Fatalf("aead.New(handle) err = %v, want nil", err) 197 } 198 gotPlaintext, err := primitive.Decrypt(ciphertext, associatedData) 199 if err != nil { 200 t.Fatalf("primitive.Decrypt(ciphertext, associatedData) err = %v, want nil", err) 201 } 202 if !bytes.Equal(plaintext, gotPlaintext) { 203 t.Fatalf("primitive.Decrypt(ciphertext, associatedData) = %q, want: %q", gotPlaintext, plaintext) 204 } 205} 206 207func TestEncryptAndDecryptWithTinkPrefix(t *testing.T) { 208 // Create an AEAD for rawAesKey with output prefix type TINK. 209 rawAesKey := random.GetRandomBytes(16) 210 handle, err := aesGCM128KeyToKeysetHandle(rawAesKey, 0x11223344, tinkpb.OutputPrefixType_TINK) 211 if err != nil { 212 t.Fatalf("aesGCM128KeyToKeysetHandle() err = %v, want nil", err) 213 } 214 primitive, err := aead.New(handle) 215 if err != nil { 216 t.Fatalf("aead.New(handle) err = %v, want nil", err) 217 } 218 219 // Encrypt and decrypt. 220 plaintext := []byte("plaintext") 221 associatedData := []byte("associatedData") 222 ciphertext, err := primitive.Encrypt(plaintext, associatedData) 223 if err != nil { 224 t.Fatalf("primitive.Encrypt(plaintext, associatedData) err = %v, want nil", err) 225 } 226 gotPlaintext, err := primitive.Decrypt(ciphertext, associatedData) 227 if err != nil { 228 t.Fatalf("primitive.Decrypt(ciphertext, associatedData) err = %v, want nil", err) 229 } 230 if !bytes.Equal(plaintext, gotPlaintext) { 231 t.Fatalf("primitive.Decrypt(ciphertext, associatedData) = %q, want: %q", gotPlaintext, plaintext) 232 } 233 234 // Check that ciphertext has the correct prefix. 235 gotPrefix := ciphertext[:5] 236 wantPrefix := []byte{0x01, 0x11, 0x22, 0x33, 0x44} 237 if !bytes.Equal(gotPrefix, wantPrefix) { 238 t.Fatalf("ciphertext[:5] = %q, want: %q", gotPrefix, wantPrefix) 239 } 240 241 // Check that subtle.NewAESGCM with rawAesKey can decrypt the ciphertext if the prefix is removed. 242 aesGCMForRawAesKey, err := subtle.NewAESGCM(rawAesKey) 243 if err != nil { 244 t.Fatalf("subtle.NewAESGCM(rawAesKey) err = %v, want nil", err) 245 } 246 gotPlaintext, err = aesGCMForRawAesKey.Decrypt(ciphertext[5:], associatedData) 247 if err != nil { 248 t.Fatalf("aesGCMForRawAesKey.Decrypt() err = %v, want nil", err) 249 } 250 if !bytes.Equal(plaintext, gotPlaintext) { 251 t.Fatalf("aesGCMForRawAesKey.Decrypt() = %q, want: %q", gotPlaintext, plaintext) 252 } 253} 254 255func TestMixedKeysetWorks(t *testing.T) { 256 rawAesKey := random.GetRandomBytes(16) 257 258 // Create a AES GCM 128 ciphertext using rawAesKey. 259 subtlePrimitive, err := subtle.NewAESGCM(rawAesKey) 260 if err != nil { 261 t.Fatalf("subtle.NewAESGCM(rawAesKey) err = %v, want nil", err) 262 } 263 plaintext := []byte("plaintext") 264 associatedData := []byte("associatedData") 265 ciphertext, err := subtlePrimitive.Encrypt(plaintext, associatedData) 266 if err != nil { 267 t.Fatalf("subtlePrimitive.Encrypt(plaintext, associatedData) err = %v, want nil", err) 268 } 269 270 // Create handle2, which is a keyset.Handle that contains a customKeyManager key of rawAesKey and 271 // a new, non-customKeyManager key. 272 handle1, err := aesGCM128KeyToKeysetHandle(rawAesKey, 123, tinkpb.OutputPrefixType_RAW) 273 if err != nil { 274 t.Fatalf("aesGCM128KeyToKeysetHandle() err = %v, want nil", err) 275 } 276 manager := keyset.NewManagerFromHandle(handle1) 277 keyID, err := manager.Add(aead.AES128CTRHMACSHA256KeyTemplate()) 278 if err != nil { 279 t.Fatalf("manager.Add(aead.AES128CTRHMACSHA256KeyTemplate()) err = %v, want nil", err) 280 } 281 err = manager.SetPrimary(keyID) 282 if err != nil { 283 t.Fatalf("manager.SetPrimary(keyID) = %v", err) 284 } 285 handle2, err := manager.Handle() 286 if err != nil { 287 t.Fatalf("manager.Handle() err = %v", err) 288 } 289 290 primitive, err := aead.New(handle2) 291 if err != nil { 292 t.Fatalf("aead.New(handle2) err = %v", err) 293 } 294 gotPlaintext, err := primitive.Decrypt(ciphertext, associatedData) 295 if err != nil { 296 t.Fatalf("primitive.Decrypt(ciphertext, associatedData) err = %v, want nil", err) 297 } 298 if !bytes.Equal(plaintext, gotPlaintext) { 299 t.Errorf("primitive.Decrypt(ciphertext, associatedData) = %q, want: %q", gotPlaintext, plaintext) 300 } 301} 302 303func TestSerializeAndParseKeysetWorks(t *testing.T) { 304 handle, err := keyset.NewHandle(aesGCM128KeyTemplate()) 305 if err != nil { 306 t.Fatalf("keyset.NewHandle(aesGCM128KeyTemplate()) err = %v, want nil", err) 307 } 308 primitive, err := aead.New(handle) 309 if err != nil { 310 t.Fatalf("aead.New(handle) err = %v, want nil", err) 311 } 312 313 plaintext := []byte("plaintext") 314 associatedData := []byte("associatedData") 315 ciphertext, err := primitive.Encrypt(plaintext, associatedData) 316 if err != nil { 317 t.Fatalf("primitive.Encrypt(plaintext, associatedData) err = %v, want nil", err) 318 } 319 320 // Serialize the keyset. 321 buff := &bytes.Buffer{} 322 err = insecurecleartextkeyset.Write(handle, keyset.NewBinaryWriter(buff)) 323 if err != nil { 324 t.Fatalf("insecurecleartextkeyset.Write(handle, keyset.NewBinaryWriter(buff)) = %v, want nil", err) 325 } 326 serializedKeyset := buff.Bytes() 327 328 // Parse the keyset. 329 parsedHandle, err := insecurecleartextkeyset.Read( 330 keyset.NewBinaryReader(bytes.NewBuffer(serializedKeyset))) 331 if err != nil { 332 t.Fatalf("insecurecleartextkeyset.Read(keyset.NewBinaryReader(bytes.NewBuffer(serializedKeyset))) = %v, want nil", err) 333 } 334 335 primitive2, err := aead.New(parsedHandle) 336 if err != nil { 337 t.Fatalf("aead.New(parsedHandle) err = %v, want nil", err) 338 } 339 340 gotPlaintext, err := primitive2.Decrypt(ciphertext, associatedData) 341 if err != nil { 342 t.Fatalf("primitive2.Decrypt(ciphertext, associatedData) err = %v, want nil", err) 343 } 344 if !bytes.Equal(plaintext, gotPlaintext) { 345 t.Errorf("primitive2.Decrypt(ciphertext, associatedData) = %q, want: %q", gotPlaintext, plaintext) 346 } 347} 348 349func init() { registry.RegisterKeyManager(&customKeyManager{}) } 350