-- TrustPin attestation check for Kong (open-source), run by the
-- pre-function plugin in the access phase. A request continues to the
-- upstream only when its TrustPin-Attestation header carries a valid access
-- token (examples/verify-token/README.md has the rules: ES256 by kid from the
-- verifier's JWKS, typ, iss, aud, exp with 60 s leeway, tp.prj, no debug
-- tokens, bind = base64url(SHA-256(Authorization))).
--
-- Uses only libraries Kong ships: lua-resty-openssl, lua-resty-http, cjson
-- and ngx.base64. pre-function needs `untrusted_lua = on` to require them.
local JWKS_URL = "http://issuer:18300/.well-known/jwks.json"  -- your verifier
local ISSUER = "https://attest.example.com"
local AUDIENCE = "tprn::project::org_example::prj_example"
-- Expected tp.prj, tprn::project::<organizationId>::<projectId>. "" means
-- AUDIENCE when it has that form (the default), otherwise no project check.
local PROJECT = ""
-- true refuses tokens without bind (attestation policy tokens.binding: required).
local REQUIRE_BIND = false
local LEEWAY = 60
local JWKS_MAX_AGE = 300  -- seconds a fetched JWKS is cached
local JWKS_COOLDOWN = 30  -- at most one refetch for an unknown kid per 30 s
local JWKS_TIMEOUT = 5000 -- ms before a JWKS fetch gives up

local cjson = require("cjson.safe")
local http = require("resty.http")
local pkey = require("resty.openssl.pkey")
local digest = require("resty.openssl.digest")
local b64 = require("ngx.base64")

local function refuse(why)
  kong.log.warn("trustpin attestation rejected: ", why)
  return kong.response.exit(401, { message = "attestation rejected" })
end

local function fetch_jwks()
  local httpc = http.new()
  httpc:set_timeouts(JWKS_TIMEOUT, JWKS_TIMEOUT, JWKS_TIMEOUT)
  local res, err = httpc:request_uri(JWKS_URL, { method = "GET" })
  if not res or res.status ~= 200 then
    return nil, "jwks fetch failed: " .. tostring(err or res.status)
  end
  local doc = cjson.decode(res.body)
  local keys = {}
  for _, jwk in ipairs((doc and doc.keys) or {}) do
    if jwk.kty == "EC" and jwk.crv == "P-256" and jwk.kid then
      keys[jwk.kid] = cjson.encode({ kty = "EC", crv = "P-256", x = jwk.x, y = jwk.y })
    end
  end
  return keys
end

-- The JWKS lives in Kong's cache (shared by the workers; concurrent misses
-- wait for one fetch) for JWKS_MAX_AGE; when a refetch fails, the previous
-- copy is kept for another JWKS_COOLDOWN. An unknown kid (key rotation)
-- refetches under a new cache key, at most once per JWKS_COOLDOWN across all
-- workers (the atomic add on the `kong` shared dict); a kid still unknown is
-- refused without fetching. (pre-function cannot invalidate a cache entry in
-- DB-less mode, hence the generation number in the key.)
local function key_for(kid)
  local dict = ngx.shared.kong
  local opts = { ttl = JWKS_MAX_AGE, resurrect_ttl = JWKS_COOLDOWN }
  local gen = dict:get("trustpin_attest_jwks_gen") or 0
  local keys, err = kong.cache:get("trustpin_attest_jwks:" .. gen, opts, fetch_jwks)
  if keys and not keys[kid] and dict:add("trustpin_attest_jwks_cooldown", true, JWKS_COOLDOWN) then
    gen = dict:incr("trustpin_attest_jwks_gen", 1, 0)
    keys, err = kong.cache:get("trustpin_attest_jwks:" .. gen, opts, fetch_jwks)
  end
  if not keys then
    return nil, err
  end
  if not keys[kid] then
    return nil, "unknown kid"
  end
  return pkey.new(keys[kid], { format = "JWK" })
end

-- ES256 signatures are raw r||s (64 bytes); OpenSSL wants DER.
local function der_signature(raw)
  local function int(b)
    b = b:gsub("^%z+", "")
    if #b == 0 then b = "\0" end
    if b:byte(1) >= 0x80 then b = "\0" .. b end
    return "\2" .. string.char(#b) .. b
  end
  local body = int(raw:sub(1, 32)) .. int(raw:sub(33, 64))
  return "\48" .. string.char(#body) .. body
end

local function check()
  local token = kong.request.get_header("TrustPin-Attestation")
  if not token then
    return "no TrustPin-Attestation header"
  end
  local h64, p64, s64 = token:match("^([%w_-]+)%.([%w_-]+)%.([%w_-]*)$")
  if not h64 then
    return "malformed token"
  end
  local header = cjson.decode(b64.decode_base64url(h64) or "")
  if type(header) ~= "table" or header.alg ~= "ES256" then
    return "alg must be ES256"
  end
  if header.typ ~= "trustpin-attest+jwt" then
    return "wrong typ"
  end
  local sig = b64.decode_base64url(s64)
  if not sig or #sig ~= 64 then
    return "bad signature encoding"
  end
  local key, err = key_for(header.kid)
  if not key then
    return err
  end
  local ok = key:verify(der_signature(sig), h64 .. "." .. p64, "sha256")
  if not ok then
    return "bad signature"
  end
  local c = cjson.decode(b64.decode_base64url(p64) or "")
  if type(c) ~= "table" then
    return "malformed payload"
  end
  if c.iss ~= ISSUER then
    return "wrong iss"
  end
  local aud_ok = c.aud == AUDIENCE
  if type(c.aud) == "table" then
    for _, a in ipairs(c.aud) do aud_ok = aud_ok or a == AUDIENCE end
  end
  if not aud_ok then
    return "wrong aud"
  end
  if type(c.exp) ~= "number" or c.exp + LEEWAY < ngx.time() then
    return "expired"
  end
  local project = PROJECT
  if project == "" and AUDIENCE:sub(1, 15) == "tprn::project::" then
    project = AUDIENCE
  end
  if project ~= "" and (type(c.tp) ~= "table" or c.tp.prj ~= project) then
    return "wrong project"
  end
  if type(c.tp) == "table" and c.tp.env == "debug" then
    return "debug token"
  end
  if c.bind == nil and REQUIRE_BIND then
    return "unbound token"
  end
  if c.bind ~= nil then
    local auth = kong.request.get_header("Authorization") or ""
    local sum = digest.new("sha256"):final(auth)
    if c.bind ~= b64.encode_base64url(sum) then
      return "bound to another session"
    end
  end
  kong.service.request.set_header("X-TrustPin-Sub", c.sub or "")
  return nil
end

local why = check()
if why then
  return refuse(why)
end