SST code adapter auth

sst-auth.ts

// auth function
import { CodeAdapter } from "sst/auth/adapter";
import { auth } from "sst/aws/auth";
import { createSessionBuilder } from "sst/auth";

// This can be whatever you want
export const session = createSessionBuilder<{
  account: {
    accountId: string;
    email: string;
  };
}>();

export const handler = auth.authorizer({
  session,
  providers: {
    code: CodeAdapter({
      onCodeRequest: async (code, claims, req) => {
        const searchParams = new URLSearchParams(req.url);
        const redirectUri = searchParams.get("redirect_uri")?.replace(process.env.AUTH_FRONTEND_URL as string, "");

        console.log("Code request", code, claims, redirectUri);

        return new Response(code, {
          status: 302,
          headers: {
            Location:
              process.env.AUTH_FRONTEND_URL +
              "/verify?" +
              new URLSearchParams({
                email: claims.email,
                redirect: redirectUri ?? "/workspace",
              }).toString(),
          },
        });
      },
      onCodeInvalid: async (code, claims) => {
        return new Response("Code is invalid " + code, {
          status: 200,
          headers: { "Content-Type": "text/plain" },
        });
      },
    }),
  },
  callbacks: {
    auth: {
      async allowClient(input) {
        return true;
      },
      async success(ctx, input) {
        // do stuff with db here
        return ctx.session({
          type: "account",
          properties: {
            accountId: "123",
            email: input.claims.email,
          },
        });
      },
    },
  },
});

// login form
function Index() {
  const [isSubmitting, setIsSubmitting] = React.useState<boolean>(false);
  const search = Route.useSearch();

  const handleSubmit = (event: React.FormEvent<HTMLFormElement>) => {
    setIsSubmitting(true);
    event.preventDefault();
    (event.target as HTMLFormElement).submit();
  };

  return (
    <form
      className="my-4 w-full"
      method="get"
      action={import.meta.env.VITE_AUTH_URL + "code/authorize" + `?redirect=${search.redirect}`}
      onSubmit={handleSubmit}
    >
      <fieldset className="group flex flex-col gap-y-4" disabled={isSubmitting}>
        <input autoComplete="off" type="email" name="email" placeholder="Email" required />
        <input type="hidden" name="client_id" value="web" />
        <input type="hidden" name="redirect_uri" value={window.location.origin + search.redirect || "/"} />
        <input type="hidden" name="response_type" value="token" />
        <div className="flex flex-col gap-y-3">
          <button disabled={isSubmitting} type="submit">
            Continue
          </button>
        </div>
      </fieldset>
    </form>
  );
}

// verify code
function Verify() {
  const search = Route.useSearch();

  return (
    <div className="mx-auto">
      <OTPInput
        name="code"
        pattern={REGEXP_ONLY_DIGITS_AND_CHARS}
        className="focus-visible:ring-0"
        onComplete={(code) => {
          window.location.href =
            import.meta.env.VITE_AUTH_URL +
            "code/callback?" +
            new URLSearchParams({
              code,
            }).toString();
        }}
        maxLength={6}
        containerClassName="group flex items-center has-[:disabled]:opacity-30"
        render={({ slots }) => (
          <>
            <div className="flex">
              {slots.slice(0, 3).map((slot, idx) => (
                <Slot key={idx} {...slot} />
              ))}
            </div>

            <FakeDash />

            <div className="flex">
              {slots.slice(3).map((slot, idx) => (
                <Slot key={idx} {...slot} />
              ))}
            </div>
          </>
        )}
      />
    </div>
  );
}

// auth provider
import * as React from "react";

export type AuthContextType = Record<
  string,
  {
    session: Session;
  }
>;

export interface AuthContext {
  isAuthenticated: boolean;
  logout: () => Promise<void>;
  authentication: AuthContextType;
}

export interface Session {
  email: string;
  accountId: string;
  token: string;
}

const AuthContext = React.createContext<AuthContext | null>(null);

const accountIdKey = "accountId";
const authenticationKey = "authentication";

function getAccountId() {
  return localStorage.getItem(accountIdKey);
}

function setLocalAccountId(accountId: string | null) {
  if (accountId) {
    localStorage.setItem(accountIdKey, accountId);
  } else {
    localStorage.removeItem(accountIdKey);
  }
}

function getAuthentication() {
  return JSON.parse(localStorage.getItem(authenticationKey) ?? "{}") as AuthContextType | null;
}

function setLocalAuthentication(authentication: AuthContextType | null) {
  if (authentication) {
    localStorage.setItem(authenticationKey, JSON.stringify(authentication));
  } else {
    localStorage.removeItem(authenticationKey);
  }
}

export function AuthProvider({ children }: { children: React.ReactNode }) {
  const [authentication, setAuthentication] = React.useState<AuthContextType>(getAuthentication() ?? {});
  const [accountId, setAccountId] = React.useState<string | null>(getAccountId());
  const isAuthenticated = !!accountId;

  React.useEffect(() => {
    const fragment = new URLSearchParams(window.location.hash.substring(1));
    const accessToken = fragment.get("access_token");

    if (!accessToken) {
      return;
    }

    const [_, payloadEncoded] = accessToken.split(".");

    if (!payloadEncoded) {
      return;
    }

    const payload = JSON.parse(atob(payloadEncoded.replace(/-/g, "+").replace(/_/g, "/"))) as {
      properties: Omit<Session, "token">;
    };

    setLocalAccountId(payload.properties.accountId);
    setAccountId(payload.properties.accountId);

    setAuthentication((currentAuthentication) => {
      const newAuthentication = {
        ...currentAuthentication,
        [payload.properties.accountId]: {
          session: {
            ...payload.properties,
            token: accessToken,
          },
        },
      };

      setLocalAuthentication(newAuthentication);

      return newAuthentication;
    });

    history.replaceState(null, "", " ");
  }, []);

  const logout = React.useCallback(async () => {
    const databases = await window.indexedDB.databases();

    databases.forEach((db) => {
      window.indexedDB.deleteDatabase(db.name!);
    });

    localStorage.clear();
    setAccountId(null);
    setAuthentication({});
  }, []);

  return <AuthContext.Provider value={{ isAuthenticated, authentication, logout }}>{children}</AuthContext.Provider>;
}

export function useAuth() {
  const context = React.useContext(AuthContext);
  if (!context) {
    throw new Error("useAuth must be used within an AuthProvider");
  }
  return context;
}