flxn-app/src/utils/supertokens.ts

import {
  createMiddleware,
  createServerFn,
  createServerOnlyFn,
} from "@tanstack/react-start";
import { getRequest, setResponseHeader } from "@tanstack/react-start/server";
import { redirect as redirect } from "@tanstack/react-router";
import UserRoles from "supertokens-node/recipe/userroles";
import UserMetadata from "supertokens-node/recipe/usermetadata";
import { getSessionForStart } from "@/lib/supertokens/recipes/start-session";
import { Logger } from "@/lib/logger";
import z from "zod";
import { serverFnLoggingMiddleware } from "./activities";
import { pbAdmin } from "@/lib/pocketbase/client";
const logger = new Logger("Middleware");

const verifySuperTokensSession = async (
  request: Request
) => {
  let session = await getSessionForStart(request, { sessionRequired: false });

  if (session?.needsRefresh) {
    logger.info("Session needs refresh - redirecting to client");
    return { context: { session: { tryRefresh: true } } };
  }
  
  const userAuthId = session?.userId;

  if (!userAuthId || !session) {
    logger.info("No authenticated user");
    return { context: { userAuthId: null, roles: [] } };
  }

  const { roles } = await fetchUserRoles(userAuthId);
  const { metadata } = await UserMetadata.getUserMetadata(userAuthId);

  return {
    context: {
      userAuthId,
      roles,
      metadata,
      phone: session.phone,
      session: {
        accessTokenPayload: session.accessTokenPayload,
        sessionHandle: session.sessionHandle,
      },
    },
  };
};

export const getSessionContext = createServerOnlyFn(async (request: Request, options?: { isServerFunction?: boolean }) => {
  const session = await verifySuperTokensSession(request);

  if (session.context.session?.tryRefresh) {
    if (options?.isServerFunction) {
      throw new Error("SESSION_REFRESH_REQUIRED");
    }

    const url = new URL(request.url);

    if (url.pathname === '/refresh-session') {
      logger.warn("Already on refresh-session page but session needs refresh - treating as unauthenticated");
      throw new Error("Unauthenticated");
    }

    const from = encodeURIComponent(url.pathname + url.search);
    throw redirect({
      to: "/refresh-session",
      search: { redirect: from }
    });
  }

  if (!session.context.userAuthId) {
    logger.error("Unauthenticated user", session.context);
    throw new Error("Unauthenticated");
  }

  const context = {
    userAuthId: session.context.userAuthId,
    roles: session.context.roles,
    metadata: session.context.metadata,
    phone: session.context.phone
  };

  try {
    const player = await pbAdmin.getPlayerByAuthId(session.context.userAuthId);
    if (player) {
      await pbAdmin.updatePlayer(player.id, {
        last_activity: new Date().toISOString(),
      });
    }
  } catch (error) {
    logger.error("Failed to update player last_activity", error);
  }

  return context;
});

export const superTokensRequestMiddleware = createMiddleware({
  type: "request",
}).server(async ({ next, request }) => {
  const context = await getSessionContext(request);
  return next({ context });
});

export const superTokensFunctionMiddleware = createMiddleware({
  type: "function",
}).server(async ({ next }) => {
  const request = getRequest();

  try {
    const context = await getSessionContext(request, { isServerFunction: true });
    return next({ context });
  } catch (error: any) {
    if (error.message === "SESSION_REFRESH_REQUIRED") {
      throw new Response(
        JSON.stringify({
          error: "SESSION_REFRESH_REQUIRED",
          message: "Session needs to be refreshed",
          shouldRetry: true
        }),
        {
          status: 440,
          headers: {
            "Content-Type": "application/json",
            "X-Session-Expired": "true"
          }
        }
      );
    }
    throw error;
  }
});

export const superTokensAdminFunctionMiddleware = createMiddleware({
  type: "function",
}).server(async ({ next }) => {
  const request = getRequest();

  try {
    const context = await getSessionContext(request, { isServerFunction: true });

    if (context.roles?.includes("Admin")) {
      return next({ context });
    }

    logger.error("Unauthorized user in admin function.", context);
    throw new Error("Unauthorized");
  } catch (error: any) {
    if (error.message === "SESSION_REFRESH_REQUIRED") {
      throw new Response(
        JSON.stringify({
          error: "SESSION_REFRESH_REQUIRED",
          message: "Session needs to be refreshed",
          shouldRetry: true
        }),
        {
          status: 440,
          headers: {
            "Content-Type": "application/json",
            "X-Session-Expired": "true"
          }
        }
      );
    }
    throw error;
  }
});

export const fetchUserRoles = async (userAuthId: string) => {
  const response = await UserRoles.getRolesForUser("public", userAuthId);
  return response;
};

export const setUserMetadata = createServerFn({ method: "POST" })
  .inputValidator(
    z
      .object({
        first_name: z
          .string()
          .min(2)
          .max(20)
          .regex(
            /^[a-zA-Z0-9\s]+$/,
            "First name must be 2-20 characters long and contain only letters and spaces"
          ),
        last_name: z
          .string()
          .min(2)
          .max(20)
          .regex(
            /^[a-zA-Z0-9\s]+$/,
            "Last name must be 2-20 characters long and contain only letters and spaces"
          ),
        player_id: z.string(),
      })
      .partial()
  )
  .middleware([superTokensFunctionMiddleware])
  .handler(async ({ context, data }) => {
    const { userAuthId, metadata } = context;
    if (!userAuthId) return;

    await UserMetadata.updateUserMetadata(userAuthId, {
      first_name: data.first_name,
      last_name: data.last_name,
      player_id: data.player_id,
    });

    return {
      metadata: {
        ...metadata,
        ...data,
      },
    };
  });

export const updateUserColorScheme = createServerFn({ method: "POST" })
  .inputValidator((data: string) => data)
  .middleware([superTokensFunctionMiddleware, serverFnLoggingMiddleware])
  .handler(async ({ context, data }) => {
    const { userAuthId, metadata } = context;
    if (!userAuthId) return;

    await UserMetadata.updateUserMetadata(userAuthId, {
      colorScheme: data,
    });

    return {
      metadata: {
        ...metadata,
        colorScheme: data,
      },
    };
  });

export const updateUserAccentColor = createServerFn({ method: "POST" })
  .inputValidator((data: string) => data)
  .middleware([superTokensFunctionMiddleware, serverFnLoggingMiddleware])
  .handler(async ({ context, data }) => {
    const { userAuthId, metadata } = context;
    if (!userAuthId) return;

    await UserMetadata.updateUserMetadata(userAuthId, {
      accentColor: data,
    });

    return {
      metadata: {
        ...metadata,
        accentColor: data,
      },
    };
  });