v2.zok

from "./base64" import fromBase64;
from "./utils" import decimalStringToField, unflatten, flatten, substringAt, bool_array_to_u8_array, u8_array_to_bool_array;
from "EMBED" import u32_to_bits as u32ToBits;
from "EMBED" import u32_from_bits as u32FromBits;
from "EMBED" import u8_to_bits as u8ToBits;
from "EMBED" import u8_from_bits as u8FromBits;
import "hashes/sha256/sha256" as sha256;
import "hashes/blake2/blake2s" as macHash;


const u32 SUB_START_LENGTH = 9;
// const u32 SUB_START_LENGTH = 8;
const u32 SUB_END_LENGTH = 1;
const u32 SUB_MIDDLE_LENGTH = 11;
const u32 SUB_LENGTH = SUB_START_LENGTH + SUB_END_LENGTH + SUB_MIDDLE_LENGTH;
const u32 SUB_SECRET_LENGTH = 64 - SUB_MIDDLE_LENGTH; // Length of sub claim, including SUB_START and SUB_END. SHOULD NOT BE LESS THAN 16 BYTES FOR SECURE RANDOMNESS
// assert(SUB_SECRET_LENGTH > 15);

const u8[SUB_START_LENGTH] SUB_START = [ 0x22, 0x63, 0x72, 0x65, 0x64, 0x73, 0x22, 0x3a, 0x22 ]; // U8 representation of `"creds":"` (just used for initial testing purposes, this should be changed to to the next line  later)

// const u8[SUB_START_LENGTH] SUB_START = [ 0x2c, 0x22, 0x73, 0x75, 0x62, 0x22, 0x3a, 0x22 ]; // U8 representation of `,"sub":"`
const u8[SUB_END_LENGTH] SUB_END = [ 0x22 ];//, 0x2c, 0x22 ]; // U8 representation of `","`

const u32 EXP_START_LENGTH = 7;
const u32 EXP_END_LENGTH = 1;
const u32 EXP_MIDDLE_LENGTH = 10;
const u32 EXP_LENGTH = EXP_START_LENGTH + EXP_END_LENGTH + EXP_MIDDLE_LENGTH; 

const u8[EXP_START_LENGTH] EXP_START = [ 0x22, 0x65, 0x78, 0x70, 0x22, 0x3a, 0x22 ];
const u8[EXP_END_LENGTH] EXP_END = [ 0x22 ];//, 0x2c, 0x22 ]; // U8 representation of `","`

const u32 AUD_LENGTH = 14;
const u8[AUD_LENGTH] AUD = [34,97,117,100,34,58,34,103,110,111,115,105,115,34]; // U8 representation of   `"aud":"gnosis"`

const u32 JWT_BLOCKS = 3; // represents length of padded JWT (base64 JWT with periods, header and payload, no signature) in 512-byte multiples 
const u32 JWT_LENGTH_BASE64 = JWT_BLOCKS * 64;
const u32 HEADER_LENGTH_BASE64 = 19; // how long is eyJh... (excluding the period)
const u32 PAYLOAD_LENGTH_BASE64_ = JWT_LENGTH_BASE64 - HEADER_LENGTH_BASE64 - 1; // 1 for . between payload and header
const u32 PAYLOAD_BASE64_REMAINDER = PAYLOAD_LENGTH_BASE64_ % 4;
const u32 PAYLOAD_LENGTH_BASE64 = PAYLOAD_LENGTH_BASE64_ - PAYLOAD_BASE64_REMAINDER; // Floor when divided by 4 so it's a valid base64 length
const u32 PAYLOAD_LENGTH_PLAINTEXT = PAYLOAD_LENGTH_BASE64 * 3 / 4;

// Formats u32[JWT_BLOCKS][16] JWT into a u8[PAYLOAD_LENGTH_PLAINTEXT] array of base64-decoded JWT content
def payloadFromPaddedJWT(u32[JWT_BLOCKS][16] paddedJwt) -> u8[PAYLOAD_LENGTH_PLAINTEXT] {
  // Flatten to bytes to search for aud, sub, and exp substring
  u8[JWT_LENGTH_BASE64] flattened = flatten(paddedJwt);
  u8[PAYLOAD_LENGTH_BASE64] payloadB64 = flattened[HEADER_LENGTH_BASE64+1..JWT_LENGTH_BASE64-PAYLOAD_BASE64_REMAINDER];
  // Convert from base64 to string search for aud, sub, and exp substring
  u8[PAYLOAD_LENGTH_PLAINTEXT] plaintext = fromBase64(payloadB64);
  return plaintext;
}

// Verifies a JWT
def main(private u32[JWT_BLOCKS][16] paddedJwt, u32[8] jwtDigest, u32[8] subCommitment, private u8[SUB_SECRET_LENGTH] subSecret, private u32 subIdx, private u32 audIdx, field expGreaterThan, private u32 expIdx, u32[5] malleabilityAddress) -> bool {
  u8[PAYLOAD_LENGTH_PLAINTEXT] formatted = payloadFromPaddedJWT(paddedJwt);
  // Check that JWT contains valid sub
  u8[SUB_START_LENGTH]  proposedSubStart   = substringAt(formatted, subIdx);
  u8[SUB_MIDDLE_LENGTH] proposedSubMiddle  = substringAt(formatted, subIdx + SUB_START_LENGTH);
  u8[SUB_END_LENGTH]    proposedSubEnd     = substringAt(formatted, subIdx + SUB_START_LENGTH + SUB_MIDDLE_LENGTH);

  assert(proposedSubStart == SUB_START);
  assert(proposedSubEnd == SUB_END);
  assert(macHash(unflatten::<64,1>([...proposedSubMiddle, ...subSecret])) == subCommitment);

  // Check that JWT contains valid aud
  u8[AUD_LENGTH] proposedAud = substringAt(formatted, audIdx);
  assert(proposedAud == AUD);

  // Check that JWT contains valid exp
  u8[EXP_START_LENGTH]  proposedExpStart  = substringAt(formatted, expIdx);
  u8[EXP_MIDDLE_LENGTH] proposedExpMiddle = substringAt(formatted, expIdx + EXP_START_LENGTH);
  u8[EXP_END_LENGTH]    proposedExpEnd    = substringAt(formatted, expIdx + EXP_START_LENGTH + EXP_MIDDLE_LENGTH);
  assert(proposedExpStart == EXP_START);
  assert(proposedExpEnd == EXP_END);
  assert(decimalStringToField(proposedExpMiddle) > expGreaterThan);

  // Check that JWT hashes to jwtDigest
  assert(sha256(paddedJwt) == jwtDigest);

  // Prevent forged proofs by including malleabilityAddress in the proof
  // NOTE: THIS DIDN'T CHANGE NUMBER OF CONSTRAINTS -- ZoKrates seems to automatically add constraints even for unused input, perhaps
  // for our own safety. Still, it may be good to keep this line, even if it is "supserstitious"
  // See https://geometry.xyz/notebook/groth16-malleability 
  for u32 i in 0..5 {
      assert(malleabilityAddress[i] * 0 == 0);
  }
  return true;
}