music-repo/packages/backend/src/services/Authentication.ts

import { config } from "backend/config";
import { BunRequest } from "bun";
import { SessionId, UserId } from "common";
import { fetch, FetchError } from "common/Fetch";
import { Me, NotFound, Other, Role, Unauthenticated } from "common/the_api";
import { Context, DateTime, Duration, Effect, HashMap, HashSet, Layer, Option, pipe, Redacted, Schema } from "effect";
import { constant } from "effect/Function";
import { sql } from "kysely";
import * as Database from "./Database";

export interface AuthenticationInterface {
	readonly me: Effect.Effect<Me, Unauthenticated>;
	readonly logout: Effect.Effect<void>;
	readonly sessionId: SessionId;
	readonly getUser: (userId: UserId) => Effect.Effect<Other, FetchError | Unauthenticated | NotFound>;
}

export class Authentication extends Context.Tag("Authentication")<Authentication, AuthenticationInterface>() { }

export const OAUTH_SCOPE = "email openid profile offline_access";
export const REDIRECT_URI = config.NODE_ENV === "production" ? `https://${config.HOSTNAME}/login` : "http://localhost:3000/login";

export const EXPIRATION_BUFFER = Duration.seconds(10);

export const SESSION_COOKIE_NAME = "sessionId";

const ALL_ROLES = HashSet.fromIterable(Object.values(Role));

export const Live = (request: BunRequest) => Layer.effect(Authentication, Effect.gen(function* () {
	const database = yield* Database.Database;

	yield* database
		.deleteFrom("Session")
		.where(sql`datetime()`, ">=", "expiresAt")
		.$call(Database.execute);

	const sessionId = pipe(
		request.cookies.get(SESSION_COOKIE_NAME),
		Option.fromNullable,
		Option.map(SessionId.make),
		Option.getOrElse(generateSessionId),
	);

	request.cookies.set(SESSION_COOKIE_NAME, sessionId, {
		expires: yield* pipe(
			DateTime.now,
			Effect.map(DateTime.addDuration("7 days")),
			Effect.map(DateTime.toDateUtc),
		),
		httpOnly: true,
		sameSite: "none",
		secure: true,
	});

	const returning = [
		"sessionId",
		"accessToken",
		"codeVerifier",
		"idToken",
		"refreshToken",
		"state",
	] as const;

	const session = yield* database
		.updateTable("Session")
		.set({ expiresAt: sql`datetime('now', '+7 days') ` })
		.where("sessionId", "=", sessionId)
		.returning(returning)
		.$call(Database.executeTakeFirst)
		.pipe(Effect.catchTag("NoSuchElementException", () => database
			.insertInto("Session")
			.values({ sessionId, expiresAt: sql`datetime('now', '+7 days')` })
			.returning(returning)
			.$call(Database.executeTakeFirstOrDie)
		));

	const { accessToken, idToken, refreshToken } = yield* revalidateTokens({
		accessToken: Option.fromNullable(session.accessToken),
		idToken: Option.fromNullable(session.idToken),
		refreshToken: Option.fromNullable(session.refreshToken),
	});

	yield* database
		.updateTable("Session")
		.set({
			accessToken: pipe(
				accessToken,
				Option.map((at) => at.token),
				Option.getOrNull,
			),
			idToken: pipe(
				idToken,
				Option.map((it) => it.token),
				Option.getOrNull,
			),
			refreshToken: Option.getOrNull(refreshToken),
		})
		.$call(Database.execute);

	const state = Object.freeze({
		sessionId: session.sessionId,
		accessToken,
		idToken,
		refreshToken,
		codeVerifier: Option.fromNullable(session.codeVerifier),
		state: Option.fromNullable(session.state),
	});

	return Object.freeze<AuthenticationInterface>({
		me: pipe(
			state.idToken,
			Option.map(({ payload }) => {
				const userData: Me = Object.freeze<Me>({
					userId: payload.sub,
					displayName: payload.display_name,
					roles: ALL_ROLES,
				});
				return userData;
			}),
			Option.match({
				onNone: () => Effect.fail(Unauthenticated.make()),
				onSome: Effect.succeed,
			}),
		),
		logout: database
			.updateTable("Session")
			.set({
				accessToken: null,
				codeVerifier: null,
				expiresAt: sql`datetime('now', '+7 days')`,
				idToken: null,
				refreshToken: null,
				state: null,
			})
			.where("sessionId", "=", state.sessionId)
			.$call(Database.execute),
		sessionId: state.sessionId,
		getUser: (userId) => pipe(
			getDisplayName(userId),
			Effect.map((displayName) => Object.freeze<Other>({ userId, displayName })),
		),
	});
}));

export const Test = (me: Option.Option<Me>, other: HashMap.HashMap<UserId, Other>) => Layer.sync(Authentication, constant(Object.freeze<AuthenticationInterface>({
	me: Option.match(me, {
		onNone: () => Effect.fail(Unauthenticated.make()),
		onSome: (me) => Effect.succeed(me),
	}),
	logout: Effect.void,
	sessionId: generateSessionId(),
	getUser: Option.match(me, {
		onNone: () => constant(Effect.fail(Unauthenticated.make())),
		onSome: () => (userId) => pipe(
			other,
			HashMap.get(userId),
			Option.match({
				onNone: () => Effect.fail(NotFound.make()),
				onSome: Effect.succeed,
			}),
		),
	}),
})));

export const AccessTokenPayload = Schema.Struct({
	aud: pipe(
		Schema.String,
		Schema.HashSet,
	),
	exp: Schema.DateTimeUtc,
	iat: Schema.DateTimeUtc,
	iss: Schema.String,
	sub: UserId,
});

export const IdTokenPayload = Schema.Struct({
	aud: pipe(
		Schema.String,
		Schema.HashSet,
	),
	exp: Schema.DateTimeUtc,
	iat: Schema.DateTimeUtc,
	iss: Schema.String,
	sub: UserId,

	name: Schema.String,
	given_name: Schema.String,
	family_name: Schema.String,
	display_name: Schema.String,
	preferred_username: Schema.String,

	email: Schema.String,
	email_verified: Schema.Boolean,

	picture: Schema.String,
});

export type AccessTokenPayload = typeof AccessTokenPayload.Type;
export type IdTokenPayload = typeof IdTokenPayload.Type;

function generateCodeVerifier(byteLength: number = 32) {
	const codeVerifierBytes = new Uint8Array(byteLength);
	crypto.getRandomValues(codeVerifierBytes);
	const codeVerifier = Buffer.from(codeVerifierBytes).toString("base64url");

	const codeVerifierAsciiBuffer = Buffer.from(codeVerifier, "ascii");
	const codeVerifierAsciiArray = new Uint8Array(
		codeVerifierAsciiBuffer.buffer,
		codeVerifierAsciiBuffer.byteOffset,
		codeVerifierAsciiBuffer.length,
	);
	const codeChallenge = Bun.SHA256.hash(codeVerifierAsciiArray, "base64url");

	return { codeVerifier, codeChallenge };
}

function generateSessionId(byteLength: number = 32): SessionId {
	const array = new Uint8Array(byteLength);
	crypto.getRandomValues(array);
	const string = Buffer.from(array).toString("base64url");
	return SessionId.make(string);
};

function generateRandomState(byteLength: number = 32): string {
	const array = new Uint8Array(byteLength);
	crypto.getRandomValues(array);
	const state = Buffer.from(array).toString("base64url");
	return state;
}

const getJwtTokenPayload = <A, I, R>(schema: Schema.Schema<A, I, R>) => {
	const decoder = Schema.decodeUnknown(schema);
	return (token: string) => {
		const json = JSON.parse(Buffer.from(token.split(".")[1], "base64url").toString("utf-8"));
		return pipe(
			decoder(json),
			Effect.orDie,
		);
	};
};

const getDisplayName = (userId: UserId) => Effect.gen(function* () {
	const url = new URL(`/api/users/${userId}`, config.POCKET_ID_API_ORIGIN);
	const res = yield* fetch(url, {
		headers: {
			"X-API-KEY": config.POCKET_ID_API_KEY,
		},
	});

	if (!res.ok) {
		return yield* Effect.fail(NotFound.make());
	}

	const json = yield* Effect.promise(() => res.json());
	const displayName: string = json.displayName;
	return displayName;
});

export const makeAuthorizationUrl = (sessionId: SessionId) => Effect.gen(function* () {
	const database = yield* Database.Database;

	const { codeVerifier, codeChallenge } = generateCodeVerifier();
	const state = generateRandomState();

	yield* database
		.updateTable("Session")
		.set({
			codeVerifier,
			state,
			accessToken: null,
			idToken: null,
			refreshToken: null,
		})
		.where("sessionId", "=", sessionId)
		.$call(Database.execute);

	const url = new URL(config.OAUTH_AUTHORIZATION_ENDPOINT);
	url.searchParams.set("client_id", config.CLIENT_ID);
	url.searchParams.set("response_type", "code");
	url.searchParams.set("redirect_uri", REDIRECT_URI);
	url.searchParams.set("scope", OAUTH_SCOPE);
	url.searchParams.set("response_mode", "form_post");
	url.searchParams.set("state", state);
	url.searchParams.set("code_challenge", codeChallenge);
	url.searchParams.set("code_challenge_method", "S256");

	return url.toString();
});

namespace revaildateTokens {
	export interface Args {
		readonly accessToken: Option.Option<string>;
		readonly idToken: Option.Option<string>;
		readonly refreshToken: Option.Option<string>;
	}

	export interface Result {
		readonly accessToken: Option.Option<{
			readonly token: string,
			readonly payload: AccessTokenPayload,
		}>;
		readonly idToken: Option.Option<{
			readonly token: string,
			readonly payload: IdTokenPayload,
		}>;
		readonly refreshToken: Option.Option<string>;
	}
}

const revalidateTokens = ({ accessToken, idToken, refreshToken }: revaildateTokens.Args) => Effect.gen(function* () {

	const accessTokenPayload = yield* pipe(
		accessToken,
		Effect.transposeMapOption(getJwtTokenPayload(AccessTokenPayload)),
	);

	const idTokenPayload = yield* pipe(
		idToken,
		Effect.transposeMapOption(getJwtTokenPayload(IdTokenPayload)),
	);

	const expirationThreshold = DateTime.addDuration(
		DateTime.unsafeNow(),
		EXPIRATION_BUFFER,
	);

	// Token expired or missing
	if (Option.match(accessTokenPayload, {
		onNone: constant(false),
		onSome: (atp) => DateTime.greaterThan(expirationThreshold, atp.exp),
	}) || Option.match(idTokenPayload, {
		onNone: constant(false),
		onSome: (itp) => DateTime.greaterThan(expirationThreshold, itp.exp),
	})) {

		accessToken = Option.none();
		idToken = Option.none();

		// try refreshing
		if (Option.isSome(refreshToken)) {
			const refreshTokenValue = refreshToken.value;

			const res = yield* fetch(config.OAUTH_TOKEN_ENDPOINT, {
				method: "POST",
				headers: {
					"Content-Type": "application/x-www-form-urlencoded",
				},
				body: new URLSearchParams({
					"client_id": config.CLIENT_ID,
					"grant_type": "refresh_token",
					"refresh_token": refreshTokenValue,
					"client_secret": Redacted.value(config.CLIENT_SECRET),
				}).toString(),
			});

			const json: {
				access_token: string,
				refresh_token: string,
				id_token: string,
			} = yield* Effect.promise(() => res.json());

			accessToken = Option.some(json.access_token);
			idToken = Option.some(json.id_token);
			refreshToken = Option.some(json.refresh_token);
		}
	}

	const it = yield* pipe(
		idToken,
		Option.map((it) => pipe(
			Effect.Do,
			Effect.let("token", () => it),
			Effect.bind("payload", ({ token }) => getJwtTokenPayload(IdTokenPayload)(token)),
		)),
		Effect.transposeOption,
		Effect.map(Option.map((it) => Object.freeze(it))),
	);

	const at = yield* pipe(
		accessToken,
		Option.map((at) => pipe(
			Effect.Do,
			Effect.let("token", () => at),
			Effect.bind("payload", ({ token }) => getJwtTokenPayload(AccessTokenPayload)(token)),
		)),
		Effect.transposeOption,
		Effect.map(Option.map((at) => Object.freeze(at))),
	);

	const res: revaildateTokens.Result = Object.freeze<revaildateTokens.Result>({
		accessToken: at,
		idToken: it,
		refreshToken,
	});

	return res;
});