tRPC server (running on hono with adapter) - running on Cloudflare Workers - websocket connection in trpc context

context.ts

import { type inferAsyncReturnType } from '@trpc/server';
import { Context as HonoContext } from 'hono';
import { getCookie } from 'hono/cookie';
import { Lucia, Session, User, verifyRequestOrigin } from 'lucia';

import { createAuth, getAllowedOriginHost } from '../auth';
import {
  NeonDBWebSocketWithSchemas,
  NeonDBWithSchemas,
  createDBWebSocket,
  createDbHTTP,
} from './db/client';

interface ApiContextProps {
  session?: Session;
  user?: User;
  auth: Lucia;
  db: NeonDBWithSchemas;
  dbWebSockets: NeonDBWebSocketWithSchemas;
  c: HonoContext;
  AWS_ACCESS_KEY: string;
  AWS_SECRET_KEY: string;
  AWS_SES_REGION: string;
  GITHUB_CLIENT_ID: string;
  GITHUB_CLIENT_SECRET: string;
  GOOGLE_CLIENT_ID: string;
  GOOGLE_CLIENT_SECRET: string;
  GOOGLE_REDIRECT_URI: string;
}

export const createContext = async (
  c: HonoContext,
  DATABASE_URL: string,
  AWS_ACCESS_KEY: string,
  AWS_SECRET_KEY: string,
  AWS_SES_REGION: string,
  GITHUB_CLIENT_ID: string,
  GITHUB_CLIENT_SECRET: string,
  GOOGLE_CLIENT_ID: string,
  GOOGLE_CLIENT_SECRET: string,
  GOOGLE_REDIRECT_URI: string,
): Promise<ApiContextProps> => {
  const db = createDbHTTP(DATABASE_URL);
  const dbWebSockets = createDBWebSocket(DATABASE_URL);
  const auth = createAuth(db);

  const enableTokens = Boolean(c.req.header('x-enable-tokens'));

  async function getSession() {
    let user: User | undefined;
    let session: Session | undefined;
    const res = {
      user,
      session,
    };

    if (!c.req) return res;

    const cookieSessionId = getCookie(c, auth.sessionCookieName);
    const bearerSessionId =
      enableTokens && c.req.header('authorization')?.split(' ')[1];

    if (!cookieSessionId && !bearerSessionId) return res;

    let authResult:
      | Awaited<ReturnType<typeof auth.validateSession>>
      | undefined;

    if (cookieSessionId && !enableTokens) {
      const originHeader = c.req.header('origin');
      const hostHeader = c.req.header('host');
      const allowedOrigin = getAllowedOriginHost(c.env.APP_URL, c.req.raw);
      if (
        originHeader &&
        hostHeader &&
        verifyRequestOrigin(originHeader, [
          hostHeader,
          ...(allowedOrigin ? [allowedOrigin] : []),
        ])
      ) {
        authResult = await auth.validateSession(cookieSessionId);
        if (authResult.session?.fresh) {
          c.header(
            'Set-Cookie',
            auth.createSessionCookie(authResult.session.id).serialize(),
            {
              append: true,
            },
          );
        }
        if (!authResult?.session) {
          c.header('Set-Cookie', auth.createBlankSessionCookie().serialize(), {
            append: true,
          });
        }
      } else {
        console.log('CSRF failed', {
          cookieSessionId,
          originHeader,
          hostHeader,
          allowedOrigin,
        });
      }
    }

    if (bearerSessionId) {
      authResult = await auth.validateSession(bearerSessionId);
    }

    res.session = authResult?.session || undefined;
    res.user = authResult?.user || undefined;
    return res;
  }

  const { session, user } = await getSession();

  return {
    auth,
    db,
    dbWebSockets,
    session,
    user,
    c,
    AWS_ACCESS_KEY,
    AWS_SECRET_KEY,
    AWS_SES_REGION,
    GITHUB_CLIENT_ID,
    GITHUB_CLIENT_SECRET,
    GOOGLE_CLIENT_ID,
    GOOGLE_CLIENT_SECRET,
    GOOGLE_REDIRECT_URI,
  };
};

export type Context = inferAsyncReturnType<typeof createContext>;

dbClients.ts

import { Pool, neon } from '@neondatabase/serverless';
import { drizzle } from 'drizzle-orm/neon-http';
import { drizzle as drizzleServerless } from 'drizzle-orm/neon-serverless';

import * as bookingSchema from './schema/booking';
import * as bookingSlotSchema from './schema/bookingSlot';
import * as emailVerificationCodeSchema from './schema/emailVerificationCode';
import * as machineSchema from './schema/machine';
import * as oauthSchema from './schema/oauth';
import * as phoneVerificationCodeSchema from './schema/phoneVerificationCode';
import * as roomSchema from './schema/room';
import * as roomScheduleSchema from './schema/roomSchedule';
import * as sessionSchema from './schema/session';
import * as subscriptionSchema from './schema/subscription';
import * as userSchema from './schema/user';

const schema = {
  ...bookingSchema,
  ...bookingSlotSchema,
  ...emailVerificationCodeSchema,
  ...machineSchema,
  ...oauthSchema,
  ...phoneVerificationCodeSchema,
  ...roomSchema,
  ...roomScheduleSchema,
  ...sessionSchema,
  ...subscriptionSchema,
  ...userSchema,
};

export const createDbHTTP = (DATABASE_URL: string) => {
  const sql = neon(DATABASE_URL);
  return drizzle(sql, {
    schema,
    logger: true,
  });
};

export const createDBWebSocket = (DATABASE_URL: string) => {
  const pool = new Pool({ connectionString: DATABASE_URL });
  return drizzleServerless(pool, {
    schema,
    logger: true,
  });
};

export type NeonDBWithSchemas = ReturnType<typeof createDbHTTP>;
export type NeonDBWebSocketWithSchemas = ReturnType<typeof createDBWebSocket>;

trpc.ts

import { TRPCError, initTRPC } from '@trpc/server';
import superJson from 'superjson';

import { type Context } from './context';

const t = initTRPC.context<Context>().create({
  transformer: superJson,
  errorFormatter({ shape }) {
    return shape;
  },
});

/**
 * This is a middleware that checks if the user is authenticated
 * @link https://trpc.io/docs/middlewares
 */
const isAuthenticated = t.middleware(({ next, ctx }) => {
  if (!ctx.session || !ctx.user) {
    throw new TRPCError({ code: 'UNAUTHORIZED', message: 'Not authenticated' });
  }
  return next({
    ctx: {
      user: ctx.user,
      session: ctx.session,
    },
  });
});

export const router = t.router;
export const publicProcedure = t.procedure;
export const protectedProcedure = t.procedure.use(isAuthenticated);

worker.ts

import { trpcServer } from '@hono/trpc-server';
import { Hono } from 'hono';
import { prettyJSON } from 'hono/pretty-json';

import { createContext } from './context';
import OAuth from './oauth/index';
import { appRouter } from './router';

export type Bindings = {
  DATABASE_URL: string;
  AWS_ACCESS_KEY: string;
  AWS_SECRET_KEY: string;
  AWS_SES_REGION: string;
  GITHUB_CLIENT_ID: string;
  GITHUB_CLIENT_SECRET: string;
  GOOGLE_CLIENT_ID: string;
  GOOGLE_CLIENT_SECRET: string;
  GOOGLE_REDIRECT_URI: string;
};

const app = new Hono<{ Bindings: Bindings }>();

// Setup TRPC server with context
app.use('/trpc/*', async (c, next) => {
  return await trpcServer({
    router: appRouter,
    createContext: async () => {
      return await createContext(
        c,
        c.env.DATABASE_URL,
        c.env.AWS_ACCESS_KEY,
        c.env.AWS_SECRET_KEY,
        c.env.AWS_SES_REGION,
        c.env.GITHUB_CLIENT_ID,
        c.env.GITHUB_CLIENT_SECRET,
        c.env.GOOGLE_CLIENT_ID,
        c.env.GOOGLE_CLIENT_SECRET,
        c.env.GOOGLE_REDIRECT_URI,
      );
    },
  })(c, next);
});

app.use('*', prettyJSON());

app.route('/oauth/', OAuth);

export default app;