Encapsulate key identifier in own object

Enforces invariants and avoids having to pass raw byte arrays around.

5 files changed, 205 insertions, 60 deletions

diff --git a/security-utils/src/main/java/com/yahoo/security/KeyId.java b/security-utils/src/main/java/com/yahoo/security/KeyId.java
new file mode 100644
index 00000000000..472651fa1bb
--- /dev/null
+++ b/security-utils/src/main/java/com/yahoo/security/KeyId.java
@@ -0,0 +1,89 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.security;
+
+import java.util.Arrays;
+import java.util.Objects;
+
+import static com.yahoo.security.ArrayUtils.fromUtf8Bytes;
+import static com.yahoo.security.ArrayUtils.toUtf8Bytes;
+
+/**
+ * Represents a named key ID comprising an arbitrary (but length-limited)
+ * sequence of valid UTF-8 bytes.
+ *
+ * @author vekterli
+ */
+public class KeyId {
+
+    // Max length MUST be possible to fit in an unsigned byte; see SealedSharedKey token encoding/decoding.
+    public static final int MAX_KEY_ID_UTF8_LENGTH = 255;
+
+    private final byte[] keyIdBytes;
+
+    private KeyId(byte[] keyIdBytes) {
+        if (keyIdBytes.length > MAX_KEY_ID_UTF8_LENGTH) {
+            throw new IllegalArgumentException("Key ID is too large to be encoded (max is %d, got %d)"
+                                               .formatted(MAX_KEY_ID_UTF8_LENGTH, keyIdBytes.length));
+        }
+        verifyByteStringRoundtripsAsValidUtf8(keyIdBytes);
+        this.keyIdBytes = keyIdBytes;
+    }
+
+    /**
+     * Construct a KeyId containing the given sequence of bytes.
+     *
+     * @param keyIdBytes array of valid UTF-8 bytes. May be zero-length, but not null.
+     *                   Note: to avoid accidental mutations, the key bytes are deep-copied.
+     * @return a new KeyId instance
+     */
+    public static KeyId ofBytes(byte[] keyIdBytes) {
+        Objects.requireNonNull(keyIdBytes);
+        return new KeyId(keyIdBytes.clone());
+    }
+
+    /**
+     * Construct a KeyId containing the UTF-8 byte representation of the given string.
+     *
+     * @param keyId a string whose UTF-8 byte representation will be the key ID. May be
+     *              zero-length but not null.
+     * @return a new KeyId instance
+     */
+    public static KeyId ofString(String keyId) {
+        Objects.requireNonNull(keyId);
+        return new KeyId(toUtf8Bytes(keyId));
+    }
+
+    /**
+     * @return the raw backing byte array. <strong>Must therefore not be mutated.</strong>
+     */
+    public byte[] asBytes() { return keyIdBytes; }
+
+    public String asString() { return fromUtf8Bytes(keyIdBytes); }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        KeyId keyId = (KeyId) o;
+        return Arrays.equals(keyIdBytes, keyId.keyIdBytes);
+    }
+
+    @Override
+    public int hashCode() {
+        return Arrays.hashCode(keyIdBytes);
+    }
+
+    @Override
+    public String toString() {
+        return "KeyId(%s)".formatted(asString());
+    }
+
+    private static void verifyByteStringRoundtripsAsValidUtf8(byte[] byteStr) {
+        String asStr   = fromUtf8Bytes(byteStr); // Replaces bad chars with a placeholder
+        byte[] asBytes = toUtf8Bytes(asStr);
+        if (!Arrays.equals(byteStr, asBytes)) {
+            throw new IllegalArgumentException("Key ID is not valid normalized UTF-8");
+        }
+    }
+
+}
diff --git a/security-utils/src/main/java/com/yahoo/security/SealedSharedKey.java b/security-utils/src/main/java/com/yahoo/security/SealedSharedKey.java
index fe1be85539c..a921b3baf87 100644
--- a/security-utils/src/main/java/com/yahoo/security/SealedSharedKey.java
+++ b/security-utils/src/main/java/com/yahoo/security/SealedSharedKey.java
@@ -19,20 +19,14 @@ import static com.yahoo.security.ArrayUtils.toUtf8Bytes;
  * This token representation is expected to be used as a convenient serialization
  * form when communicating shared keys.
  */
-public record SealedSharedKey(byte[] keyId, byte[] enc, byte[] ciphertext) {
+public record SealedSharedKey(KeyId keyId, byte[] enc, byte[] ciphertext) {
 
     /** Current encoding version of opaque sealed key tokens. Must be less than 256. */
     public static final int CURRENT_TOKEN_VERSION = 1;
-    public static final int MAX_KEY_ID_UTF8_LENGTH = 255;
     /** Encryption context for v1 tokens is always a 32-byte X25519 public key */
     public static final int MAX_ENC_CONTEXT_LENGTH = 255;
 
     public SealedSharedKey {
-        if (keyId.length > MAX_KEY_ID_UTF8_LENGTH) {
-            throw new IllegalArgumentException("Key ID is too large to be encoded (max is %d, got %d)"
-                    .formatted(MAX_KEY_ID_UTF8_LENGTH, keyId.length));
-        }
-        verifyByteStringRoundtripsAsValidUtf8(keyId);
         if (enc.length > MAX_ENC_CONTEXT_LENGTH) {
             throw new IllegalArgumentException("Encryption context is too large to be encoded (max is %d, got %d)"
                     .formatted(MAX_ENC_CONTEXT_LENGTH, enc.length));
@@ -44,11 +38,12 @@ public record SealedSharedKey(byte[] keyId, byte[] enc, byte[] ciphertext) {
      * reconstruct the SealedSharedKey instance when passed verbatim to fromTokenString().
      */
     public String toTokenString() {
+        byte[] keyIdBytes = keyId.asBytes();
         // u8 token version || u8 length(key id) || key id || u8 length(enc) || enc || ciphertext
-        ByteBuffer encoded = ByteBuffer.allocate(1 + 1 + keyId.length + 1 + enc.length + ciphertext.length);
+        ByteBuffer encoded = ByteBuffer.allocate(1 + 1 + keyIdBytes.length + 1 + enc.length + ciphertext.length);
         encoded.put((byte)CURRENT_TOKEN_VERSION);
-        encoded.put((byte)keyId.length);
-        encoded.put(keyId);
+        encoded.put((byte)keyIdBytes.length);
+        encoded.put(keyIdBytes);
         encoded.put((byte)enc.length);
         encoded.put(enc);
         encoded.put(ciphertext);
@@ -76,24 +71,15 @@ public record SealedSharedKey(byte[] keyId, byte[] enc, byte[] ciphertext) {
                                                .formatted(CURRENT_TOKEN_VERSION, version));
         }
         int keyIdLen = Byte.toUnsignedInt(decoded.get());
-        byte[] keyId = new byte[keyIdLen];
-        decoded.get(keyId);
-        verifyByteStringRoundtripsAsValidUtf8(keyId);
+        byte[] keyIdBytes = new byte[keyIdLen];
+        decoded.get(keyIdBytes);
         int encLen = Byte.toUnsignedInt(decoded.get());
         byte[] enc = new byte[encLen];
         decoded.get(enc);
         byte[] ciphertext = new byte[decoded.remaining()];
         decoded.get(ciphertext);
 
-        return new SealedSharedKey(keyId, enc, ciphertext);
-    }
-
-    private static void verifyByteStringRoundtripsAsValidUtf8(byte[] byteStr) {
-        String asStr   = fromUtf8Bytes(byteStr); // Replaces bad chars with a placeholder
-        byte[] asBytes = toUtf8Bytes(asStr);
-        if (!Arrays.equals(byteStr, asBytes)) {
-            throw new IllegalArgumentException("Key ID is not valid normalized UTF-8");
-        }
+        return new SealedSharedKey(KeyId.ofBytes(keyIdBytes), enc, ciphertext);
     }
 
     public int tokenVersion() { return CURRENT_TOKEN_VERSION; }
diff --git a/security-utils/src/main/java/com/yahoo/security/SharedKeyGenerator.java b/security-utils/src/main/java/com/yahoo/security/SharedKeyGenerator.java
index 47936dab114..8a1a7dd3688 100644
--- a/security-utils/src/main/java/com/yahoo/security/SharedKeyGenerator.java
+++ b/security-utils/src/main/java/com/yahoo/security/SharedKeyGenerator.java
@@ -60,17 +60,17 @@ public class SharedKeyGenerator {
         }
     }
 
-    public static SecretSharedKey generateForReceiverPublicKey(PublicKey receiverPublicKey, byte[] keyId) {
+    public static SecretSharedKey generateForReceiverPublicKey(PublicKey receiverPublicKey, KeyId keyId) {
         var secretKey = generateRandomSecretAesKey();
         // We protect the integrity of the key ID by passing it as AAD.
-        var sealed = HPKE.sealBase((XECPublicKey) receiverPublicKey, EMPTY_BYTES, keyId, secretKey.getEncoded());
+        var sealed = HPKE.sealBase((XECPublicKey) receiverPublicKey, EMPTY_BYTES, keyId.asBytes(), secretKey.getEncoded());
         var sealedSharedKey = new SealedSharedKey(keyId, sealed.enc(), sealed.ciphertext());
         return new SecretSharedKey(secretKey, sealedSharedKey);
     }
 
     public static SecretSharedKey fromSealedKey(SealedSharedKey sealedKey, PrivateKey receiverPrivateKey) {
         byte[] secretKeyBytes = HPKE.openBase(sealedKey.enc(), (XECPrivateKey) receiverPrivateKey,
-                                              EMPTY_BYTES, sealedKey.keyId(), sealedKey.ciphertext());
+                                              EMPTY_BYTES, sealedKey.keyId().asBytes(), sealedKey.ciphertext());
         return new SecretSharedKey(new SecretKeySpec(secretKeyBytes, "AES"), sealedKey);
     }
 
diff --git a/security-utils/src/test/java/com/yahoo/security/KeyIdTest.java b/security-utils/src/test/java/com/yahoo/security/KeyIdTest.java
new file mode 100644
index 00000000000..d5ed71be5e6
--- /dev/null
+++ b/security-utils/src/test/java/com/yahoo/security/KeyIdTest.java
@@ -0,0 +1,82 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.security;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+/**
+ * @author vekterli
+ */
+public class KeyIdTest {
+
+    @Test
+    void equality_predicated_on_key_id_byte_string() {
+        var id0s = KeyId.ofString("");
+        var id1s = KeyId.ofString("1");
+        var id2s = KeyId.ofString("12");
+        assertEquals(id0s, id0s);
+        assertEquals(id1s, id1s);
+        assertEquals(id2s, id2s);
+        assertNotEquals(id0s, id1s);
+        assertNotEquals(id1s, id0s);
+        assertNotEquals(id1s, id2s);
+        assertNotEquals(id0s, id2s);
+        var id0b = KeyId.ofBytes(new byte[0]);
+        var id1b = KeyId.ofBytes(new byte[]{ '1' });
+        var id2b = KeyId.ofBytes(new byte[]{ '1', '2' });
+        assertEquals(id0s, id0b);
+        assertEquals(id1s, id1b);
+        assertEquals(id2s, id2b);
+    }
+
+    @Test
+    void accessors_return_expected_values() {
+        byte[] fooBytes = new byte[]{'f','o','o'};
+        byte[] barBytes = new byte[]{'b','a','r'};
+
+        var id1 = KeyId.ofString("foo");
+        assertEquals("foo", id1.asString());
+        assertArrayEquals(fooBytes, id1.asBytes());
+
+        var id2 = KeyId.ofBytes(barBytes);
+        assertEquals("bar", id2.asString());
+        assertArrayEquals(barBytes, id2.asBytes());
+    }
+
+    @Test
+    void key_id_bytes_are_deep_copied_when_constructed_from_raw_byte_array() {
+        byte[] keyBytes = new byte[]{'f','o','o'};
+        byte[] expected = keyBytes.clone();
+        var id = KeyId.ofBytes(keyBytes);
+        keyBytes[0] = 'b';
+        assertArrayEquals(expected, id.asBytes());
+    }
+
+    @Test
+    void can_construct_largest_possible_key_id() {
+        byte[] okIdBytes = new byte[KeyId.MAX_KEY_ID_UTF8_LENGTH];
+        Arrays.fill(okIdBytes, (byte)'A');
+        var okId = KeyId.ofBytes(okIdBytes);
+        assertArrayEquals(okIdBytes, okId.asBytes());
+    }
+
+    @Test
+    void too_big_key_id_throws() {
+        byte[] tooBigIdBytes = new byte[KeyId.MAX_KEY_ID_UTF8_LENGTH + 1];
+        Arrays.fill(tooBigIdBytes, (byte)'A');
+        assertThrows(IllegalArgumentException.class, () -> KeyId.ofBytes(tooBigIdBytes));
+    }
+
+    @Test
+    void malformed_utf8_key_id_is_rejected_on_construction() {
+        byte[] malformedIdBytes = new byte[]{ (byte)0xC0 }; // First part of a 2-byte continuation without trailing byte
+        assertThrows(IllegalArgumentException.class, () -> KeyId.ofBytes(malformedIdBytes));
+    }
+
+}
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 f635121eda0..74b4ca0854b 100644
--- a/security-utils/src/test/java/com/yahoo/security/SharedKeyTest.java
+++ b/security-utils/src/test/java/com/yahoo/security/SharedKeyTest.java
@@ -21,7 +21,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class SharedKeyTest {
 
-    private static final byte[] KEY_ID_1 = new byte[] {1};
+    private static final KeyId KEY_ID_1 = KeyId.ofString("1");
 
     @Test
     void generated_secret_key_is_128_bit_aes() {
@@ -49,7 +49,7 @@ public class SharedKeyTest {
     void token_v1_representation_is_stable() {
         var receiverPrivate = KeyUtils.fromBase64EncodedX25519PrivateKey("4qGcntygFn_a3uqeBa1PbDlygQ-cpOuNznTPIz9ftWE");
         var receiverPublic  = KeyUtils.fromBase64EncodedX25519PublicKey( "ROAH_S862tNMpbJ49lu1dPXFCPHFIXZK30pSrMZEmEg");
-        byte[] keyId        = toUtf8Bytes("my key ID");
+        var keyId           = KeyId.ofString("my key ID");
 
         // Token generated for the above receiver public key, with the below expected shared secret (in hex)
         var publicToken = "AQlteSBrZXkgSUQgAtTxJJdmv3eUoW5Z3NJSdZ3poKPEkW0SJOGQXP6CaC5XfyAVoUlK_NyYIMsJKyNYKU6WmagZpVG2zQGFJoqiFA";
@@ -58,7 +58,7 @@ public class SharedKeyTest {
         var theirSealed = SealedSharedKey.fromTokenString(publicToken);
         var theirShared = SharedKeyGenerator.fromSealedKey(theirSealed, receiverPrivate);
 
-        assertArrayEquals(keyId, theirSealed.keyId());
+        assertEquals(keyId, theirSealed.keyId());
         assertEquals(expectedSharedSecret, hex(theirShared.secretKey().getEncoded()));
     }
 
@@ -73,21 +73,21 @@ public class SharedKeyTest {
 
     @Test
     void token_carries_opaque_key_id_bytes_as_metadata() {
-        byte[] keyId    = toUtf8Bytes("hello key id world");
+        var keyId       = KeyId.ofString("hello key id world");
         var keyPair     = KeyUtils.generateX25519KeyPair();
         var myShared    = SharedKeyGenerator.generateForReceiverPublicKey(keyPair.getPublic(), keyId);
         var publicToken = myShared.sealedSharedKey().toTokenString();
         var theirShared = SealedSharedKey.fromTokenString(publicToken);
-        assertArrayEquals(theirShared.keyId(), keyId);
+        assertEquals(theirShared.keyId(), keyId);
     }
 
     @Test
     void key_id_integrity_is_protected_by_aad() {
-        byte[] goodId = toUtf8Bytes("my key 1");
-        var keyPair   = KeyUtils.generateX25519KeyPair();
-        var myShared  = SharedKeyGenerator.generateForReceiverPublicKey(keyPair.getPublic(), goodId);
-        var mySealed  = myShared.sealedSharedKey();
-        byte[] badId  = toUtf8Bytes("my key 2");
+        var goodId   = KeyId.ofString("my key 1");
+        var keyPair  = KeyUtils.generateX25519KeyPair();
+        var myShared = SharedKeyGenerator.generateForReceiverPublicKey(keyPair.getPublic(), goodId);
+        var mySealed = myShared.sealedSharedKey();
+        var badId    = KeyId.ofString("my key 2");
 
         var tamperedShared = new SealedSharedKey(badId, mySealed.enc(), mySealed.ciphertext());
         // Should not be able to unseal the token since the AAD auth tag won't be correct
@@ -96,44 +96,32 @@ public class SharedKeyTest {
     }
 
     @Test
-    void key_id_encoding_size_is_bounded() {
-        byte[] okId = new byte[SealedSharedKey.MAX_KEY_ID_UTF8_LENGTH];
-        Arrays.fill(okId, (byte)'A');
+    void can_encode_and_decode_largest_possible_key_id() {
+        byte[] okIdBytes = new byte[KeyId.MAX_KEY_ID_UTF8_LENGTH];
+        Arrays.fill(okIdBytes, (byte)'A');
+        var okId     = KeyId.ofBytes(okIdBytes);
         var keyPair  = KeyUtils.generateX25519KeyPair();
         var myShared = SharedKeyGenerator.generateForReceiverPublicKey(keyPair.getPublic(), okId);
-        assertArrayEquals(okId, myShared.sealedSharedKey().keyId());
+        assertEquals(okId, myShared.sealedSharedKey().keyId());
 
         var asToken = myShared.sealedSharedKey().toTokenString();
         var decoded = SealedSharedKey.fromTokenString(asToken);
-        assertArrayEquals(okId, decoded.keyId());
-
-        byte[] tooBigId = new byte[SealedSharedKey.MAX_KEY_ID_UTF8_LENGTH + 1];
-        Arrays.fill(tooBigId, (byte)'A');
-        assertThrows(IllegalArgumentException.class,
-                     () -> SharedKeyGenerator.generateForReceiverPublicKey(keyPair.getPublic(), tooBigId));
-    }
-
-    @Test
-    void malformed_utf8_key_id_is_rejected_on_construction() {
-        byte[] malformedId = new byte[]{ (byte)0xC0 }; // First part of a 2-byte continuation without trailing byte
-        var keyPair = KeyUtils.generateX25519KeyPair();
-        assertThrows(IllegalArgumentException.class,
-                     () -> SharedKeyGenerator.generateForReceiverPublicKey(keyPair.getPublic(), malformedId));
+        assertEquals(okId, decoded.keyId());
     }
 
     // TODO make this test less implementation specific if possible...
     @Test
     void malformed_utf8_key_id_is_rejected_on_parsing() {
-        byte[] goodId = new byte[] { (byte)'A' };
-        var keyPair   = KeyUtils.generateX25519KeyPair();
-        var myShared  = SharedKeyGenerator.generateForReceiverPublicKey(keyPair.getPublic(), goodId);
+        var goodId   = KeyId.ofBytes(new byte[] { (byte)'A' });
+        var keyPair  = KeyUtils.generateX25519KeyPair();
+        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 (see above test)
+        // Since the key ID is only 1 bytes long, patch it with a bad UTF-8 value
         byte[] tokenBytes = Base64.getUrlDecoder().decode(myShared.sealedSharedKey().toTokenString());
-        tokenBytes[2] = (byte)0xC0;
-        var tokenStr = Base64.getUrlEncoder().encodeToString(tokenBytes);
-        assertThrows(IllegalArgumentException.class, () -> SealedSharedKey.fromTokenString(tokenStr));
+        tokenBytes[2] = (byte)0xC0; // First part of a 2-byte continuation without trailing byte
+        var patchedTokenStr = Base64.getUrlEncoder().encodeToString(tokenBytes);
+        assertThrows(IllegalArgumentException.class, () -> SealedSharedKey.fromTokenString(patchedTokenStr));
     }
 
     static byte[] streamEncryptString(String data, SecretSharedKey secretSharedKey) throws IOException {