1 files changed, 61 insertions, 7 deletions

diff --git a/security-utils/src/test/java/com/yahoo/security/SharedKeyTest.java b/security-utils/src/test/java/com/yahoo/security/SharedKeyTest.java
index 25324ad7317..35b52d13b1d 100644
--- a/security-utils/src/test/java/com/yahoo/security/SharedKeyTest.java
+++ b/security-utils/src/test/java/com/yahoo/security/SharedKeyTest.java
@@ -59,20 +59,74 @@ public class SharedKeyTest {
     }
 
     @Test
-    void token_v1_representation_is_stable() {
+    void resealed_token_preserves_token_version_of_source_token() {
+        var originalPrivate = KeyUtils.fromBase58EncodedX25519PrivateKey("GFg54SaGNCmcSGufZCx68SKLGuAFrASoDeMk3t5AjU6L");
+        var v1Token         = "OntP9gRVAjXeZIr4zkYqRJFcnA993v7ZEE7VbcNs1NcR3HdE7Mpwlwi3r3anF1kVa5fn7O1CyeHQpBWpdayUTKkrtyFepG6WJrZdE";
+
+        var originalSealed = SealedSharedKey.fromTokenString(v1Token);
+        var originalSecret = SharedKeyGenerator.fromSealedKey(originalSealed, originalPrivate);
+
+        var secondaryReceiverKp = KeyUtils.generateX25519KeyPair();
+        var resealedShared = SharedKeyGenerator.reseal(originalSecret, secondaryReceiverKp.getPublic(), KEY_ID_2);
+
+        var theirSealed = SealedSharedKey.fromTokenString(resealedShared.sealedSharedKey().toTokenString());
+        assertEquals(1, theirSealed.tokenVersion());
+    }
+
+    @Test
+    void token_v1_representation_is_stable() throws IOException {
         var receiverPrivate = KeyUtils.fromBase58EncodedX25519PrivateKey("GFg54SaGNCmcSGufZCx68SKLGuAFrASoDeMk3t5AjU6L");
         var receiverPublic  = KeyUtils.fromBase58EncodedX25519PublicKey( "5drrkakYLjYSBpr5Haknh13EiCYL36ndMzK4gTJo6pwh");
         var keyId           = KeyId.ofString("my key ID");
 
-        // Token generated for the above receiver public key, with the below expected shared secret (in hex)
+        // V1 token generated for the above receiver public key, with the below expected shared secret (in hex)
         var publicToken = "OntP9gRVAjXeZIr4zkYqRJFcnA993v7ZEE7VbcNs1NcR3HdE7Mpwlwi3r3anF1kVa5fn7O1CyeHQpBWpdayUTKkrtyFepG6WJrZdE";
         var expectedSharedSecret = "1b33b4dcd6a94e5a4a1ee6d208197d01";
 
         var theirSealed = SealedSharedKey.fromTokenString(publicToken);
         var theirShared = SharedKeyGenerator.fromSealedKey(theirSealed, receiverPrivate);
 
+        assertEquals(1, theirSealed.tokenVersion());
+        assertEquals(keyId, theirSealed.keyId());
+        assertEquals(expectedSharedSecret, hex(theirShared.secretKey().getEncoded()));
+
+        // Encryption with v1 tokens must use AES-GCM 128
+        var plaintext = "it's Bocchi time";
+        var expectedCiphertext = "a2ba842b2e0769a4a2948c4236d4ae921f1dd05c2e094dcde9699eeefcc3d7ae";
+        byte[] ct = streamEncryptString(plaintext, theirShared);
+        assertEquals(expectedCiphertext, hex(ct));
+
+        // Decryption with v1 tokens must use AES-GCM 128
+        var decrypted = streamDecryptString(ct, theirShared);
+        assertEquals(plaintext, decrypted);
+    }
+
+    @Test
+    void token_v2_representation_is_stable() throws IOException {
+        var receiverPrivate = KeyUtils.fromBase58EncodedX25519PrivateKey("GFg54SaGNCmcSGufZCx68SKLGuAFrASoDeMk3t5AjU6L");
+        var receiverPublic  = KeyUtils.fromBase58EncodedX25519PublicKey( "5drrkakYLjYSBpr5Haknh13EiCYL36ndMzK4gTJo6pwh");
+        var keyId           = KeyId.ofString("my key ID");
+
+        // V2 token generated for the above receiver public key, with the below expected shared secret (in hex)
+        var publicToken = "mjA83HYuulZW5SWV8FKz4m3b3m9zU8mTrX9n6iY4wZaA6ZNr8WnBZwOU4KQqhPCORPlzSYk4svlonzPZIb3Bjbqr2ePYKLOpdGhCO";
+        var expectedSharedSecret = "205af82154690fd7b6d56a977563822c";
+
+        var theirSealed = SealedSharedKey.fromTokenString(publicToken);
+        var theirShared = SharedKeyGenerator.fromSealedKey(theirSealed, receiverPrivate);
+
+        assertEquals(2, theirSealed.tokenVersion());
         assertEquals(keyId, theirSealed.keyId());
         assertEquals(expectedSharedSecret, hex(theirShared.secretKey().getEncoded()));
+
+        // Encryption with v2 tokens must use ChaCha20-Poly1305
+        var plaintext = "it's Bocchi time";
+        var expectedCiphertext = "ea19dd0ac3ea6d76dc4e96430b0d5902a21cb3a27fa99490f4dcc391eaf5cec4";
+        byte[] ct = streamEncryptString(plaintext, theirShared);
+        assertEquals(expectedCiphertext, hex(ct));
+
+        // Decryption with v2 tokens must use ChaCha20-Poly1305
+        var decrypted = streamDecryptString(ct, theirShared);
+        assertEquals(plaintext, decrypted);
     }
 
     @Test
@@ -102,7 +156,7 @@ public class SharedKeyTest {
         var mySealed = myShared.sealedSharedKey();
         var badId    = KeyId.ofString("my key 2");
 
-        var tamperedShared = new SealedSharedKey(badId, mySealed.enc(), mySealed.ciphertext());
+        var tamperedShared = new SealedSharedKey(SealedSharedKey.CURRENT_TOKEN_VERSION, badId, mySealed.enc(), mySealed.ciphertext());
         // Should not be able to unseal the token since the AAD auth tag won't be correct
         assertThrows(RuntimeException.class, // TODO consider distinct exception class
                      () -> SharedKeyGenerator.fromSealedKey(tamperedShared, keyPair.getPrivate()));
@@ -130,7 +184,7 @@ public class SharedKeyTest {
         var myShared = SharedKeyGenerator.generateForReceiverPublicKey(keyPair.getPublic(), goodId);
 
         // token header is u8 version || u8 key id length || key id bytes ...
-        // Since the key ID is only 1 bytes long, patch it with a bad UTF-8 value
+        // Since the key ID is only 1 byte long, patch it with a bad UTF-8 value
         byte[] tokenBytes = Base62.codec().decode(myShared.sealedSharedKey().toTokenString());
         tokenBytes[2] = (byte)0xC0; // First part of a 2-byte continuation without trailing byte
         var patchedTokenStr = Base62.codec().encode(tokenBytes);
@@ -138,7 +192,7 @@ public class SharedKeyTest {
     }
 
     static byte[] streamEncryptString(String data, SecretSharedKey secretSharedKey) throws IOException {
-        var cipher = SharedKeyGenerator.makeAesGcmEncryptionCipher(secretSharedKey);
+        var cipher = secretSharedKey.makeEncryptionCipher();
         var outStream = new ByteArrayOutputStream();
         try (var cipherStream = cipher.wrapOutputStream(outStream)) {
             cipherStream.write(data.getBytes(StandardCharsets.UTF_8));
@@ -148,7 +202,7 @@ public class SharedKeyTest {
     }
 
     static String streamDecryptString(byte[] encrypted, SecretSharedKey secretSharedKey) throws IOException {
-        var cipher   = SharedKeyGenerator.makeAesGcmDecryptionCipher(secretSharedKey);
+        var cipher   = secretSharedKey.makeDecryptionCipher();
         var inStream = new ByteArrayInputStream(encrypted);
         var total    = ByteBuffer.allocate(encrypted.length); // Assume decrypted form can't be _longer_
         byte[] tmp   = new byte[8]; // short buf to test chunking
@@ -198,7 +252,7 @@ public class SharedKeyTest {
     }
 
     private static void doOutputStreamCipherDecrypt(SecretSharedKey myShared, byte[] encrypted) throws Exception {
-        var cipher = SharedKeyGenerator.makeAesGcmDecryptionCipher(myShared);
+        var cipher = myShared.makeDecryptionCipher();
         var outStream = new ByteArrayOutputStream();
         try (var cipherStream = cipher.wrapOutputStream(outStream)) {
             cipherStream.write(encrypted);