# 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.jwtIsPresent STRING;
set var.jwtIsPresent = false;

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 618 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 618 "jwt:algorithm-not-supported";
    } else if (var.jwtKeyData == "" || var.jwtKeyID == "") {
      error 618 "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);
      }

      # Convert from standard base64 with padding to URL-safe base64 with no padding, consistent with the JWT specification.
      set var.jwtCorrectSig = std.replace_suffix(std.replaceall(std.replaceall(var.jwtCorrectSig, "+", "-"), "/", "_"), "=", "");
      
      if (digest.secure_is_equal(var.jwtCorrectSig, var.jwtSig)) {
        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 618 "jwt:signature-fail";
    }
    log "JWT Signature verified";
  }

  if (var.jwtExpires == 0) {
    if (var.jwtOptionTimeInvalidBehavior == "anon") {
      set var.jwtSigVerified = false;
    } else {
      error 618 "jwt:expires-not-present-or-valid";
    }
  } else 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 618 "jwt:time-out-of-bounds";
    }
    log "Checked JWT time validity";
  }

  if (var.jwtPath ~ {"^(\*)?(\/[^\*]+)(/\*)?$"}) {
    if (

      # Exact match, eg "/index.html"
      (!re.group.1 && !re.group.3 && var.jwtPath != req.url.path) ||

      # Prefix match eg "/products/*"
      (!re.group.1 && re.group.3 && !std.prefixof(req.url.path, re.group.2)) ||

      # Suffix match eg "*/protected.html"
      (re.group.1 && !re.group.3 && !std.suffixof(req.url.path, re.group.2)) ||

      # Both eg "*/somedirectory/*"
      (re.group.1 && re.group.3 && !std.strstr(req.url.path, re.group.2))
    ) {
      if (var.jwtOptionPathInvalidBehavior == "anon") {
        set var.jwtSigVerified = false;
      } else {
        error 618 "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 618 "jwt:anonymous";
    }
  }
} else {
  if (var.jwtOptionAnonAccess == "allow") {
    set req.http.Auth-State = "anonymous";
  } else {
    error 618 "jwt:anonymous";
  }
}

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