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 package com.google.crypto.tink.custom; 18 19 import static com.google.common.truth.Truth.assertThat; 20 import static java.nio.charset.StandardCharsets.UTF_8; 21 22 import com.google.crypto.tink.Aead; 23 import com.google.crypto.tink.CleartextKeysetHandle; 24 import com.google.crypto.tink.InsecureSecretKeyAccess; 25 import com.google.crypto.tink.KeyManager; 26 import com.google.crypto.tink.KeysetHandle; 27 import com.google.crypto.tink.Parameters; 28 import com.google.crypto.tink.Registry; 29 import com.google.crypto.tink.TinkProtoKeysetFormat; 30 import com.google.crypto.tink.TinkProtoParametersFormat; 31 import com.google.crypto.tink.aead.AeadConfig; 32 import com.google.crypto.tink.aead.ChaCha20Poly1305Parameters; 33 import com.google.crypto.tink.proto.KeyData; 34 import com.google.crypto.tink.proto.KeyData.KeyMaterialType; 35 import com.google.crypto.tink.proto.KeyStatusType; 36 import com.google.crypto.tink.proto.KeyTemplate; 37 import com.google.crypto.tink.proto.Keyset; 38 import com.google.crypto.tink.proto.OutputPrefixType; 39 import com.google.crypto.tink.subtle.AesGcmJce; 40 import com.google.crypto.tink.subtle.Random; 41 import com.google.protobuf.ByteString; 42 import com.google.protobuf.BytesValue; 43 import com.google.protobuf.ExtensionRegistryLite; 44 import com.google.protobuf.InvalidProtocolBufferException; 45 import com.google.protobuf.StringValue; 46 import java.security.GeneralSecurityException; 47 import java.util.Arrays; 48 import org.junit.BeforeClass; 49 import org.junit.Test; 50 import org.junit.runner.RunWith; 51 import org.junit.runners.JUnit4; 52 53 /** This test creates a custom Aead KeyManager and uses it. */ 54 @RunWith(JUnit4.class) 55 public final class CustomAeadKeyManagerTest { 56 57 /** 58 * A custom implementation of {@link com.google.crypto.tink.KeyManager} for AES GCM 128. 59 * 60 * <p>It only implements the methods of the KeyManager interface that are needed. 61 */ 62 static class MyCustomKeyManager implements KeyManager<Aead> { 63 64 private static final String TYPE_URL = 65 "type.googleapis.com/google.crypto.tink.testonly.CustomAeadKey"; 66 67 private static final String AEAD_AES_128_GCM = "AEAD_AES_128_GCM"; 68 69 @Override getPrimitive(ByteString serializedKey)70 public Aead getPrimitive(ByteString serializedKey) throws GeneralSecurityException { 71 try { 72 BytesValue key = 73 BytesValue.parseFrom(serializedKey, ExtensionRegistryLite.getEmptyRegistry()); 74 byte[] keyValue = key.getValue().toByteArray(); 75 if (keyValue.length != 16) { 76 throw new GeneralSecurityException("unexpected length of keyValue"); 77 } 78 return new AesGcmJce(keyValue); 79 } catch (InvalidProtocolBufferException e) { 80 throw new GeneralSecurityException(e); 81 } 82 } 83 84 @Override getKeyType()85 public String getKeyType() { 86 return TYPE_URL; 87 } 88 89 @Override getPrimitiveClass()90 public Class<Aead> getPrimitiveClass() { 91 return Aead.class; 92 } 93 94 @Override newKeyData(ByteString serializedKeyFormat)95 public KeyData newKeyData(ByteString serializedKeyFormat) throws GeneralSecurityException { 96 // serializedKeyFormat is a StringValue proto. The only allowed string is "AEAD_AES_128_GCM". 97 try { 98 StringValue keyFormat = 99 StringValue.parseFrom(serializedKeyFormat, ExtensionRegistryLite.getEmptyRegistry()); 100 if (!keyFormat.getValue().equals(AEAD_AES_128_GCM)) { 101 throw new GeneralSecurityException("unknown algorithm"); 102 } 103 byte[] rawAesKey = Random.randBytes(16); 104 BytesValue value = BytesValue.of(ByteString.copyFrom(rawAesKey)); 105 return KeyData.newBuilder() 106 .setTypeUrl(getKeyType()) 107 .setValue(value.toByteString()) 108 .setKeyMaterialType(KeyMaterialType.SYMMETRIC) 109 .build(); 110 } catch (InvalidProtocolBufferException e) { 111 throw new GeneralSecurityException(e); 112 } 113 } 114 aesGcm128Parameters()115 static Parameters aesGcm128Parameters() throws GeneralSecurityException { 116 StringValue format = StringValue.of(AEAD_AES_128_GCM); 117 KeyTemplate template = 118 KeyTemplate.newBuilder() 119 .setValue(format.toByteString()) 120 .setTypeUrl(TYPE_URL) 121 .setOutputPrefixType(OutputPrefixType.RAW) 122 .build(); 123 return TinkProtoParametersFormat.parse(template.toByteArray()); 124 } 125 aesGcm128KeyToKeysetHandle( byte[] rawAesKey, int keyId, OutputPrefixType outputPrefixType)126 static KeysetHandle aesGcm128KeyToKeysetHandle( 127 byte[] rawAesKey, int keyId, OutputPrefixType outputPrefixType) 128 throws GeneralSecurityException { 129 if (rawAesKey.length != 16) { 130 throw new IllegalArgumentException("unexpected raw key length"); 131 } 132 BytesValue value = BytesValue.of(ByteString.copyFrom(rawAesKey)); 133 Keyset keyset = 134 Keyset.newBuilder() 135 .addKey( 136 Keyset.Key.newBuilder() 137 .setStatus(KeyStatusType.ENABLED) 138 .setOutputPrefixType(outputPrefixType) 139 .setKeyId(keyId) 140 .setKeyData( 141 KeyData.newBuilder() 142 .setTypeUrl(TYPE_URL) 143 .setValue(value.toByteString()) 144 .setKeyMaterialType(KeyMaterialType.SYMMETRIC) 145 .build()) 146 .build()) 147 .setPrimaryKeyId(keyId) 148 .build(); 149 return CleartextKeysetHandle.fromKeyset(keyset); 150 } 151 } 152 153 @BeforeClass setUpClass()154 public static void setUpClass() throws Exception { 155 AeadConfig.register(); 156 Registry.registerKeyManager(new MyCustomKeyManager(), /* newKeyAllowed= */ true); 157 } 158 159 @Test createEncryptAndDecrypt_success()160 public void createEncryptAndDecrypt_success() throws Exception { 161 Parameters aesGcm128Parameters = MyCustomKeyManager.aesGcm128Parameters(); 162 KeysetHandle handle = KeysetHandle.generateNew(aesGcm128Parameters); 163 Aead aead = handle.getPrimitive(Aead.class); 164 165 byte[] plaintext = "plaintext".getBytes(UTF_8); 166 byte[] associatedData = "associatedData".getBytes(UTF_8); 167 byte[] ciphertext = aead.encrypt(plaintext, associatedData); 168 byte[] decrypted = aead.decrypt(ciphertext, associatedData); 169 assertThat(decrypted).isEqualTo(plaintext); 170 } 171 172 @Test importExistingKey_decrypts()173 public void importExistingKey_decrypts() throws Exception { 174 byte[] rawAesKey = Random.randBytes(16); 175 Aead jceAead = new AesGcmJce(rawAesKey); 176 byte[] plaintext = "plaintext".getBytes(UTF_8); 177 byte[] associatedData = "associatedData".getBytes(UTF_8); 178 byte[] ciphertext = jceAead.encrypt(plaintext, associatedData); 179 180 KeysetHandle handle = 181 MyCustomKeyManager.aesGcm128KeyToKeysetHandle( 182 rawAesKey, /* keyId= */ 0x11223344, OutputPrefixType.RAW); 183 Aead aead = handle.getPrimitive(Aead.class); 184 byte[] decrypted = aead.decrypt(ciphertext, associatedData); 185 assertThat(decrypted).isEqualTo(plaintext); 186 } 187 188 @Test encryptAndDecryptWithTinkPrefix_success()189 public void encryptAndDecryptWithTinkPrefix_success() throws Exception { 190 // Create a new key and import it with output prefix type TINK with a fixed key ID. 191 byte[] rawAesKey = Random.randBytes(16); 192 int keyId = 0x11223344; 193 KeysetHandle handle = 194 MyCustomKeyManager.aesGcm128KeyToKeysetHandle(rawAesKey, keyId, OutputPrefixType.TINK); 195 196 Aead aead = handle.getPrimitive(Aead.class); 197 byte[] plaintext = "plaintext".getBytes(UTF_8); 198 byte[] associatedData = "associatedData".getBytes(UTF_8); 199 byte[] ciphertext = aead.encrypt(plaintext, associatedData); 200 assertThat(aead.decrypt(ciphertext, associatedData)).isEqualTo(plaintext); 201 202 // Check that ciphertext generated using OutputPrefixType.TINK has a 5 byte prefix: 203 // the first byte is always 0x01, and the next 4 bytes are the big-endian encoded key ID. 204 byte[] prefix = Arrays.copyOf(ciphertext, 5); 205 assertThat(prefix) 206 .isEqualTo(new byte[] {(byte) 0x01, (byte) 0x11, (byte) 0x22, (byte) 0x33, (byte) 0x44}); 207 208 // Check that AesGcmJce can decrypt using the raw key, if the prefix is removed. 209 byte[] ciphertextWithoutPrefix = Arrays.copyOfRange(ciphertext, 5, ciphertext.length); 210 Aead jceAead = new AesGcmJce(rawAesKey); 211 assertThat(jceAead.decrypt(ciphertextWithoutPrefix, associatedData)).isEqualTo(plaintext); 212 } 213 214 @Test keysetWithCustomAndTinkKeys_decrypts()215 public void keysetWithCustomAndTinkKeys_decrypts() throws Exception { 216 byte[] rawAesKey = Random.randBytes(16); 217 Aead jceAead = new AesGcmJce(rawAesKey); 218 byte[] plaintext = "plaintext".getBytes(UTF_8); 219 byte[] associatedData = "associatedData".getBytes(UTF_8); 220 byte[] ciphertext = jceAead.encrypt(plaintext, associatedData); 221 222 // Create keyset handle with normal Tink key 223 KeysetHandle handleWithTinkKey = 224 KeysetHandle.generateNew( 225 ChaCha20Poly1305Parameters.create(ChaCha20Poly1305Parameters.Variant.TINK)); 226 Aead aead2 = handleWithTinkKey.getPrimitive(Aead.class); 227 byte[] ciphertext2 = aead2.encrypt(plaintext, associatedData); 228 229 KeysetHandle handle = 230 MyCustomKeyManager.aesGcm128KeyToKeysetHandle( 231 rawAesKey, /* keyId= */ 0x11223344, OutputPrefixType.RAW); 232 // Create keyset handle with both the custom key and the normal Tink key 233 KeysetHandle handle2 = 234 KeysetHandle.newBuilder(handle) 235 .addEntry(KeysetHandle.importKey(handleWithTinkKey.getAt(0).getKey()).makePrimary()) 236 .build(); 237 238 // Decrypt both ciphertexts 239 Aead aead = handle2.getPrimitive(Aead.class); 240 assertThat(aead.decrypt(ciphertext, associatedData)).isEqualTo(plaintext); 241 assertThat(aead.decrypt(ciphertext2, associatedData)).isEqualTo(plaintext); 242 } 243 244 @Test serializeAndParse_decrypts()245 public void serializeAndParse_decrypts() throws Exception { 246 Parameters aesGcm128Parameters = MyCustomKeyManager.aesGcm128Parameters(); 247 KeysetHandle handle = KeysetHandle.generateNew(aesGcm128Parameters); 248 Aead aead = handle.getPrimitive(Aead.class); 249 byte[] plaintext = "plaintext".getBytes(UTF_8); 250 byte[] associatedData = "associatedData".getBytes(UTF_8); 251 byte[] ciphertext = aead.encrypt(plaintext, associatedData); 252 253 byte[] serializedKeyset = 254 TinkProtoKeysetFormat.serializeKeyset(handle, InsecureSecretKeyAccess.get()); 255 256 KeysetHandle handle2 = 257 TinkProtoKeysetFormat.parseKeyset(serializedKeyset, InsecureSecretKeyAccess.get()); 258 Aead aead2 = handle2.getPrimitive(Aead.class); 259 byte[] decrypted = aead2.decrypt(ciphertext, associatedData); 260 assertThat(decrypted).isEqualTo(plaintext); 261 } 262 } 263