Cloud KMS Signers

import { KMSClient, SignCommand } from '@aws-sdk/client-kms'

const kms = new KMSClient({})

const KEY_ID = process.env.KMS_KEY_ID
if (!KEY_ID) throw new Error('Missing env var KMS_KEY_ID')

// secp256k1 curve parameters
const SECP256K1_P = BigInt('0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F')
const SECP256K1_N = BigInt('0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141')
const SECP256K1_HALF_N = SECP256K1_N / 2n
const SECP256K1_Gx = BigInt('0x79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798')
const SECP256K1_Gy = BigInt('0x483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8')

export const handler = async (event) => {
  const body = typeof event.body === 'string' ? JSON.parse(event.body) : (event.body ?? event)

  const digestInput = body.userOpHash ?? body.useropHash ?? body.hash ?? body.digest
  if (!digestInput) {
    return resp(400, {
      error: 'Missing userOpHash (0x..32 bytes) or digest (0x..32 bytes or base64 32 bytes)'
    })
  }

  const signerAddress = body.signerAddress
  if (!signerAddress || !signerAddress.startsWith('0x') || signerAddress.length !== 42) {
    return resp(400, {
      error: 'Missing or invalid signerAddress (must be 0x-prefixed 20-byte hex address)'
    })
  }

  const digestBytes = parse32ByteDigest(digestInput)
  if (!digestBytes) return resp(400, { error: 'Digest must be 32 bytes (hex 0x.. or base64)' })

  // Apply EIP-191 personal sign prefix before signing
  // This computes: keccak256("\x19Ethereum Signed Message:\n32" + digestBytes)
  const ethSignedHash = hashMessage(digestBytes)

  const out = await kms.send(
    new SignCommand({
      KeyId: KEY_ID,
      Message: ethSignedHash,
      MessageType: 'DIGEST',
      SigningAlgorithm: 'ECDSA_SHA_256'
    })
  )

  if (!out.Signature) return resp(500, { error: 'KMS did not return a signature' })

  let { r, s } = derToRS(Buffer.from(out.Signature))

  // Enforce low-s (EIP-2)
  let sBig = BigInt(s)
  let sFlipped = false
  if (sBig > SECP256K1_HALF_N) {
    sBig = SECP256K1_N - sBig
    s = toPaddedHex(sBig, 32)
    sFlipped = true
  }

  // Determine the correct recovery parameter by trying both and comparing addresses
  const rBig = BigInt(r)
  const hashBig = BigInt('0x' + Buffer.from(ethSignedHash).toString('hex'))

  let correctV = null
  const recoveredAddresses = []
  for (const recoveryParam of [0, 1]) {
    try {
      const recoveredPubKey = ecrecover(hashBig, rBig, sBig, recoveryParam)
      if (!recoveredPubKey) {
        recoveredAddresses.push({ v: recoveryParam + 27, error: 'recovery returned null' })
        continue
      }
      
      const recoveredAddress = pubKeyToAddress(recoveredPubKey)
      recoveredAddresses.push({ v: recoveryParam + 27, address: recoveredAddress })
      
      if (recoveredAddress.toLowerCase() === signerAddress.toLowerCase()) {
        correctV = recoveryParam + 27
        break
      }
    } catch (err) {
      recoveredAddresses.push({ v: recoveryParam + 27, error: err.message })
    }
  }

  if (correctV === null) {
    return resp(500, { 
      error: 'Failed to determine correct signature recovery parameter. Verify signerAddress matches the KMS key.',
      debug: {
        expectedAddress: signerAddress,
        recoveredAddresses,
        r,
        s
      }
    })
  }

  const signature = hexConcat([r, s, toPaddedHex(BigInt(correctV), 1)])

  return {
    statusCode: 200,
    headers: { 'content-type': 'application/json' },
    body: JSON.stringify({ signature })
  }
}

function resp(statusCode, body) {
  return { statusCode, headers: { 'content-type': 'application/json' }, body: JSON.stringify(body) }
}

// EIP-191 personal sign: keccak256("\x19Ethereum Signed Message:\n" + len + message)
function hashMessage(messageBytes) {
  const prefix = Buffer.from(`\x19Ethereum Signed Message:\n${messageBytes.length}`, 'utf8')
  const prefixedMessage = Buffer.concat([prefix, messageBytes])
  return Buffer.from(keccak256(prefixedMessage))
}

function parse32ByteDigest(input) {
  if (typeof input !== 'string') return null
  if (input.startsWith('0x')) {
    try {
      const bytes = Buffer.from(input.slice(2), 'hex')
      return bytes.length === 32 ? bytes : null
    } catch {
      return null
    }
  }
  try {
    const bytes = Buffer.from(input, 'base64')
    return bytes.length === 32 ? bytes : null
  } catch {
    return null
  }
}

function derToRS(der) {
  let i = 0
  if (der[i++] !== 0x30) throw new Error('Invalid DER: expected SEQUENCE')
  const seq = readDerLen(der, i)
  i = seq.next
  if (der[i++] !== 0x02) throw new Error('Invalid DER: expected INTEGER (r)')
  const rLen = readDerLen(der, i)
  i = rLen.next
  const rBytes = der.slice(i, i + rLen.len)
  i += rLen.len
  if (der[i++] !== 0x02) throw new Error('Invalid DER: expected INTEGER (s)')
  const sLen = readDerLen(der, i)
  i = sLen.next
  const sBytes = der.slice(i, i + sLen.len)
  return { r: bytesToPaddedHex(stripLeadingZero(rBytes), 32), s: bytesToPaddedHex(stripLeadingZero(sBytes), 32) }
}

function readDerLen(buf, offset) {
  const first = buf[offset]
  if (first < 0x80) return { len: first, next: offset + 1 }
  const n = first & 0x7f
  let len = 0
  for (let j = 0; j < n; j++) len = (len << 8) | buf[offset + 1 + j]
  return { len, next: offset + 1 + n }
}

function stripLeadingZero(b) {
  return (b.length > 1 && b[0] === 0x00) ? b.slice(1) : b
}

function bytesToPaddedHex(b, size) {
  return '0x' + Buffer.from(b).toString('hex').padStart(size * 2, '0')
}

function toPaddedHex(n, size) {
  let hex = n.toString(16)
  if (hex.length > size * 2) hex = hex.slice(-size * 2)
  return '0x' + hex.padStart(size * 2, '0')
}

function hexConcat(parts) {
  return '0x' + parts.map((p) => (p.startsWith('0x') ? p.slice(2) : p)).join('')
}

// ============================================================================
// Pure JavaScript secp256k1 ecrecover implementation
// ============================================================================

// Modular arithmetic helpers
function mod(a, m) {
  const result = a % m
  return result >= 0n ? result : result + m
}

function modInverse(a, m) {
  a = mod(a, m) // Ensure a is positive before computing inverse
  if (a === 0n) throw new Error('No modular inverse for 0')
  let [old_r, r] = [a, m]
  let [old_s, s] = [1n, 0n]
  while (r !== 0n) {
    const q = old_r / r
    ;[old_r, r] = [r, old_r - q * r]
    ;[old_s, s] = [s, old_s - q * s]
  }
  return mod(old_s, m)
}

function modPow(base, exp, m) {
  let result = 1n
  base = mod(base, m)
  while (exp > 0n) {
    if (exp % 2n === 1n) result = mod(result * base, m)
    exp = exp / 2n
    base = mod(base * base, m)
  }
  return result
}

// Elliptic curve point operations (affine coordinates)
function pointAdd(p1, p2) {
  if (p1 === null) return p2
  if (p2 === null) return p1
  
  const [x1, y1] = p1
  const [x2, y2] = p2
  
  if (x1 === x2) {
    if (y1 !== y2) return null // point at infinity
    // Point doubling
    const s = mod(3n * x1 * x1 * modInverse(2n * y1, SECP256K1_P), SECP256K1_P)
    const x3 = mod(s * s - 2n * x1, SECP256K1_P)
    const y3 = mod(s * (x1 - x3) - y1, SECP256K1_P)
    return [x3, y3]
  }
  
  const s = mod((y2 - y1) * modInverse(x2 - x1, SECP256K1_P), SECP256K1_P)
  const x3 = mod(s * s - x1 - x2, SECP256K1_P)
  const y3 = mod(s * (x1 - x3) - y1, SECP256K1_P)
  return [x3, y3]
}

function pointMul(k, point) {
  let result = null
  let addend = point
  while (k > 0n) {
    if (k & 1n) result = pointAdd(result, addend)
    addend = pointAdd(addend, addend)
    k >>= 1n
  }
  return result
}

// Recover y-coordinate from x-coordinate on secp256k1
function recoverY(x, isOdd) {
  // y² = x³ + 7 (mod p)
  const ySquared = mod(modPow(x, 3n, SECP256K1_P) + 7n, SECP256K1_P)
  // y = ySquared^((p+1)/4) mod p (works because p ≡ 3 mod 4)
  const y = modPow(ySquared, (SECP256K1_P + 1n) / 4n, SECP256K1_P)
  // Verify
  if (mod(y * y, SECP256K1_P) !== ySquared) return null
  // Return the y with correct parity
  const yIsOdd = (y & 1n) === 1n
  return yIsOdd === isOdd ? y : mod(-y, SECP256K1_P)
}

// ecrecover: recover public key from signature
function ecrecover(msgHash, r, s, recoveryParam) {
  // Recover the point R from r and recoveryParam
  const isOdd = (recoveryParam & 1) === 1
  const x = r
  const y = recoverY(x, isOdd)
  if (y === null) return null
  
  const R = [x, y]
  const G = [SECP256K1_Gx, SECP256K1_Gy]
  
  // Public key P = r^(-1) * (s*R - e*G)
  const rInv = modInverse(r, SECP256K1_N)
  const u1 = mod(-msgHash * rInv, SECP256K1_N)
  const u2 = mod(s * rInv, SECP256K1_N)
  
  const point1 = pointMul(u1, G)
  const point2 = pointMul(u2, R)
  const pubKeyPoint = pointAdd(point1, point2)
  
  if (pubKeyPoint === null) return null
  
  // Return uncompressed public key (65 bytes: 0x04 || x || y)
  const xHex = pubKeyPoint[0].toString(16).padStart(64, '0')
  const yHex = pubKeyPoint[1].toString(16).padStart(64, '0')
  return Buffer.from('04' + xHex + yHex, 'hex')
}

function pubKeyToAddress(pubKey) {
  // Skip the 0x04 prefix and hash the 64 bytes
  const hash = keccak256(pubKey.slice(1))
  return '0x' + Buffer.from(hash.slice(12)).toString('hex')
}

// ============================================================================
// Pure JavaScript Keccak-256 implementation (no external dependencies)
// ============================================================================
const KECCAK_ROUNDS = 24
const ROTC = [1,3,6,10,15,21,28,36,45,55,2,14,27,41,56,8,25,43,62,18,39,61,20,44]
const PILN = [10,7,11,17,18,3,5,16,8,21,24,4,15,23,19,13,12,2,20,14,22,9,6,1]
const RNDC = [
  0x0000000000000001n, 0x0000000000008082n, 0x800000000000808an, 0x8000000080008000n,
  0x000000000000808bn, 0x0000000080000001n, 0x8000000080008081n, 0x8000000000008009n,
  0x000000000000008an, 0x0000000000000088n, 0x0000000080008009n, 0x000000008000000an,
  0x000000008000808bn, 0x800000000000008bn, 0x8000000000008089n, 0x8000000000008003n,
  0x8000000000008002n, 0x8000000000000080n, 0x000000000000800an, 0x800000008000000an,
  0x8000000080008081n, 0x8000000000008080n, 0x0000000080000001n, 0x8000000080008008n
]

function rotl64(x, n) {
  n = BigInt(n)
  return ((x << n) | (x >> (64n - n))) & 0xffffffffffffffffn
}

function keccakF(state) {
  for (let round = 0; round < KECCAK_ROUNDS; round++) {
    // Theta
    const C = new Array(5)
    for (let x = 0; x < 5; x++) {
      C[x] = state[x] ^ state[x + 5] ^ state[x + 10] ^ state[x + 15] ^ state[x + 20]
    }
    const D = new Array(5)
    for (let x = 0; x < 5; x++) {
      D[x] = C[(x + 4) % 5] ^ rotl64(C[(x + 1) % 5], 1)
    }
    for (let x = 0; x < 5; x++) {
      for (let y = 0; y < 5; y++) {
        state[x + y * 5] ^= D[x]
      }
    }

    // Rho and Pi
    let t = state[1]
    for (let i = 0; i < 24; i++) {
      const j = PILN[i]
      const temp = state[j]
      state[j] = rotl64(t, ROTC[i])
      t = temp
    }

    // Chi
    for (let y = 0; y < 5; y++) {
      const row = [state[y * 5], state[y * 5 + 1], state[y * 5 + 2], state[y * 5 + 3], state[y * 5 + 4]]
      for (let x = 0; x < 5; x++) {
        state[y * 5 + x] = row[x] ^ ((~row[(x + 1) % 5]) & row[(x + 2) % 5])
      }
    }

    // Iota
    state[0] ^= RNDC[round]
  }
}

function keccak256(input) {
  const inputBytes = Buffer.isBuffer(input) ? input : Buffer.from(input)
  const rate = 136 // (1600 - 256 * 2) / 8 for keccak-256
  const state = new Array(25).fill(0n)

  // Pad the input: append 0x01, then zeros, then 0x80
  const padLen = rate - (inputBytes.length % rate)
  const padded = Buffer.alloc(inputBytes.length + padLen)
  inputBytes.copy(padded)
  padded[inputBytes.length] = 0x01
  padded[padded.length - 1] |= 0x80

  // Absorb
  for (let offset = 0; offset < padded.length; offset += rate) {
    for (let i = 0; i < rate / 8; i++) {
      const lane = padded.readBigUInt64LE(offset + i * 8)
      state[i] ^= lane
    }
    keccakF(state)
  }

  // Squeeze (just 32 bytes for keccak-256)
  const output = Buffer.alloc(32)
  for (let i = 0; i < 4; i++) {
    output.writeBigUInt64LE(state[i], i * 8)
  }
  return output
}