headplane/app/server/web/sessions.ts

import { eq } from "drizzle-orm";
import { LibSQLDatabase } from "drizzle-orm/libsql/driver";
import { EncryptJWT, jwtDecrypt } from "jose";
import { createHash } from "node:crypto";
import { open, readFile, rm } from "node:fs/promises";
import { resolve } from "node:path";
import { createCookie } from "react-router";
import { ulid } from "ulidx";

import log from "~/utils/log";

import { users } from "../db/schema";
import { Capabilities, Roles } from "./roles";

export interface AuthSession {
  state: "auth";
  api_key: string;
  user: {
    subject: string;
    name: string;
    email?: string;
    username?: string;
    picture?: string;
  };
}

interface JWTSession {
  api_key: string;
  user: {
    subject: string;
    name: string;
    email?: string;
    username?: string;
    picture?: string;
  };
}

export interface OidcFlowSession {
  state: "flow";
  oidc: {
    state: string;
    nonce: string;
    code_verifier: string;
    redirect_uri: string;
  };
}

interface AuthSessionOptions {
  secret: string;
  db: LibSQLDatabase;
  oidcUsersFile?: string;
  cookie: {
    name: string;
    secure: boolean;
    maxAge: number;
    domain?: string;
  };
}

class Sessionizer {
  private options: AuthSessionOptions;

  constructor(options: AuthSessionOptions) {
    this.options = options;
  }

  // This throws on the assumption that auth is already checked correctly
  // on something that wraps the route calling auth. The top-level routes
  // that call this are wrapped with try/catch to handle the error.
  async auth(request: Request) {
    return decodeSession(request, this.options);
  }

  async createSession(payload: JWTSession, maxAge = this.options.cookie.maxAge) {
    // TODO: What the hell is this garbage
    return createSession(payload, {
      ...this.options,
      cookie: {
        ...this.options.cookie,
        maxAge,
      },
    });
  }

  async destroySession() {
    return destroySession(this.options);
  }

  async roleForSubject(subject: string): Promise<keyof typeof Roles | undefined> {
    const [user] = await this.options.db
      .select()
      .from(users)
      .where(eq(users.sub, subject))
      .limit(1);

    if (!user) {
      return;
    }

    // We need this in string form based on Object.keys of the roles
    for (const [key, value] of Object.entries(Roles)) {
      if (value === user.caps) {
        return key as keyof typeof Roles;
      }
    }
  }

  // Given an OR of capabilities, check if the session has the required
  // capabilities. If not, return false. Can throw since it calls auth()
  async check(request: Request, capabilities: Capabilities) {
    const session = await this.auth(request);

    // This is the subject we set on API key based sessions. API keys
    // inherently imply admin access so we return true for all checks.
    if (session.user.subject === "unknown-non-oauth") {
      return true;
    }

    const [user] = await this.options.db
      .select()
      .from(users)
      .where(eq(users.sub, session.user.subject))
      .limit(1);

    if (!user) {
      return false;
    }

    return (capabilities & user.caps) === capabilities;
  }

  // Updates the capabilities and roles of a subject
  // Creates the user record if it doesn't exist yet
  async reassignSubject(subject: string, role: keyof typeof Roles) {
    // Check if we are owner
    const subjectRole = await this.roleForSubject(subject);
    if (subjectRole === "owner") {
      return false;
    }

    // Use upsert to handle users who exist in Headscale but haven't
    // logged into Headplane yet (no DB record)
    await this.options.db
      .insert(users)
      .values({
        id: ulid(),
        sub: subject,
        caps: Roles[role],
        onboarded: false,
      })
      .onConflictDoUpdate({
        target: users.sub,
        set: { caps: Roles[role] },
      });

    return true;
  }
}

async function createSession(payload: JWTSession, options: AuthSessionOptions) {
  const now = Math.floor(Date.now() / 1000);
  const secret = createHash("sha256").update(options.secret, "utf8").digest();
  const jwt = await new EncryptJWT({
    ...payload,
  })
    .setProtectedHeader({ alg: "dir", enc: "A256GCM", typ: "JWT" })
    .setIssuedAt()
    .setExpirationTime(now + options.cookie.maxAge)
    .setIssuer("urn:tale:headplane")
    .setAudience("urn:tale:headplane")
    .setJti(ulid())
    .encrypt(secret);

  const cookie = createCookie(options.cookie.name, {
    ...options.cookie,
    path: __PREFIX__,
  });

  return cookie.serialize(jwt);
}

async function decodeSession(request: Request, options: AuthSessionOptions) {
  const cookieHeader = request.headers.get("cookie");
  if (cookieHeader === null) {
    throw new Error("No session cookie found");
  }

  const cookie = createCookie(options.cookie.name, {
    ...options.cookie,
    path: __PREFIX__,
  });

  const cookieValue = (await cookie.parse(cookieHeader)) as string | null;
  if (cookieValue === null) {
    throw new Error("Session cookie is empty");
  }

  const secret = createHash("sha256").update(options.secret, "utf8").digest();
  const { payload } = await jwtDecrypt(cookieValue, secret, {
    issuer: "urn:tale:headplane",
    audience: "urn:tale:headplane",
  });

  // Safe since we encode the session directly into the JWT
  return payload as unknown as JWTSession;
}

async function destroySession(options: AuthSessionOptions) {
  const cookie = createCookie(options.cookie.name, {
    ...options.cookie,
    path: __PREFIX__,
  });

  return cookie.serialize("", {
    expires: new Date(0),
  });
}

export async function createSessionStorage(options: AuthSessionOptions) {
  if (options.oidcUsersFile) {
    await migrateUserDatabase(options.oidcUsersFile, options.db);
  }

  return new Sessionizer(options);
}

async function migrateUserDatabase(path: string, db: LibSQLDatabase) {
  const realPath = resolve(path);

  try {
    const handle = await open(realPath, "a+");
    await handle.close();
  } catch (error) {
    if (error != null && typeof error === "object" && "code" in error && error.code === "ENOENT") {
      log.debug("config", "No old user database file found at %s", realPath);
      return;
    }

    log.warn("config", "Failed to migrate old user database at %s", realPath);
    log.warn("config", "This is not an error, but existing users will not be migrated");
    log.warn("config", "Unable to open user database file: %s", String(error));
    log.debug("config", "Error details: %s", error);
    return;
  }

  log.info("config", "Found old user database file at %s", realPath);
  log.info("config", "Migrating user database to the new SQL database");

  let migratableUsers: {
    u: string;
    c: number;
    oo?: boolean;
  }[];

  try {
    const data = await readFile(realPath, "utf8");
    if (data.trim().length === 0) {
      log.info("config", "Old user database file is empty, nothing to migrate");
      log.info("config", "You SHOULD remove oidc.user_storage_file from your config!");
      await rm(realPath, { force: true });
      return;
    }

    const users = JSON.parse(data.trim()) as {
      u?: string;
      c?: number;
      oo?: boolean;
    }[];

    migratableUsers = users.filter((user) => user.u !== undefined && user.c !== undefined) as {
      u: string;
      c: number;
      oo?: boolean;
    }[];
  } catch (error) {
    log.warn("config", "Error reading old user database file: %s", error);
    log.warn("config", "Not migrating any users");
    return;
  }

  if (migratableUsers.length === 0) {
    log.info("config", "No users found in the old database to migrate");
    return;
  }

  log.info("config", "Migrating %d users from the old database", migratableUsers.length);

  const updated = await db
    .insert(users)
    .values(
      migratableUsers.map((user) => ({
        id: ulid(),
        sub: user.u,
        caps: user.c,
        onboarded: user.oo ?? false,
      })),
    )
    .onConflictDoNothing({
      target: users.sub,
    })
    .returning();

  log.info("config", "Migrated %d users successfully", updated.length);
  log.info("config", "Removed old user database file %s", realPath);
  await rm(realPath, { force: true });
}