lambda.ts

import * as lib from 'synapse:lib'
import * as aws from 'terraform-provider:aws'
import * as Lambda from '@aws-sdk/client-lambda'

interface LambdaOptions {
    arch?: 'arm64' | 'x64'
    env?: Record<string, string>
    bucket?: aws.S3Bucket
}

let assetBucket: aws.S3Bucket
function getAssetBucket() {
    return assetBucket ??= new aws.S3Bucket()
}

export class LambdaFunction<T extends any[] = any[], U = unknown> {
    private readonly client = new Lambda.Lambda({})
    public readonly resource: aws.LambdaFunction
    public readonly principal: aws.IamRole

    public constructor(target: (...args: T) => Promise<U> | U, opt?: LambdaOptions) {
        const arch = opt?.arch ?? 'arm64'
        const handler = `handler.default`
        const environment = opt?.env ?? {}

        const entrypoint = new lib.Bundle(wrap(target), {
            os: 'linux',
            arch,
            includeAssets: true, // Currently required when using Zig modules
        })

        const role = new aws.IamRole({
            assumeRolePolicy: JSON.stringify(spPolicy('lambda.amazonaws.com')),
            managedPolicyArns: ['arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'],
        })

        this.principal = role

        const bucket = opt?.bucket ?? getAssetBucket()

        const code = new lib.Archive(entrypoint)
        const obj = new aws.S3Object({
            bucket: bucket.bucket,
            key: entrypoint.destination,
            source: code.filePath,
        })

        const architectures = arch === 'arm64' ? [arch] 
            : arch === 'x64' ? ['x86_64'] : undefined

        const fn = new aws.LambdaFunction({
            functionName: lib.generateIdentifier(aws.LambdaFunction, 'functionName', 64),
            s3Bucket: obj.bucket,
            s3Key: obj.key,
            handler,
            architectures,
            timeout: 30,
            memorySize: 1024,
            packageType: 'Zip',
            runtime: 'nodejs20.x',
            role: role.arn,
            environment: {
                variables: environment,
            },
        })
        this.resource = fn
    }

    public async invoke(...args: T): Promise<U> {
        const payload = args.length > 0 ? Buffer.from(JSON.stringify({ args })) : undefined
        const resp = await this.client.invoke({
            FunctionName: this.resource.functionName,
            Payload: payload,
        })

        const resultString = Buffer.from(resp.Payload).toString('utf-8')
        if (resultString === 'null') {
            return void 0 as U
        }

        if (resp.FunctionError) {
            const errResp = JSON.parse(resultString) as { errorType: string; errorMessage: string; trace: string }
            const err = new Error(errResp.errorMessage)
            err.name = errResp.errorType
            throw err
        }

        return JSON.parse(resultString)
    }
}

type LambdaHandler = (event: any) => Promise<any>

function wrap<T, U extends any[]>(fn: (...args: U) => T): LambdaHandler {
    return async (event) => {
        if (typeof event !== 'object' || !Array.isArray(event?.args)) {
            throw new Error('Got unexpected event payload')
        }

        return fn(...event.args)
    }
}

const spPolicy = (principal: string | string[]) => ({
    Version: "2012-10-17",
    Statement: [{
        Effect: "Allow",
        Action: "sts:AssumeRole",
        Principal: { Service: principal },
    }]
})