# Declare some locally-scoped variables to help us with the
# processing of the authentication cookie
declare local var.jwtSource STRING;
declare local var.jwtHeader STRING;
declare local var.jwtHeaderDecoded STRING;
declare local var.jwtPayload STRING;
declare local var.jwtPayloadDecoded STRING;
declare local var.jwtSig STRING;
declare local var.jwtSigDecoded STRING;
declare local var.jwtStringToSign STRING;
declare local var.jwtCorrectSig STRING;
declare local var.jwtSigVerified BOOL;
declare local var.jwtKeyID STRING;
declare local var.jwtAlgo STRING;
declare local var.jwtKeyData STRING;
declare local var.jwtNotBefore INTEGER;
declare local var.jwtExpires INTEGER;
declare local var.jwtPath STRING;
declare local var.jwtTag STRING;

declare local var.jwtOptionTimeInvalidBehavior STRING;
set var.jwtOptionTimeInvalidBehavior = "block"; # Choose from 'anon' or 'block'
declare local var.jwtOptionPathInvalidBehavior STRING;
set var.jwtOptionPathInvalidBehavior = "anon"; # Choose from 'anon' or 'block'
declare local var.jwtTokenSource STRING;
set var.jwtTokenSource = "cookie"; # Choose from 'cookie' or 'query'
declare local var.jwtOptionAnonAccess STRING;
set var.jwtOptionAnonAccess = "allow"; # Choose from 'allow' or 'deny'

if (req.http.Fastly-JWT-Error) {
  error 718 req.http.Fastly-JWT-Error;
}

if (var.jwtTokenSource == "cookie") {
  set var.jwtSource = req.http.Cookie:auth;
} else { 
	set var.jwtSource = subfield(req.url.qs, "auth", "&");
}

if (var.jwtSource ~ "^([A-Za-z0-9-_=]+)\.([A-Za-z0-9-_=]+)\.([A-Za-z0-9-_.+/=]*)$") {
  set var.jwtHeader = re.group.1;
  set var.jwtHeaderDecoded = digest.base64url_decode(var.jwtHeader);
  set var.jwtPayload = re.group.2;
  set var.jwtPayloadDecoded = digest.base64url_decode(var.jwtPayload);
  set var.jwtSig = re.group.3;
  set var.jwtSigDecoded = digest.base64url_decode(var.jwtSig);
}

set var.jwtAlgo = if(var.jwtHeaderDecoded ~ {"\{(?:.+,)?\s*"alg" *: *"([^"]*)""}, re.group.1, "");
set var.jwtKeyID = if(var.jwtPayloadDecoded ~ {"\{(?:.+,)?\s*"key" *: *"([^"]*)""}, re.group.1, "");
set var.jwtKeyData = digest.base64_decode(table.lookup(solution_jwt_keys, var.jwtKeyID, ""));
set var.jwtStringToSign = var.jwtHeader "." var.jwtPayload;
set var.jwtNotBefore = std.atoi(if(var.jwtPayloadDecoded ~ {"\{(?:.+,)?\s*"nbf" *: *"?(\d+)[",\}]"}, re.group.1, "0"));
set var.jwtExpires = std.atoi(if(var.jwtPayloadDecoded ~ {"\{(?:.+,)?\s*"exp" *: *"?(\d+)[",\}]"}, re.group.1, "0"));
set var.jwtPath = if(var.jwtPayloadDecoded ~ {"\{(?:.+,)?\s*"path" *: *"([^\"]+)""}, re.group.1, "");
set var.jwtTag = if(var.jwtPayloadDecoded ~ {"\{(?:.+,)?\s*"tag" *: *"([^\"]+)""}, re.group.1, "");

if (var.jwtHeader) {
  if (var.jwtAlgo !~ "^(HS256|HS512|RS256|RS512)$") {
    error 718 "jwt:algorithm-not-supported";
  } else if (var.jwtKeyData == "" || var.jwtKeyID == "") {
    error 718 "jwt:key-not-found";
  } else if (std.prefixof(var.jwtAlgo, "HS")) {
    if (var.jwtAlgo == "HS256") {
      set var.jwtCorrectSig = digest.hmac_sha256_base64(var.jwtKeyData, var.jwtStringToSign);
    } else {
      set var.jwtCorrectSig = digest.hmac_sha512_base64(var.jwtKeyData, var.jwtStringToSign);
    }
    set var.jwtCorrectSig = digest.base64_decode(var.jwtCorrectSig);
    if (var.jwtCorrectSig == var.jwtSigDecoded) {
      set var.jwtSigVerified = true;
    } else {
      set var.jwtSigVerified = false;
    }
  } else if (var.jwtAlgo == "RS256") {
    set var.jwtSigVerified = digest.rsa_verify(sha256, var.jwtKeyData, var.jwtStringToSign, var.jwtSig, url);
  } else if (var.jwtAlgo == "RS512") {
    set var.jwtSigVerified = digest.rsa_verify(sha512, var.jwtKeyData, var.jwtStringToSign, var.jwtSig, url);
  }
  if (!var.jwtSigVerified) {
    error 718 "jwt:signature-fail";
  }
  log "JWT Signature verified";
}

if ((var.jwtNotBefore > 0 && !time.is_after(now, std.integer2time(var.jwtNotBefore))) || (var.jwtExpires > 0 && time.is_after(now, std.integer2time(var.jwtExpires))))  {
  if (var.jwtOptionTimeInvalidBehavior == "anon") {
    set var.jwtSigVerified = false;
  } else {
    error 718 "jwt:time-out-of-bounds";
  }
  log "Checked JWT time validity";
}

if (var.jwtPath ~ {"^([^\*]+)?(\*?)([^\*]+?)?$"} && var.jwtPath != "")  {
  if (
    (re.group.1 && !re.group.2 && !re.group.3 && var.jwtPath != re.group.1) ||
    (!re.group.1 && re.group.2 && !re.group.3 && !std.prefixof(var.jwtPath, re.group.1)) ||
    (!re.group.1 && re.group.2 && re.group.3 && !std.suffixof(var.jwtPath, re.group.3)) ||
    (re.group.1 && re.group.2 && re.group.3 && (!std.prefixof(var.jwtPath, re.group.3) || !std.suffixof(var.jwtPath, re.group.3)))
  ) {
    if (var.jwtOptionPathInvalidBehavior == "anon") {
      set var.jwtSigVerified = false;
    } else {
      error 718 "jwt:path-mismatch";
    }
  }
  log "Checked path constraint";
}

if (var.jwtTag != "") {
  set req.http.auth-require-tag = var.jwtTag;
}

if (var.jwtSigVerified) {
  set req.http.Auth-State = "authenticated";
  set req.http.Auth-UserID = if(var.jwtPayloadDecoded ~ {"\{(?:.+,)?\s*"uid" *: *"([^\"]+)""}, re.group.1, "");
  set req.http.Auth-Groups = if(var.jwtPayloadDecoded ~ {"\{(?:.+,)?\s*"groups" *: *"([^\"]+)""}, re.group.1, "");
  set req.http.Auth-Name = if(var.jwtPayloadDecoded ~ {"\{(?:.+,)?\s*"name" *: *"([^\"]+)""}, re.group.1, "");
  set req.http.Auth-Is-Admin = if (var.jwtPayloadDecoded ~ {"\{(?:.+,)?\s*"admin" *: *true"}, "1", "0");
} else {
  if (var.jwtOptionAnonAccess == "allow") {
    set req.http.Auth-State = "anonymous";
    unset req.http.Auth-UserID;
    unset req.http.Auth-Groups;
    unset req.http.Auth-Name;
    unset req.http.Auth-Is-Admin;
  } else {
    error 718 "jwt:anonymous";
  }
}

if (var.jwtTokenSource == "cookie") {
  unset req.http.Cookie:auth;
}