-- TrustPin attestation checks that Envoy's jwt_authn filter does not do
-- (examples/verify-token/README.md has the full rules). jwt_authn has
-- already verified the ES256 signature against the verifier's JWKS, iss,
-- aud and exp (60 s skew) and put the token's header and payload into
-- dynamic metadata; this filter checks:
--   * header typ is trustpin-attest+jwt (rejects key certificates),
--   * tp.prj is PROJECT,
--   * tp.env is not "debug",
--   * bind, when present, equals base64url(SHA-256(Authorization)); with
--     REQUIRE_BIND, a token without bind is refused.
-- Envoy's LuaJIT has no SHA-256, hence the small implementation below.

-- Expected tp.prj, tprn::project::<organizationId>::<projectId>: the
-- audience in envoy.yaml when that is the default tprn name. "" skips the
-- check (only for a custom audience that names a single project).
local PROJECT = "tprn::project::org_example::prj_example"
-- true refuses tokens without bind (attestation policy tokens.binding: required).
local REQUIRE_BIND = false

local bit = require("bit")
local band, bor, bxor, bnot = bit.band, bit.bor, bit.bxor, bit.bnot
local rshift, lshift, ror, tobit = bit.rshift, bit.lshift, bit.ror, bit.tobit

local K = {
  0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
  0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
  0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
  0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
  0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
  0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
  0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
  0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
}

local function sha256(msg)
  local h0, h1, h2, h3 = 0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a
  local h4, h5, h6, h7 = 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19
  local len = #msg
  msg = msg .. "\128" .. string.rep("\0", (55 - len) % 64)
  local bits = len * 8
  local hi = math.floor(bits / 0x100000000)
  local lo = bits % 0x100000000
  msg = msg .. string.char(
    rshift(hi, 24) % 256, rshift(hi, 16) % 256, rshift(hi, 8) % 256, hi % 256,
    math.floor(lo / 0x1000000) % 256, math.floor(lo / 0x10000) % 256, math.floor(lo / 0x100) % 256, lo % 256)
  local w = {}
  for chunk = 1, #msg, 64 do
    for i = 0, 15 do
      local a, b, c, d = msg:byte(chunk + i * 4, chunk + i * 4 + 3)
      w[i] = bor(lshift(a, 24), lshift(b, 16), lshift(c, 8), d)
    end
    for i = 16, 63 do
      local s0 = bxor(ror(w[i - 15], 7), ror(w[i - 15], 18), rshift(w[i - 15], 3))
      local s1 = bxor(ror(w[i - 2], 17), ror(w[i - 2], 19), rshift(w[i - 2], 10))
      w[i] = tobit(w[i - 16] + s0 + w[i - 7] + s1)
    end
    local a, b, c, d, e, f, g, h = h0, h1, h2, h3, h4, h5, h6, h7
    for i = 0, 63 do
      local S1 = bxor(ror(e, 6), ror(e, 11), ror(e, 25))
      local ch = bxor(band(e, f), band(bnot(e), g))
      local t1 = tobit(h + S1 + ch + K[i + 1] + w[i])
      local S0 = bxor(ror(a, 2), ror(a, 13), ror(a, 22))
      local maj = bxor(band(a, b), band(a, c), band(b, c))
      local t2 = tobit(S0 + maj)
      h, g, f, e, d, c, b, a = g, f, e, tobit(d + t1), c, b, a, tobit(t1 + t2)
    end
    h0, h1, h2, h3 = tobit(h0 + a), tobit(h1 + b), tobit(h2 + c), tobit(h3 + d)
    h4, h5, h6, h7 = tobit(h4 + e), tobit(h5 + f), tobit(h6 + g), tobit(h7 + h)
  end
  local out = {}
  for _, v in ipairs({ h0, h1, h2, h3, h4, h5, h6, h7 }) do
    out[#out + 1] = string.char(band(rshift(v, 24), 255), band(rshift(v, 16), 255), band(rshift(v, 8), 255), band(v, 255))
  end
  return table.concat(out)
end

local B64 = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"

local function base64url(s)
  local out = {}
  for i = 1, #s, 3 do
    local a, b, c = s:byte(i, i + 2)
    local n = lshift(a, 16) + lshift(b or 0, 8) + (c or 0)
    out[#out + 1] = B64:sub(rshift(n, 18) % 64 + 1, rshift(n, 18) % 64 + 1)
    out[#out + 1] = B64:sub(rshift(n, 12) % 64 + 1, rshift(n, 12) % 64 + 1)
    if b then out[#out + 1] = B64:sub(rshift(n, 6) % 64 + 1, rshift(n, 6) % 64 + 1) end
    if c then out[#out + 1] = B64:sub(n % 64 + 1, n % 64 + 1) end
  end
  return table.concat(out)
end

local function refuse(handle, why)
  handle:logWarn("trustpin attestation rejected: " .. why)
  handle:respond({ [":status"] = "401", ["content-type"] = "text/plain" }, "attestation rejected\n")
end

function envoy_on_request(handle)
  local md = handle:streamInfo():dynamicMetadata():get("envoy.filters.http.jwt_authn")
  if md == nil or md["payload"] == nil or md["header"] == nil then
    return refuse(handle, "no verified token")
  end
  if md["header"]["typ"] ~= "trustpin-attest+jwt" then
    return refuse(handle, "wrong typ")
  end
  local payload = md["payload"]
  if payload["exp"] == nil then
    return refuse(handle, "no exp")
  end
  if PROJECT ~= "" and (payload["tp"] == nil or payload["tp"]["prj"] ~= PROJECT) then
    return refuse(handle, "wrong project")
  end
  if payload["tp"] ~= nil and payload["tp"]["env"] == "debug" then
    return refuse(handle, "debug token")
  end
  if payload["bind"] == nil and REQUIRE_BIND then
    return refuse(handle, "unbound token")
  end
  if payload["bind"] ~= nil then
    local auth = handle:headers():get("authorization") or ""
    if payload["bind"] ~= base64url(sha256(auth)) then
      return refuse(handle, "bound to another session")
    end
  end
  handle:headers():replace("x-trustpin-sub", payload["sub"] or "")
end