From 963b65b3cc19147a5a270b8efe627332ffa2b824 Mon Sep 17 00:00:00 2001
From: Valentin Vidic <vvidic@google.com>
Date: Thu, 7 Mar 2024 22:16:42 +0100
Subject: [PATCH] Fix handling of prefix type in login CLI

Accept any index without a check or do a MSB check with the public key.
---
 cli/commands.c | 16 ++++++++++++----
 cli/test.sh    | 12 ++++++++----
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/cli/commands.c b/cli/commands.c
index eabc9d3..120a6fc 100644
--- a/cli/commands.c
+++ b/cli/commands.c
@@ -359,10 +359,18 @@ int login(int argc, char **argv) {
     fprintf(stderr, "handshake size is invalid in path %s\n", path);
     goto out;
   }
-  if ((handshake[0] & 0x80) != 0) {
-    fprintf(stderr,
-            "only \"service-key-indicator\" prefix type is supported\n");
-    goto out;
+  if ((handshake[0] & 0x80) == 0) {
+    uint8_t public_key[GLOME_MAX_PUBLIC_KEY_LENGTH] = {0};
+    if (glome_derive_key(private_key, public_key)) {
+      fprintf(stderr, "unable to generate a public key\n");
+      goto out;
+    }
+    // Most significant bit is not set for X25519 key (see RFC 7748).
+    uint8_t public_key_msb = public_key[GLOME_MAX_PUBLIC_KEY_LENGTH - 1];
+    if (handshake[0] != public_key_msb) {
+      fprintf(stderr, "unexpected public key prefix\n");
+      goto out;
+    }
   }
   uint8_t peer_key[GLOME_MAX_PRIVATE_KEY_LENGTH] = {0};
   memcpy(peer_key, handshake + 1, GLOME_MAX_PUBLIC_KEY_LENGTH);
diff --git a/cli/test.sh b/cli/test.sh
index a0152f9..2e4299c 100755
--- a/cli/test.sh
+++ b/cli/test.sh
@@ -69,14 +69,18 @@ for n in 1 2; do
 done
 
 key="$t/vector-2/a"
-path="v2/R4cvQ1u4uJ0OOtYqouURB07hleHDnvaogAFBi-ZW48N2/myhost/exec=%2Fbin%2Fsh/"
 expected_tag="ZmxczN4x3g4goXu-A2AuuEEVftgS6xM-6gYj-dRrlis="
-tag=$("$binary" login --key "$key" "$path")
-if [ "$tag" != "$expected_tag" ]; then
+for path in \
+  "v2/R4cvQ1u4uJ0OOtYqouURB07hleHDnvaogAFBi-ZW48N2/myhost/exec=%2Fbin%2Fsh/" \
+  "v2/x4cvQ1u4uJ0OOtYqouURB07hleHDnvaogAFBi-ZW48N2/myhost/exec=%2Fbin%2Fsh/"
+do
+  tag=$("$binary" login --key "$key" "$path")
+  if [ "$tag" != "$expected_tag" ]; then
     echo "Generated wrong tag for test path $path" >&2
     echo "$expected_tag <- expected" >&2
     echo "$tag <- actual" >&2
     errors=$((errors + 1))
-fi
+  fi
+done
 
 exit "$errors"