authProcedure.ts

import { cookies } from "next/headers"
import { db } from "~/server/db/client"
import { procedureFactory } from "~/server/rpc/procedureFactory"
import { z } from "zod"

const getAuthContext = async () => {
  const cookie = cookies().get("auth")
  if (!cookie) {
    return { user: null }
  }

  const user = await db
    .selectFrom("session")
    .where("key", "=", cookie.value)
    .where("expiringAt", ">", new Date())
    .innerJoin("auth_user", "auth_user.id", "session.userId")
    .select(["auth_user.id", "auth_user.createdAt", "auth_user.username"])
    .executeTakeFirst()

  return {
    user: user ?? null,
  }
}

export const procedure = procedureFactory(getAuthContext)

rpc.ts

import { type z } from "zod"

type ContextT<ContextProviderReturnT> = () =>
  | Promise<ContextProviderReturnT>
  | ContextProviderReturnT

type ContextInputT<ContextProviderReturnT> = {
  ctx: ContextProviderReturnT
}

type CallbackInputT<SchemaT> = {
  data: z.output<z.ZodType<SchemaT>>
}

type CallbackT<SchemaT, ContextProviderReturnT> =
  ContextInputT<ContextProviderReturnT> & CallbackInputT<SchemaT>

type AuthT<ContextProviderReturnT> = ContextInputT<ContextProviderReturnT>

type PermissionT<SchemaT, ContextProviderReturnT> =
  ContextInputT<ContextProviderReturnT> & CallbackInputT<SchemaT>

type CodeT = 200 | 400 | 401 | 403 | 500

type ProcedureReturnT<SchemaT, CallbackReturnT> = Promise<
  | {
      success: false
      code: CodeT
      error?: z.ZodFormattedError<z.ZodType<SchemaT>>
    }
  | {
      success: true
      code: CodeT
      out: CallbackReturnT
    }
>

type ProcedureT<SchemaT, CallbackReturnT, ContextProviderReturnT> = {
  schema: z.ZodType<SchemaT>
  auth: (params: AuthT<ContextProviderReturnT>) => Promise<boolean> | boolean
  permission: (
    params: PermissionT<SchemaT, ContextProviderReturnT>,
  ) => Promise<boolean> | boolean
  callback: (
    params: CallbackT<SchemaT, ContextProviderReturnT>,
  ) => Promise<CallbackReturnT> | CallbackReturnT
}

export const procedureFactory =
  <ContextProviderReturnT>(contextProvider: ContextT<ContextProviderReturnT>) =>
  <T, R>({
    schema,
    auth,
    permission,
    callback,
  }: ProcedureT<T, R, ContextProviderReturnT>) =>
  async (input?: unknown): ProcedureReturnT<T, R> => {
    const ctx = await contextProvider()

    const authPassed = await auth({ ctx })
    if (!authPassed) {
      return {
        success: false,
        code: 401,
      }
    }

    const inputValidationResult = schema.safeParse(input)
    if (!inputValidationResult.success) {
      return {
        success: false,
        code: 400,
        error: inputValidationResult.error.format(),
      }
    }

    const data = inputValidationResult.data

    const hasPermission = await permission({ ctx, data })
    if (!hasPermission) {
      return {
        success: false,
        code: 403,
      }
    }

    try {
      return {
        success: true,
        code: 200,
        out: await callback({ ctx, data }),
      }
    } catch (e) {
      console.log(e)
      return {
        success: false,
        code: 500,
      }
    }
  }