Introduce PMBTOKENS key abstractions.

Start tidying up the TRUST_TOKENS/PMBTOKENS split.

Change-Id: Iabcbc864f4016dfcb22438387446b04d31b64beb
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/41064
Reviewed-by: Steven Valdez <svaldez@google.com>
Commit-Queue: Steven Valdez <svaldez@google.com>
diff --git a/crypto/trust_token/internal.h b/crypto/trust_token/internal.h
index 582dc49..9b1d2a9 100644
--- a/crypto/trust_token/internal.h
+++ b/crypto/trust_token/internal.h
@@ -39,18 +39,13 @@
 // protocol.
 #define PMBTOKEN_NONCE_SIZE 64
 
-// Structure representing a single Trust Token public key with the specified ID.
-struct trust_token_client_key_st {
-  uint32_t id;
+typedef struct {
   EC_RAW_POINT pub0;
   EC_RAW_POINT pub1;
   EC_RAW_POINT pubs;
-};
+} PMBTOKEN_CLIENT_KEY;
 
-// Structure representing a single Trust Token private key with the specified
-// ID.
-struct trust_token_issuer_key_st {
-  uint32_t id;
+typedef struct {
   EC_SCALAR x0;
   EC_SCALAR y0;
   EC_SCALAR x1;
@@ -60,7 +55,7 @@
   EC_RAW_POINT pub0;
   EC_RAW_POINT pub1;
   EC_RAW_POINT pubs;
-};
+} PMBTOKEN_ISSUER_KEY;
 
 // PMBTOKEN_PRETOKEN represents the intermediate state a client keeps during a
 // PMBToken issuance operation.
@@ -87,44 +82,68 @@
 // PMBTOKEN_TOKEN_free releases the memory associated with |token|.
 void PMBTOKEN_TOKEN_free(PMBTOKEN_TOKEN *token);
 
-// pmbtoken_compute_public computes the public keypairs from the private
-// keypairs in |key|. It returns one on success and zero on failure.
-int pmbtoken_compute_public(struct trust_token_issuer_key_st *key);
+// pmbtoken_generate_key generates a fresh keypair and writes their serialized
+// forms into |out_private| and |out_public|. It returns one on success and zero
+// on failure.
+int pmbtoken_generate_key(CBB *out_private, CBB *out_public);
+
+// pmbtoken_client_key_from_bytes decodes a client key from |in| and sets |key|
+// to the resulting key. It returns one on success and zero
+// on failure.
+int pmbtoken_client_key_from_bytes(PMBTOKEN_CLIENT_KEY *key, const uint8_t *in,
+                                   size_t len);
+
+// pmbtoken_issuer_key_from_bytes decodes a issuer key from |in| and sets |key|
+// to the resulting key. It returns one on success and zero
+// on failure.
+int pmbtoken_issuer_key_from_bytes(PMBTOKEN_ISSUER_KEY *key, const uint8_t *in,
+                                   size_t len);
 
 // pmbtoken_blind generates a new blinded pretoken based on the configuration of
 // |ctx| as per the first stage of the AT.Usr operation and returns the
 // resulting pretoken.
 PMBTOKEN_PRETOKEN *pmbtoken_blind(void);
 
-// pmbtoken_sign signs a blinded point based on the configuration of |ctx|
-// and the key specified by |key_id| with a private metadata value of
-// |private_metadata| as per the AT.Sig operation and stores the resulting nonce
-// and points in |*out_s|, |*out_Wp|, and |*out_Wsp| and the resulting DLEQ
-// proof in |*out_proof|. The caller takes ownership of |*out_proof| and is
+// pmbtoken_sign signs a blinded point with |key| and a private metadata value
+// of |private_metadata| as per the AT.Sig operation and stores the resulting
+// nonce and points in |*out_s|, |*out_Wp|, and |*out_Wsp| and the resulting
+// DLEQ proof in |*out_proof|. The caller takes ownership of |*out_proof| and is
 // responsible for freeing it using |OPENSSL_free|. It returns one on success
 // and zero on failure.
-int pmbtoken_sign(const TRUST_TOKEN_ISSUER *ctx,
+int pmbtoken_sign(const PMBTOKEN_ISSUER_KEY *key,
                   uint8_t out_s[PMBTOKEN_NONCE_SIZE], EC_RAW_POINT *out_Wp,
                   EC_RAW_POINT *out_Wsp, uint8_t **out_proof,
                   size_t *out_proof_len, const EC_RAW_POINT *Tp,
-                  uint32_t key_id, uint8_t private_metadata);
+                  uint8_t private_metadata);
 
 // pmbtoken_unblind unblinds the result of an AT.Sig operation as per the final
 // stage of the AT.Usr operation and sets |*out_token| to the resulting token.
 // It returns one on success and zero on failure.
-int pmbtoken_unblind(PMBTOKEN_TOKEN *out_token,
-                     const struct trust_token_client_key_st *key,
+int pmbtoken_unblind(const PMBTOKEN_CLIENT_KEY *key, PMBTOKEN_TOKEN *out_token,
                      const uint8_t s[PMBTOKEN_NONCE_SIZE],
                      const EC_RAW_POINT *Wp, const EC_RAW_POINT *Wsp,
                      const uint8_t *proof, size_t proof_len,
                      const PMBTOKEN_PRETOKEN *pretoken);
 
-// pmbtoken_read verifies the validity of a PMBToken |token| using the key
-// specified by |key_id| and stores the value of the private metadata
-// bit in |*out_private_metadata|. It returns one if the token is valid and zero
-// otherwise.
-int pmbtoken_read(const TRUST_TOKEN_ISSUER *ctx, uint8_t *out_private_metadata,
-                  const PMBTOKEN_TOKEN *token, uint32_t key_id);
+// pmbtoken_read verifies a PMBToken |token| using |key| and stores the value of
+// the private metadata bit in |*out_private_metadata|. It returns one if the
+// token is valid and zero otherwise.
+int pmbtoken_read(const PMBTOKEN_ISSUER_KEY *key, uint8_t *out_private_metadata,
+                  const PMBTOKEN_TOKEN *token);
+
+
+// Structure representing a single Trust Token public key with the specified ID.
+struct trust_token_client_key_st {
+  uint32_t id;
+  PMBTOKEN_CLIENT_KEY key;
+};
+
+// Structure representing a single Trust Token private key with the specified
+// ID.
+struct trust_token_issuer_key_st {
+  uint32_t id;
+  PMBTOKEN_ISSUER_KEY key;
+};
 
 struct trust_token_client_st {
   // max_batchsize is the maximum supported batchsize.
diff --git a/crypto/trust_token/pmbtoken.c b/crypto/trust_token/pmbtoken.c
index 192eba9..5cac5dd 100644
--- a/crypto/trust_token/pmbtoken.c
+++ b/crypto/trust_token/pmbtoken.c
@@ -12,6 +12,8 @@
  * OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
  * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. */
 
+#include <openssl/trust_token.h>
+
 #include <openssl/bn.h>
 #include <openssl/bytestring.h>
 #include <openssl/ec.h>
@@ -20,7 +22,6 @@
 #include <openssl/nid.h>
 #include <openssl/rand.h>
 #include <openssl/sha.h>
-#include <openssl/trust_token.h>
 
 #include "../ec_extra/internal.h"
 #include "../fipsmodule/bn/internal.h"
@@ -199,10 +200,23 @@
                            len) == len;
 }
 
-int TRUST_TOKEN_generate_key(uint8_t *out_priv_key, size_t *out_priv_key_len,
-                             size_t max_priv_key_len, uint8_t *out_pub_key,
-                             size_t *out_pub_key_len, size_t max_pub_key_len,
-                             uint32_t id) {
+static int cbs_get_prefixed_point(CBS *cbs, const EC_GROUP *group,
+                                  EC_RAW_POINT *out) {
+  CBS child;
+  return CBS_get_u16_length_prefixed(cbs, &child) &&
+         ec_point_from_uncompressed(group, out, CBS_data(&child),
+                                    CBS_len(&child));
+}
+
+void PMBTOKEN_PRETOKEN_free(PMBTOKEN_PRETOKEN *pretoken) {
+  OPENSSL_free(pretoken);
+}
+
+void PMBTOKEN_TOKEN_free(PMBTOKEN_TOKEN *token) {
+  OPENSSL_free(token);
+}
+
+int pmbtoken_generate_key(CBB *out_private, CBB *out_public) {
   EC_GROUP *group = EC_GROUP_new_by_curve_name(NID_secp521r1);
   if (group == NULL) {
     return 0;
@@ -217,66 +231,78 @@
     return 0;
   }
 
-  int ret = 0;
-  CBB cbb;
-  CBB_zero(&cbb);
-  size_t scalar_len = BN_num_bytes(&group->order);
-  if (!CBB_init_fixed(&cbb, out_priv_key, max_priv_key_len) ||
-      !CBB_add_u32(&cbb, id)) {
-    OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_BUFFER_TOO_SMALL);
-    goto err;
-  }
-
   const EC_SCALAR *scalars[] = {&x0, &y0, &x1, &y1, &xs, &ys};
+  size_t scalar_len = BN_num_bytes(&group->order);
   for (size_t i = 0; i < OPENSSL_ARRAY_SIZE(scalars); i++) {
     uint8_t *buf;
-    if (!CBB_add_space(&cbb, &buf, scalar_len)) {
+    if (!CBB_add_space(out_private, &buf, scalar_len)) {
       OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_BUFFER_TOO_SMALL);
-      goto err;
+      return 0;
     }
     ec_scalar_to_bytes(group, buf, &scalar_len, scalars[i]);
   }
 
-  if (!CBB_finish(&cbb, NULL, out_priv_key_len)) {
+  // TODO(https://crbug.com/boringssl/331): When updating the key format, remove
+  // the redundant length prefixes.
+  CBB child;
+  if (!CBB_add_u16_length_prefixed(out_public, &child) ||
+      !point_to_cbb(&child, group, &pub0) ||
+      !CBB_add_u16_length_prefixed(out_public, &child) ||
+      !point_to_cbb(&child, group, &pub1) ||
+      !CBB_add_u16_length_prefixed(out_public, &child) ||
+      !point_to_cbb(&child, group, &pubs) ||
+      !CBB_flush(out_public)) {
     OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_BUFFER_TOO_SMALL);
-    goto err;
+    return 0;
   }
 
-  CBB pub_cbb;
-  if (!CBB_init_fixed(&cbb, out_pub_key, max_pub_key_len) ||
-      !CBB_add_u32(&cbb, id) ||
-      !CBB_add_u16_length_prefixed(&cbb, &pub_cbb) ||
-      !point_to_cbb(&pub_cbb, group, &pub0) ||
-      !CBB_add_u16_length_prefixed(&cbb, &pub_cbb) ||
-      !point_to_cbb(&pub_cbb, group, &pub1) ||
-      !CBB_add_u16_length_prefixed(&cbb, &pub_cbb) ||
-      !point_to_cbb(&pub_cbb, group, &pubs) ||
-      !CBB_finish(&cbb, NULL, out_pub_key_len)) {
-    OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_BUFFER_TOO_SMALL);
-    goto err;
-  }
-
-  ret = 1;
-
-err:
-  CBB_cleanup(&cbb);
-  return ret;
+  return 1;
 }
 
-void PMBTOKEN_PRETOKEN_free(PMBTOKEN_PRETOKEN *pretoken) {
-  OPENSSL_free(pretoken);
-}
-
-void PMBTOKEN_TOKEN_free(PMBTOKEN_TOKEN *token) {
-  OPENSSL_free(token);
-}
-
-int pmbtoken_compute_public(struct trust_token_issuer_key_st *key) {
+int pmbtoken_client_key_from_bytes(PMBTOKEN_CLIENT_KEY *key, const uint8_t *in,
+                                   size_t len) {
   EC_GROUP *group = EC_GROUP_new_by_curve_name(NID_secp521r1);
   if (group == NULL) {
     return 0;
   }
 
+  // TODO(https://crbug.com/boringssl/331): When updating the key format, remove
+  // the redundant length prefixes.
+  CBS cbs;
+  CBS_init(&cbs, in, len);
+  if (!cbs_get_prefixed_point(&cbs, group, &key->pub0) ||
+      !cbs_get_prefixed_point(&cbs, group, &key->pub1) ||
+      !cbs_get_prefixed_point(&cbs, group, &key->pubs) ||
+      CBS_len(&cbs) != 0) {
+    OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_DECODE_FAILURE);
+    return 0;
+  }
+
+  return 1;
+}
+
+int pmbtoken_issuer_key_from_bytes(PMBTOKEN_ISSUER_KEY *key, const uint8_t *in,
+                                   size_t len) {
+  EC_GROUP *group = EC_GROUP_new_by_curve_name(NID_secp521r1);
+  if (group == NULL) {
+    return 0;
+  }
+
+  CBS cbs, tmp;
+  CBS_init(&cbs, in, len);
+  size_t scalar_len = BN_num_bytes(&group->order);
+  EC_SCALAR *scalars[] = {&key->x0, &key->y0, &key->x1,
+                          &key->y1, &key->xs, &key->ys};
+  for (size_t i = 0; i < OPENSSL_ARRAY_SIZE(scalars); i++) {
+    if (!CBS_get_bytes(&cbs, &tmp, scalar_len) ||
+        !ec_scalar_from_bytes(group, scalars[i], CBS_data(&tmp),
+                              CBS_len(&tmp))) {
+      OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_DECODE_FAILURE);
+     return 0;
+    }
+  }
+
+  // Recompute the public key.
   EC_RAW_POINT h;
   if (!get_h(&h) ||
       !mul_twice_base(group, &key->pubs, &key->xs, &h, &key->ys) ||
@@ -471,8 +497,7 @@
 // DLEQOR2 with only one value (n=1).
 
 static int dleq_generate(const EC_GROUP *group, uint8_t **out_proof,
-                         size_t *out_proof_len,
-                         const struct trust_token_issuer_key_st *priv,
+                         size_t *out_proof_len, const PMBTOKEN_ISSUER_KEY *priv,
                          const EC_RAW_POINT *T, const EC_RAW_POINT *S,
                          const EC_RAW_POINT *W, const EC_RAW_POINT *Ws,
                          uint8_t private_metadata) {
@@ -634,8 +659,7 @@
 }
 
 static int dleq_verify(const EC_GROUP *group, const uint8_t *proof,
-                       size_t proof_len,
-                       const struct trust_token_client_key_st *pub,
+                       size_t proof_len, const PMBTOKEN_CLIENT_KEY *pub,
                        const EC_RAW_POINT *T, const EC_RAW_POINT *S,
                        const EC_RAW_POINT *W, const EC_RAW_POINT *Ws) {
   EC_RAW_POINT h;
@@ -721,33 +745,16 @@
   return 1;
 }
 
-int pmbtoken_sign(const TRUST_TOKEN_ISSUER *ctx,
+int pmbtoken_sign(const PMBTOKEN_ISSUER_KEY *key,
                   uint8_t out_s[PMBTOKEN_NONCE_SIZE], EC_RAW_POINT *out_Wp,
                   EC_RAW_POINT *out_Wsp, uint8_t **out_proof,
                   size_t *out_proof_len, const EC_RAW_POINT *Tp,
-                  uint32_t key_id, uint8_t private_metadata) {
+                  uint8_t private_metadata) {
   EC_GROUP *group = EC_GROUP_new_by_curve_name(NID_secp521r1);
   if (group == NULL) {
     return 0;
   }
 
-  if (ctx->num_keys == 0) {
-    OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_NO_KEYS_CONFIGURED);
-    return 0;
-  }
-  const struct trust_token_issuer_key_st *key = NULL;
-  for (size_t i = 0; i < ctx->num_keys; i++) {
-    if (ctx->keys[i].id == key_id) {
-      key = &ctx->keys[i];
-      break;
-    }
-  }
-
-  if (key == NULL) {
-    OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_INVALID_KEY_ID);
-    return 0;
-  }
-
   EC_SCALAR xb, yb;
   BN_ULONG mask = ((BN_ULONG)0) - (private_metadata & 1);
   ec_scalar_select(group, &xb, mask, &key->x1, &key->x0);
@@ -769,8 +776,7 @@
                        out_Wsp, private_metadata);
 }
 
-int pmbtoken_unblind(PMBTOKEN_TOKEN *out_token,
-                     const struct trust_token_client_key_st *key,
+int pmbtoken_unblind(const PMBTOKEN_CLIENT_KEY *key, PMBTOKEN_TOKEN *out_token,
                      const uint8_t s[PMBTOKEN_NONCE_SIZE],
                      const EC_RAW_POINT *Wp, const EC_RAW_POINT *Wsp,
                      const uint8_t *proof, size_t proof_len,
@@ -799,30 +805,13 @@
   return 1;
 }
 
-int pmbtoken_read(const TRUST_TOKEN_ISSUER *ctx, uint8_t *out_private_metadata,
-                  const PMBTOKEN_TOKEN *token, uint32_t key_id) {
+int pmbtoken_read(const PMBTOKEN_ISSUER_KEY *key, uint8_t *out_private_metadata,
+                  const PMBTOKEN_TOKEN *token) {
   EC_GROUP *group = EC_GROUP_new_by_curve_name(NID_secp521r1);
   if (group == NULL) {
     return 0;
   }
 
-  if (ctx->num_keys == 0) {
-    OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_NO_KEYS_CONFIGURED);
-    return 0;
-  }
-  const struct trust_token_issuer_key_st *key = NULL;
-  for (size_t i = 0; i < ctx->num_keys; i++) {
-    if (ctx->keys[i].id == key_id) {
-      key = &ctx->keys[i];
-      break;
-    }
-  }
-
-  if (key == NULL) {
-    OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_INVALID_KEY_ID);
-    return 0;
-  }
-
   EC_RAW_POINT T;
   if (!hash_t(group, &T, token->t)) {
     return 0;
diff --git a/crypto/trust_token/trust_token.c b/crypto/trust_token/trust_token.c
index 6f5756a..e460803 100644
--- a/crypto/trust_token/trust_token.c
+++ b/crypto/trust_token/trust_token.c
@@ -73,6 +73,41 @@
   OPENSSL_free(token);
 }
 
+int TRUST_TOKEN_generate_key(uint8_t *out_priv_key, size_t *out_priv_key_len,
+                             size_t max_priv_key_len, uint8_t *out_pub_key,
+                             size_t *out_pub_key_len, size_t max_pub_key_len,
+                             uint32_t id) {
+  // Prepend the key ID in front of the PMBTokens format.
+  int ret = 0;
+  CBB priv_cbb, pub_cbb;
+  CBB_zero(&priv_cbb);
+  CBB_zero(&pub_cbb);
+  if (!CBB_init_fixed(&priv_cbb, out_priv_key, max_priv_key_len) ||
+      !CBB_init_fixed(&pub_cbb, out_pub_key, max_pub_key_len) ||
+      !CBB_add_u32(&priv_cbb, id) ||
+      !CBB_add_u32(&pub_cbb, id)) {
+    OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_BUFFER_TOO_SMALL);
+    goto err;
+  }
+
+  if (!pmbtoken_generate_key(&priv_cbb, &pub_cbb)) {
+    goto err;
+  }
+
+  if (!CBB_finish(&priv_cbb, NULL, out_priv_key_len) ||
+      !CBB_finish(&pub_cbb, NULL, out_pub_key_len)) {
+    OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_BUFFER_TOO_SMALL);
+    goto err;
+  }
+
+  ret = 1;
+
+err:
+  CBB_cleanup(&priv_cbb);
+  CBB_cleanup(&pub_cbb);
+  return ret;
+}
+
 TRUST_TOKEN_CLIENT *TRUST_TOKEN_CLIENT_new(size_t max_batchsize) {
   if (max_batchsize > 0xffff) {
     // The protocol supports only two-byte token counts.
@@ -107,26 +142,18 @@
 
 int TRUST_TOKEN_CLIENT_add_key(TRUST_TOKEN_CLIENT *ctx, size_t *out_key_index,
                                const uint8_t *key, size_t key_len) {
-  EC_GROUP *group = EC_GROUP_new_by_curve_name(NID_secp521r1);
-  if (group == NULL) {
-    return 0;
-  }
-
   if (ctx->num_keys == OPENSSL_ARRAY_SIZE(ctx->keys)) {
     OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_TOO_MANY_KEYS);
     return 0;
   }
 
   struct trust_token_client_key_st *key_s = &ctx->keys[ctx->num_keys];
-
   CBS cbs;
   CBS_init(&cbs, key, key_len);
   uint32_t key_id;
   if (!CBS_get_u32(&cbs, &key_id) ||
-      !cbs_get_raw_point(&cbs, group, &key_s->pub0) ||
-      !cbs_get_raw_point(&cbs, group, &key_s->pub1) ||
-      !cbs_get_raw_point(&cbs, group, &key_s->pubs) ||
-      CBS_len(&cbs) != 0) {
+      !pmbtoken_client_key_from_bytes(&key_s->key, CBS_data(&cbs),
+                                      CBS_len(&cbs))) {
     OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_DECODE_FAILURE);
     return 0;
   }
@@ -245,7 +272,7 @@
 
     PMBTOKEN_PRETOKEN *pretoken = sk_PMBTOKEN_PRETOKEN_value(ctx->pretokens, i);
     PMBTOKEN_TOKEN pmbtoken;
-    if (!pmbtoken_unblind(&pmbtoken, key, s, &Wp, &Wsp, CBS_data(&proof),
+    if (!pmbtoken_unblind(&key->key, &pmbtoken, s, &Wp, &Wsp, CBS_data(&proof),
                           CBS_len(&proof), pretoken)) {
       goto err;
     }
@@ -397,41 +424,22 @@
 
 int TRUST_TOKEN_ISSUER_add_key(TRUST_TOKEN_ISSUER *ctx, const uint8_t *key,
                                size_t key_len) {
-  EC_GROUP *group = EC_GROUP_new_by_curve_name(NID_secp521r1);
-  if (group == NULL) {
-    return 0;
-  }
-
   if (ctx->num_keys == OPENSSL_ARRAY_SIZE(ctx->keys)) {
     OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_TOO_MANY_KEYS);
     return 0;
   }
 
-  CBS cbs, tmp;
+  struct trust_token_issuer_key_st *key_s = &ctx->keys[ctx->num_keys];
+  CBS cbs;
   CBS_init(&cbs, key, key_len);
   uint32_t key_id;
-  if (!CBS_get_u32(&cbs, &key_id)) {
+  if (!CBS_get_u32(&cbs, &key_id) ||
+      !pmbtoken_issuer_key_from_bytes(&key_s->key, CBS_data(&cbs),
+                                      CBS_len(&cbs))) {
     OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_DECODE_FAILURE);
     return 0;
   }
 
-  size_t scalar_len = BN_num_bytes(&group->order);
-  struct trust_token_issuer_key_st *key_s = &(ctx->keys[ctx->num_keys]);
-  EC_SCALAR *scalars[] = {&key_s->x0, &key_s->y0, &key_s->x1,
-                          &key_s->y1, &key_s->xs, &key_s->ys};
-  for (size_t i = 0; i < OPENSSL_ARRAY_SIZE(scalars); i++) {
-    if (!CBS_get_bytes(&cbs, &tmp, scalar_len) ||
-        !ec_scalar_from_bytes(group, scalars[i], CBS_data(&tmp),
-                              CBS_len(&tmp))) {
-      OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_DECODE_FAILURE);
-     return 0;
-    }
-  }
-
-  if (!pmbtoken_compute_public(key_s)) {
-    return 0;
-  }
-
   key_s->id = key_id;
   ctx->num_keys += 1;
   return 1;
@@ -460,6 +468,16 @@
   return 1;
 }
 
+static const struct trust_token_issuer_key_st *trust_token_issuer_get_key(
+    const TRUST_TOKEN_ISSUER *ctx, uint32_t key_id) {
+  for (size_t i = 0; i < ctx->num_keys; i++) {
+    if (ctx->keys[i].id == key_id) {
+      return &ctx->keys[i];
+    }
+  }
+  return NULL;
+}
+
 int TRUST_TOKEN_ISSUER_issue(const TRUST_TOKEN_ISSUER *ctx, uint8_t **out,
                              size_t *out_len, size_t *out_tokens_issued,
                              const uint8_t *request, size_t request_len,
@@ -474,15 +492,9 @@
     max_issuance = ctx->max_batchsize;
   }
 
-  int found_public_metadata = 0;
-  for (size_t i = 0; i < ctx->num_keys; i++) {
-    if (ctx->keys[i].id == public_metadata) {
-      found_public_metadata = 1;
-      break;
-    }
-  }
-
-  if (!found_public_metadata || private_metadata > 1) {
+  const struct trust_token_issuer_key_st *key =
+      trust_token_issuer_get_key(ctx, public_metadata);
+  if (key == NULL || private_metadata > 1) {
     OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_INVALID_METADATA);
     return 0;
   }
@@ -524,8 +536,8 @@
     EC_RAW_POINT Wp, Wsp;
     uint8_t *proof = NULL;
     size_t proof_len;
-    if (!pmbtoken_sign(ctx, s, &Wp, &Wsp, &proof, &proof_len, &Tp,
-                       public_metadata, private_metadata)) {
+    if (!pmbtoken_sign(&key->key, s, &Wp, &Wsp, &proof, &proof_len, &Tp,
+                       private_metadata)) {
       goto err;
     }
 
@@ -628,13 +640,20 @@
 
   // Parse the token. If there is an error, treat it as an invalid token.
   PMBTOKEN_TOKEN pmbtoken;
-  if (!CBS_get_u32(&token_cbs, &public_metadata) ||
+  if (!CBS_get_u32(&token_cbs, &public_metadata)) {
+    OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_INVALID_TOKEN);
+    return 0;
+  }
+
+  const struct trust_token_issuer_key_st *key =
+      trust_token_issuer_get_key(ctx, public_metadata);
+  if (key == NULL ||
       !CBS_copy_bytes(&token_cbs, pmbtoken.t, PMBTOKEN_NONCE_SIZE) ||
       !cbs_get_raw_point(&token_cbs, group, &pmbtoken.S) ||
       !cbs_get_raw_point(&token_cbs, group, &pmbtoken.W) ||
       !cbs_get_raw_point(&token_cbs, group, &pmbtoken.Ws) ||
       CBS_len(&token_cbs) != 0 ||
-      !pmbtoken_read(ctx, &private_metadata, &pmbtoken, public_metadata)) {
+      !pmbtoken_read(&key->key, &private_metadata, &pmbtoken)) {
     OPENSSL_PUT_ERROR(TRUST_TOKEN, TRUST_TOKEN_R_INVALID_TOKEN);
     return 0;
   }
diff --git a/crypto/trust_token/trust_token_test.cc b/crypto/trust_token/trust_token_test.cc
index 2da8b08..e312dc8 100644
--- a/crypto/trust_token/trust_token_test.cc
+++ b/crypto/trust_token/trust_token_test.cc
@@ -595,8 +595,8 @@
   bssl::UniquePtr<uint8_t> free_issue_msg(issue_msg);
 
   struct trust_token_issuer_key_st *key = &issuer->keys[0];
-  EC_SCALAR *scalars[] = {&key->x0, &key->y0, &key->x1,
-                          &key->y1, &key->xs, &key->ys};
+  EC_SCALAR *scalars[] = {&key->key.x0, &key->key.y0, &key->key.x1,
+                          &key->key.y1, &key->key.xs, &key->key.ys};
   int corrupted_key = std::get<1>(GetParam());
 
   // Corrupt private key scalar.