oprfServerless.ts

import crypto from 'crypto';
import { APIGatewayProxyHandler, APIGatewayProxyResult } from 'aws-lambda';
import { DynamoDB } from 'aws-sdk';
import base64url from 'base64url';
import {
  base64urlToUtf8,
  hexToUtf8,
  isValidBase64UrlInput,
  isValidHexInput,
  isValidId,
  isValidUtf8Input,
  utf8ToHex,
} from './encodingUtils';

const OPRF = require('oprf');

const DEFAULT_TABLE_NAME = 'oprf-users';
const TABLE_NAME = process.env.TABLE_NAME ?? DEFAULT_TABLE_NAME;
const RESPONSE_DELAY_MS = 1000;

const dynamoDB = new DynamoDB.DocumentClient();
const oprfInstance = new OPRF();

type OprfRequestBody = {
  id?: unknown;
  input?: unknown;
};

function createResponse(statusCode: number, payload: Record<string, unknown>): APIGatewayProxyResult {
  return {
    statusCode,
    headers: {
      'Content-Type': 'application/json',
      'Cache-Control': 'no-store',
    },
    body: JSON.stringify(payload),
  };
}

function parseBody(body: string | null): OprfRequestBody {
  if (!body) {
    return {};
  }

  try {
    const parsed = JSON.parse(body) as OprfRequestBody;
    return parsed ?? {};
  } catch {
    return {};
  }
}

async function getOrCreateSecretKey(id: string): Promise<Uint8Array> {
  const result = await dynamoDB
    .get({
      TableName: TABLE_NAME,
      Key: { id },
    })
    .promise();

  if (result.Item?.secretKey) {
    return new Uint8Array(result.Item.secretKey as number[]);
  }

  const secretKey = oprfInstance.generateRandomScalar();
  await dynamoDB
    .put({
      TableName: TABLE_NAME,
      Item: { id, secretKey: Array.from(secretKey) },
    })
    .promise();

  return secretKey;
}

async function createRandomId(): Promise<string> {
  while (true) {
    const randomId = crypto.randomBytes(16).toString('base64url');
    const result = await dynamoDB
      .get({
        TableName: TABLE_NAME,
        Key: { id: randomId },
      })
      .promise();

    if (!result.Item?.id) {
      return randomId;
    }
  }
}

export const handler: APIGatewayProxyHandler = async (event) => {
  const timeout = new Promise((resolve) => setTimeout(resolve, RESPONSE_DELAY_MS));
  const parsedBody = parseBody(event.body);
  const getParams = event.queryStringParameters;

  if (!getParams && !parsedBody.input) {
    return createResponse(400, {
      error:
        'Provide a masked point via "input" in UTF-8, Base64Url, or Hex. Optionally provide "id" to reuse a salt key.',
    });
  }

  const rawId = getParams?.id ?? parsedBody.id;
  const rawInput = getParams?.input ?? parsedBody.input;

  if (rawInput === undefined || rawInput === null || rawInput === '') {
    return createResponse(400, { error: 'A masked point is required as "input".' });
  }

  if (rawId !== undefined && rawId !== null && !isValidId(rawId)) {
    return createResponse(400, { error: 'Invalid ID format.' });
  }

  if (typeof rawInput !== 'string') {
    return createResponse(400, { error: 'Invalid input format.' });
  }

  const inputValue = rawInput;
  const inputLength = (inputValue as string).length;
  const inputIsUtf8 = isValidUtf8Input(inputValue);
  const inputIsBase64Url = isValidBase64UrlInput(inputValue);
  const inputIsHex = isValidHexInput(inputValue);

  if (!inputIsUtf8 && !inputIsBase64Url && !inputIsHex) {
    return createResponse(400, { error: `Invalid input format. Input length: ${inputLength}` });
  }

  const encodedMaskedPoint = inputIsBase64Url
    ? base64urlToUtf8(inputValue)
    : inputIsHex
      ? hexToUtf8(inputValue)
      : inputValue;

  const id = typeof rawId === 'string' ? rawId : await createRandomId();

  await oprfInstance.ready;
  const secretKey = await getOrCreateSecretKey(id);

  let maskedPoint: Uint8Array;
  try {
    maskedPoint = oprfInstance.decodePoint(encodedMaskedPoint, 'UTF-8');
  } catch {
    return createResponse(400, { error: 'Decoding the masked point failed.' });
  }

  if (!oprfInstance.isValidPoint(maskedPoint)) {
    return createResponse(400, { error: 'Masked point is not valid on the curve.' });
  }

  let saltedPoint: Uint8Array;
  try {
    saltedPoint = oprfInstance.scalarMult(maskedPoint, secretKey);
  } catch {
    return createResponse(400, { error: 'Scalar multiplication failed.' });
  }

  if (!oprfInstance.isValidPoint(saltedPoint)) {
    return createResponse(400, { error: 'Salted point is not valid on the curve.' });
  }

  const encodedSaltedPoint = oprfInstance.encodePoint(saltedPoint, 'UTF-8');
  const output = inputIsBase64Url
    ? base64url(encodedSaltedPoint)
    : inputIsHex
      ? utf8ToHex(encodedSaltedPoint)
      : encodedSaltedPoint;

  await timeout;
  return createResponse(200, { id, output });
};