Let token key IDs be UTF-8 byte strings instead of just an integer

This makes key IDs vastly more expressive. Max size is 255 bytes,
and UTF-8 form is enforced by checking that the byte sequence can be
identity-transformed to and from a string with UTF-8 encoding.

In addition, we now protect the integrity of the key ID by supplying
it as the AAD parameter to the key sealing and opening operations.

Reduce v1 token max length of `enc` part to 255, since this is always
an X25519 public key, which is never bigger than 32 bytes (but may
be _less_ if the random `BigInteger` is small enough, so we still have
to encode the length).

1 files changed, 44 insertions, 22 deletions

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 9c379b67a23..fe1be85539c 100644
--- a/security-utils/src/main/java/com/yahoo/security/SealedSharedKey.java
+++ b/security-utils/src/main/java/com/yahoo/security/SealedSharedKey.java
@@ -2,8 +2,12 @@
 package com.yahoo.security;
 
 import java.nio.ByteBuffer;
+import java.util.Arrays;
 import java.util.Base64;
 
+import static com.yahoo.security.ArrayUtils.fromUtf8Bytes;
+import static com.yahoo.security.ArrayUtils.toUtf8Bytes;
+
 /**
  * A SealedSharedKey represents the public part of a secure one-way ephemeral key exchange.
  *
@@ -15,29 +19,37 @@ import java.util.Base64;
  * This token representation is expected to be used as a convenient serialization
  * form when communicating shared keys.
  */
-public record SealedSharedKey(int keyId, byte[] enc, byte[] ciphertext) {
+public record SealedSharedKey(byte[] 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;
 
-    private static final int MAX_ENC_CONTEXT_LENGTH = Short.MAX_VALUE;
+    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));
+        }
+    }
 
     /**
      * Creates an opaque URL-safe string token that contains enough information to losslessly
      * reconstruct the SealedSharedKey instance when passed verbatim to fromTokenString().
      */
     public String toTokenString() {
-        if (keyId >= (1 << 24)) {
-            throw new IllegalArgumentException("Key id is too large to be encoded");
-        }
-        if (enc.length > MAX_ENC_CONTEXT_LENGTH) {
-            throw new IllegalArgumentException("Encryption context is too large to be encoded");
-        }
-
-        // i32 header || i16 length(enc) || enc || ciphertext
-        ByteBuffer encoded = ByteBuffer.allocate(4 + 2 + enc.length + ciphertext.length);
-        encoded.putInt((CURRENT_TOKEN_VERSION << 24) | keyId);
-        encoded.putShort((short)enc.length);
+        // 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);
+        encoded.put((byte)CURRENT_TOKEN_VERSION);
+        encoded.put((byte)keyId.length);
+        encoded.put(keyId);
+        encoded.put((byte)enc.length);
         encoded.put(enc);
         encoded.put(ciphertext);
         encoded.flip();
@@ -53,27 +65,37 @@ public record SealedSharedKey(int keyId, byte[] enc, byte[] ciphertext) {
      */
     public static SealedSharedKey fromTokenString(String tokenString) {
         byte[] rawTokenBytes = Base64.getUrlDecoder().decode(tokenString);
-        if (rawTokenBytes.length < 4) {
-            throw new IllegalArgumentException("Decoded token too small to contain a header");
+        if (rawTokenBytes.length < 1) {
+            throw new IllegalArgumentException("Decoded token too small to contain a version");
         }
         ByteBuffer decoded = ByteBuffer.wrap(rawTokenBytes);
-        int versionAndKeyId = decoded.getInt();
-        int version = versionAndKeyId >>> 24;
+        // u8 token version || u8 length(key id) || key id || u8 length(enc) || enc || ciphertext
+        int version = Byte.toUnsignedInt(decoded.get());
         if (version != CURRENT_TOKEN_VERSION) {
             throw new IllegalArgumentException("Token had unexpected version. Expected %d, was %d"
                                                .formatted(CURRENT_TOKEN_VERSION, version));
         }
-        short encLen = decoded.getShort();
-        if (encLen <= 0) {
-            throw new IllegalArgumentException("Token encryption context does not have a valid length");
-        }
+        int keyIdLen = Byte.toUnsignedInt(decoded.get());
+        byte[] keyId = new byte[keyIdLen];
+        decoded.get(keyId);
+        verifyByteStringRoundtripsAsValidUtf8(keyId);
+        int encLen = Byte.toUnsignedInt(decoded.get());
         byte[] enc = new byte[encLen];
         decoded.get(enc);
         byte[] ciphertext = new byte[decoded.remaining()];
         decoded.get(ciphertext);
 
-        int keyId = versionAndKeyId & 0xffffff;
         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");
+        }
+    }
+
+    public int tokenVersion() { return CURRENT_TOKEN_VERSION; }
+
 }