lambda.ts

import path from "path"
import { APIGatewayProxyEvent } from "aws-lambda"
import { streamifyResponse, ResponseStream } from "lambda-stream"
import { S3Client, GetObjectCommand, HeadObjectCommand } from "@aws-sdk/client-s3"
import { Upload } from "@aws-sdk/lib-storage"

const {
	THIRD_PARTY_ASSETS_BUCKET,
} = process.env

const ASSET_ORIGIN = "https://example.com"

const s3 = new S3Client({})

function readChunks(reader: ReadableStreamDefaultReader<Uint8Array>) {
	return {
		async *[Symbol.asyncIterator]() {
			let readResult = await reader.read()

			while (!readResult.done) {
				yield readResult.value
				readResult = await reader.read()
			}
		},
	}
}

async function doesObjectExist(key: string): Promise<boolean> {
	const command = new HeadObjectCommand({
		Bucket: THIRD_PARTY_ASSETS_BUCKET,
		Key: key,
	})

	try {
		await s3.send(command)

		return true
	} catch (error) {
		// TODO(anthony): Find out correct error type.
		if (error.errorType === "NotFound") {
			return false
		}
	}

	return false
}

async function writeRequestBodyToResponseStream(body: ReadableStream<Uint8Array>, responseStream: ResponseStream) {
	const reader = body.getReader()

	for await (const chunk of readChunks(reader)) {
		responseStream.write(chunk)
	}
}

async function writeGetObjectBodyToResponseStream(
	body: SdkStream<internal.Readable | ReadableStream<any> | Blob | undefined>,
	responseStream: ResponseStream
) {
	return new Promise((resolve, reject) => {
		body.on("data", chunk => {
			responseStream.write(chunk)
		})

		body.on("error", reject)
		body.on("end", resolve)
	})
}

export const handler = streamifyResponse(thirdPartyAssetsHandler)

async function thirdPartyAssetsHandler(event: APIGatewayProxyEvent, responseStream: ResponseStream): Promise<void> {
	const filePath = "cool_file.woff2"
	const key = path.join("third-party-assets", filePath)

	const cachedFileExists = await doesObjectExist(key)

	if (cachedFileExists) {
		const getCommand = new GetObjectCommand({
			Bucket: THIRD_PARTY_ASSETS_BUCKET,
			Key: key,
		})
		const { Body: body } = await s3.send(getCommand)

		if (!body) {
			return
		}

		responseStream.setContentType("font/woff2")
		await writeGetObjectBodyToResponseStream(body, responseStream)
		responseStream.end()
		return
	}

	const remoteFilePath = path.join(ASSET_ORIGIN, filePath)
	const response = await fetch(remoteFilePath)

	if (!response.body || response.body === null) {
		return
	}

	const streamUpload = new Upload({
		client: s3,
		params: {
			Bucket: THIRD_PARTY_ASSETS_BUCKET,
			Key: key,
			Body: response.body,
		},
	})

	responseStream.setContentType("font/woff2")

	await streamUpload.done()
	await writeRequestBodyToResponseStream(response.body, responseStream)

  // This doesn't work:
	await Promise.all([writeRequestBodyToResponseStream(response.body, responseStream), streamUpload.done()])

  // This does but is not in parallel:
  await streamUpload.done()
	await writeRequestBodyToResponseStream(response.body, responseStream)
	responseStream.end()
}