|
| 1 | +#!/usr/bin/env bash |
| 2 | +# |
| 3 | +# test-mlkem.sh — round-trip the Java SDK's pure ML-KEM (FIPS 203) key |
| 4 | +# wrapping against a locally running OpenTDF platform. |
| 5 | +# |
| 6 | +# Per algorithm: encrypt → assert manifest → KAS rewrap → decrypt → diff. |
| 7 | +# |
| 8 | +# Differs from test-hybrid-pqc.sh in three places: |
| 9 | +# * Wire format: pure ML-KEM is raw concat (mlkem_ct || AES-GCM blob), not |
| 10 | +# ASN.1 SEQUENCE. The manifest check below validates the wrappedKey |
| 11 | +# length matches ciphertextSize + 28 bytes (IV(12) + tag(16) + DEK(?)). |
| 12 | +# * keyAccess[0].type == "wrapped" (NOT "hybrid-wrapped"). Pure ML-KEM |
| 13 | +# reuses the RSA slot; the KAS disambiguates by registered key algorithm. |
| 14 | +# * SPKI OIDs are the NIST FIPS 203 ones (2.16.840.1.101.3.4.4.{2,3}). |
| 15 | +# |
| 16 | +# Prereqs: |
| 17 | +# * Local platform up at $PLATFORM_ENDPOINT with ML-KEM KAS keys registered |
| 18 | +# for mlkem:768 and (optionally) mlkem:1024 |
| 19 | +# * java, mvn (JDK 17), unzip, jq on PATH |
| 20 | +# * grpcurl optional (used only for the pre-flight key-publication check) |
| 21 | +# |
| 22 | +# Usage: |
| 23 | +# scripts/test-mlkem.sh # full run, both variants |
| 24 | +# scripts/test-mlkem.sh --skip-build # reuse existing jar |
| 25 | +# scripts/test-mlkem.sh --skip-kas-check # skip grpcurl pre-flight |
| 26 | +# scripts/test-mlkem.sh --algorithms MLKEM768Key # subset |
| 27 | +# PLATFORM_ENDPOINT=http://localhost:8080 scripts/test-mlkem.sh |
| 28 | +# |
| 29 | +# See scripts/README.md for a full prereq + troubleshooting guide. |
| 30 | + |
| 31 | +set -euo pipefail |
| 32 | + |
| 33 | +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" |
| 34 | +REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)" |
| 35 | +JAR="$REPO_ROOT/cmdline/target/cmdline.jar" |
| 36 | + |
| 37 | +PLATFORM_ENDPOINT="${PLATFORM_ENDPOINT:-http://localhost:8080}" |
| 38 | +KAS_URL="${KAS_URL:-$PLATFORM_ENDPOINT}" |
| 39 | +CLIENT_ID="${CLIENT_ID:-opentdf-sdk}" |
| 40 | +CLIENT_SECRET="${CLIENT_SECRET:-secret}" |
| 41 | +DATA_ATTR="${DATA_ATTR:-https://example.com/attr/attr1/value/value1}" |
| 42 | +ALGORITHMS=(MLKEM768Key MLKEM1024Key) |
| 43 | +SKIP_BUILD=0 |
| 44 | +SKIP_KAS_CHECK=0 |
| 45 | + |
| 46 | +# With `set -u`, a bare `$2` for a value-taking flag with no argument would |
| 47 | +# crash with "unbound variable" instead of the documented exit 2 misuse path. |
| 48 | +require_opt_value() { |
| 49 | + local opt="$1" |
| 50 | + local val="${2-}" |
| 51 | + if [[ -z "$val" || "$val" == --* ]]; then |
| 52 | + echo "missing value for $opt" >&2 |
| 53 | + exit 2 |
| 54 | + fi |
| 55 | +} |
| 56 | + |
| 57 | +while [[ $# -gt 0 ]]; do |
| 58 | + case "$1" in |
| 59 | + --skip-build) SKIP_BUILD=1; shift ;; |
| 60 | + --skip-kas-check) SKIP_KAS_CHECK=1; shift ;; |
| 61 | + --algorithms) require_opt_value "$1" "${2-}"; IFS=, read -r -a ALGORITHMS <<< "$2"; shift 2 ;; |
| 62 | + --platform-endpoint) require_opt_value "$1" "${2-}"; PLATFORM_ENDPOINT="$2"; shift 2 ;; |
| 63 | + --kas-url) require_opt_value "$1" "${2-}"; KAS_URL="$2"; shift 2 ;; |
| 64 | + --attr) require_opt_value "$1" "${2-}"; DATA_ATTR="$2"; shift 2 ;; |
| 65 | + --client-id) require_opt_value "$1" "${2-}"; CLIENT_ID="$2"; shift 2 ;; |
| 66 | + --client-secret) require_opt_value "$1" "${2-}"; CLIENT_SECRET="$2"; shift 2 ;; |
| 67 | + -h|--help) sed -n '2,/^$/p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;; |
| 68 | + *) echo "unknown option: $1" >&2; exit 2 ;; |
| 69 | + esac |
| 70 | +done |
| 71 | + |
| 72 | +# Map KeyType enum name → the mlkem:* algorithm string the KAS expects. |
| 73 | +# Function form (instead of `declare -A`) so this works on macOS bash 3.2. |
| 74 | +alg_to_string() { |
| 75 | + case "$1" in |
| 76 | + MLKEM768Key) echo "mlkem:768" ;; |
| 77 | + MLKEM1024Key) echo "mlkem:1024" ;; |
| 78 | + *) return 1 ;; |
| 79 | + esac |
| 80 | +} |
| 81 | + |
| 82 | +# Map KeyType enum name → expected SPKI OID inside the standard PUBLIC KEY PEM |
| 83 | +# (NIST FIPS 203). The pre-flight extracts via openssl asn1parse and compares. |
| 84 | +alg_to_oid() { |
| 85 | + case "$1" in |
| 86 | + MLKEM768Key) echo "2.16.840.1.101.3.4.4.2" ;; |
| 87 | + MLKEM1024Key) echo "2.16.840.1.101.3.4.4.3" ;; |
| 88 | + *) return 1 ;; |
| 89 | + esac |
| 90 | +} |
| 91 | + |
| 92 | +# Map KeyType enum name → ML-KEM ciphertext size (FIPS 203). |
| 93 | +alg_to_ct_size() { |
| 94 | + case "$1" in |
| 95 | + MLKEM768Key) echo 1088 ;; |
| 96 | + MLKEM1024Key) echo 1568 ;; |
| 97 | + *) return 1 ;; |
| 98 | + esac |
| 99 | +} |
| 100 | + |
| 101 | +WORK_DIR="$(mktemp -d -t mlkem-XXXXXX)" |
| 102 | +trap 'rm -rf "$WORK_DIR"' EXIT |
| 103 | + |
| 104 | +if [[ -t 1 ]]; then |
| 105 | + GREEN=$'\033[0;32m'; RED=$'\033[0;31m'; YELLOW=$'\033[0;33m'; RESET=$'\033[0m' |
| 106 | +else |
| 107 | + GREEN=''; RED=''; YELLOW=''; RESET='' |
| 108 | +fi |
| 109 | +pass() { echo "${GREEN}[OK]${RESET} $*"; } |
| 110 | +fail() { echo "${RED}[FAIL]${RESET} $*"; } |
| 111 | +info() { echo "${YELLOW}[..]${RESET} $*"; } |
| 112 | + |
| 113 | +require() { command -v "$1" >/dev/null 2>&1 || { fail "missing required tool: $1"; exit 2; }; } |
| 114 | +require java; require unzip; require jq |
| 115 | +[[ $SKIP_BUILD -eq 1 ]] || require mvn |
| 116 | + |
| 117 | +# Portable base64 decode: GNU/BusyBox accept `-d`, BSD/macOS prior to 12 use `-D`. |
| 118 | +if printf 'MA==\n' | base64 -d >/dev/null 2>&1; then |
| 119 | + BASE64_DECODE_FLAG="-d" |
| 120 | +elif printf 'MA==\n' | base64 -D >/dev/null 2>&1; then |
| 121 | + BASE64_DECODE_FLAG="-D" |
| 122 | +else |
| 123 | + fail "neither 'base64 -d' nor 'base64 -D' works on this system"; exit 2 |
| 124 | +fi |
| 125 | +b64decode() { base64 "$BASE64_DECODE_FLAG"; } |
| 126 | + |
| 127 | +run_cmdline() { |
| 128 | + java -jar "$JAR" \ |
| 129 | + --client-id="$CLIENT_ID" \ |
| 130 | + --client-secret="$CLIENT_SECRET" \ |
| 131 | + --platform-endpoint="$PLATFORM_ENDPOINT" \ |
| 132 | + -h "$@" |
| 133 | +} |
| 134 | + |
| 135 | +##### 1. Build |
| 136 | +if [[ $SKIP_BUILD -eq 0 ]]; then |
| 137 | + info "Building cmdline (mvn clean install -DskipTests)" |
| 138 | + build_log="$WORK_DIR/build.log" |
| 139 | + if ! (cd "$REPO_ROOT" && mvn --batch-mode clean install -DskipTests) > "$build_log" 2>&1; then |
| 140 | + fail "Maven build failed. Tail of build log:" |
| 141 | + tail -40 "$build_log" | sed 's/^/ /' |
| 142 | + if grep -q "Buf API token" "$build_log" 2>/dev/null; then |
| 143 | + fail "Hint: run 'buf registry login' or export BUF_INPUT_HTTPS_USERNAME / BUF_INPUT_HTTPS_PASSWORD before retrying." |
| 144 | + fi |
| 145 | + exit 1 |
| 146 | + fi |
| 147 | + pass "Build complete" |
| 148 | +else |
| 149 | + info "Skipping build (--skip-build)" |
| 150 | +fi |
| 151 | +[[ -f "$JAR" ]] || { fail "jar not found at $JAR — run without --skip-build"; exit 1; } |
| 152 | + |
| 153 | +##### 2. Pre-flight: confirm KAS publishes ML-KEM keys |
| 154 | +if [[ $SKIP_KAS_CHECK -eq 0 ]] && command -v grpcurl >/dev/null 2>&1; then |
| 155 | + info "Pre-flight: querying KAS for ML-KEM public keys" |
| 156 | + host="${PLATFORM_ENDPOINT#http://}"; host="${host#https://}" |
| 157 | + for alg_name in "${ALGORITHMS[@]}"; do |
| 158 | + if ! alg=$(alg_to_string "$alg_name"); then |
| 159 | + fail "unknown algorithm: $alg_name"; exit 2 |
| 160 | + fi |
| 161 | + resp=$(grpcurl -plaintext -d "{\"algorithm\":\"$alg\"}" \ |
| 162 | + "$host" kas.AccessService/PublicKey 2>&1 || true) |
| 163 | + pem=$(jq -r '.publicKey // empty' <<<"$resp" 2>/dev/null || true) |
| 164 | + if [[ -z "$pem" ]]; then |
| 165 | + fail "$alg: KAS returned no publicKey. Response was:" |
| 166 | + echo "$resp" | head -5 | sed 's/^/ /' |
| 167 | + fail "Is the platform running with the ML-KEM-capable KAS branch and the key registered?" |
| 168 | + exit 1 |
| 169 | + fi |
| 170 | + first_line=$(echo "$pem" | head -1) |
| 171 | + if [[ "$first_line" != *"BEGIN PUBLIC KEY"* ]]; then |
| 172 | + fail "$alg: KAS returned a non-SPKI PEM (first line: $first_line)" |
| 173 | + exit 1 |
| 174 | + fi |
| 175 | + expected_oid=$(alg_to_oid "$alg_name") |
| 176 | + if command -v openssl >/dev/null 2>&1; then |
| 177 | + actual_oid=$(printf '%s\n' "$pem" | openssl asn1parse 2>/dev/null \ |
| 178 | + | awk '/OBJECT/ {sub(/^:/, "", $NF); print $NF; exit}') |
| 179 | + if [[ -z "$actual_oid" ]]; then |
| 180 | + fail "$alg: could not extract SPKI OID via openssl asn1parse" |
| 181 | + exit 1 |
| 182 | + fi |
| 183 | + if [[ "$actual_oid" != "$expected_oid" ]]; then |
| 184 | + fail "$alg: SPKI OID mismatch — expected $expected_oid, got $actual_oid" |
| 185 | + exit 1 |
| 186 | + fi |
| 187 | + pass "$alg: KAS returns SPKI PEM with OID $actual_oid" |
| 188 | + else |
| 189 | + pass "$alg: KAS returns SPKI PEM (openssl not available; OID not verified)" |
| 190 | + fi |
| 191 | + done |
| 192 | +else |
| 193 | + info "Skipping KAS pre-flight check" |
| 194 | +fi |
| 195 | + |
| 196 | +##### 3. Round-trip each algorithm |
| 197 | +PAYLOAD="$WORK_DIR/payload" |
| 198 | +printf 'pure ml-kem round-trip payload @ %s\n' "$(date)" > "$PAYLOAD" |
| 199 | +PAYLOAD_BYTES=$(wc -c < "$PAYLOAD" | tr -d ' ') |
| 200 | +info "Test payload: $PAYLOAD_BYTES bytes" |
| 201 | +echo " --- plaintext ---" |
| 202 | +sed 's/^/ /' < "$PAYLOAD" |
| 203 | +echo " --- end plaintext ---" |
| 204 | + |
| 205 | +failures=() |
| 206 | +for alg_name in "${ALGORITHMS[@]}"; do |
| 207 | + tdf="$WORK_DIR/test-${alg_name}.tdf" |
| 208 | + out="$WORK_DIR/out-${alg_name}" |
| 209 | + enc_log="$WORK_DIR/encrypt-${alg_name}.log" |
| 210 | + dec_log="$WORK_DIR/decrypt-${alg_name}.log" |
| 211 | + |
| 212 | + info "[$alg_name] encrypt" |
| 213 | + if ! run_cmdline encrypt \ |
| 214 | + --kas-url="$KAS_URL" \ |
| 215 | + --mime-type=text/plain \ |
| 216 | + --attr="$DATA_ATTR" \ |
| 217 | + --autoconfigure=false \ |
| 218 | + --encap-key-type="$alg_name" \ |
| 219 | + -f "$PAYLOAD" > "$tdf" 2> "$enc_log"; then |
| 220 | + fail "$alg_name: encrypt failed" |
| 221 | + sed 's/^/ /' < "$enc_log" |
| 222 | + failures+=("$alg_name (encrypt)") |
| 223 | + continue |
| 224 | + fi |
| 225 | + |
| 226 | + info "[$alg_name] verify manifest" |
| 227 | + manifest_entry=$(unzip -l "$tdf" 2>/dev/null | awk '/manifest\.json$/ {print $NF; exit}') |
| 228 | + if [[ -z "$manifest_entry" ]]; then |
| 229 | + fail "$alg_name: no manifest.json entry inside $tdf" |
| 230 | + failures+=("$alg_name (manifest entry missing)") |
| 231 | + continue |
| 232 | + fi |
| 233 | + manifest=$(unzip -p "$tdf" "$manifest_entry") |
| 234 | + # In Manifest.java, the Java field `keyType` is annotated with |
| 235 | + # @SerializedName("type"), so the JSON key is "type" (not "keyType"). |
| 236 | + keyType=$(jq -r '.encryptionInformation.keyAccess[0].type' <<<"$manifest") |
| 237 | + ephem=$(jq -r '.encryptionInformation.keyAccess[0].ephemeralPublicKey // ""' <<<"$manifest") |
| 238 | + wrapped=$(jq -r '.encryptionInformation.keyAccess[0].wrappedKey // ""' <<<"$manifest") |
| 239 | + if [[ "$keyType" != "wrapped" ]]; then |
| 240 | + fail "$alg_name: type='$keyType' (expected 'wrapped')" |
| 241 | + echo " keyAccess[0]:" |
| 242 | + jq '.encryptionInformation.keyAccess[0]' <<<"$manifest" 2>/dev/null | sed 's/^/ /' |
| 243 | + failures+=("$alg_name (bad type: $keyType)") |
| 244 | + continue |
| 245 | + fi |
| 246 | + if [[ -n "$ephem" ]]; then |
| 247 | + fail "$alg_name: ephemeralPublicKey unexpectedly set ('$ephem')" |
| 248 | + failures+=("$alg_name (stray ephemeralPublicKey)") |
| 249 | + continue |
| 250 | + fi |
| 251 | + if [[ -z "$wrapped" ]]; then |
| 252 | + fail "$alg_name: wrappedKey is empty" |
| 253 | + failures+=("$alg_name (empty wrappedKey)") |
| 254 | + continue |
| 255 | + fi |
| 256 | + # Pure ML-KEM wire format: ML-KEM ciphertext (fixed size per variant) |
| 257 | + # || AES-GCM(IV(12) || DEK(32) || tag(16)) = ciphertextSize + 60 bytes. |
| 258 | + expected_ct_size=$(alg_to_ct_size "$alg_name") |
| 259 | + expected_len=$((expected_ct_size + 12 + 32 + 16)) |
| 260 | + actual_len=$(b64decode <<<"$wrapped" 2>/dev/null | wc -c | tr -d ' ') |
| 261 | + if [[ "$actual_len" != "$expected_len" ]]; then |
| 262 | + fail "$alg_name: wrappedKey length $actual_len bytes != expected $expected_len (ct=$expected_ct_size + 60)" |
| 263 | + failures+=("$alg_name (bad wrappedKey length)") |
| 264 | + continue |
| 265 | + fi |
| 266 | + pass "$alg_name: manifest OK (wrapped, $actual_len-byte envelope, no ephemeralPublicKey)" |
| 267 | + echo " --- keyAccess[0] (KAO) ---" |
| 268 | + jq '.encryptionInformation.keyAccess[0]' <<<"$manifest" | sed 's/^/ /' |
| 269 | + echo " --- end keyAccess[0] ---" |
| 270 | + |
| 271 | + info "[$alg_name] decrypt (rewrap via KAS)" |
| 272 | + if ! run_cmdline decrypt -f "$tdf" > "$out" 2> "$dec_log"; then |
| 273 | + fail "$alg_name: decrypt failed" |
| 274 | + sed 's/^/ /' < "$dec_log" |
| 275 | + failures+=("$alg_name (decrypt)") |
| 276 | + continue |
| 277 | + fi |
| 278 | + if ! diff -q "$PAYLOAD" "$out" >/dev/null; then |
| 279 | + fail "$alg_name: decrypted payload differs from original" |
| 280 | + echo " --- expected (first 200 bytes) ---" |
| 281 | + head -c 200 "$PAYLOAD" | sed 's/^/ /' |
| 282 | + echo |
| 283 | + echo " --- got (first 200 bytes) ---" |
| 284 | + head -c 200 "$out" | sed 's/^/ /' |
| 285 | + echo |
| 286 | + failures+=("$alg_name (payload mismatch)") |
| 287 | + continue |
| 288 | + fi |
| 289 | + pass "$alg_name: round-trip OK" |
| 290 | + out_bytes=$(wc -c < "$out" | tr -d ' ') |
| 291 | + echo " --- decrypted ($out_bytes bytes) ---" |
| 292 | + sed 's/^/ /' < "$out" |
| 293 | + echo " --- end decrypted ---" |
| 294 | +done |
| 295 | + |
| 296 | +echo |
| 297 | +if [[ ${#failures[@]} -eq 0 ]]; then |
| 298 | + echo "${GREEN}All ${#ALGORITHMS[@]} ML-KEM algorithm(s) passed round-trip.${RESET}" |
| 299 | + exit 0 |
| 300 | +else |
| 301 | + echo "${RED}FAILURES (${#failures[@]}):${RESET}" |
| 302 | + printf ' - %s\n' "${failures[@]}" |
| 303 | + exit 1 |
| 304 | +fi |
0 commit comments