Basic websocket and auth for newt

server/auth/newt.ts

@@ -0,0 +1,80 @@
+export * from "./verifySession";
+export * from "./unauthorizedResponse";
+
+import {
+    encodeHexLowerCase,
+} from "@oslojs/encoding";
+import { sha256 } from "@oslojs/crypto/sha2";
+import { Newt, newts, newtSessions, NewtSession } from "@server/db/schema";
+import db from "@server/db";
+import { eq } from "drizzle-orm";
+import config from "@server/config";
+
+export const SESSION_COOKIE_NAME = "session";
+export const SESSION_COOKIE_EXPIRES = 1000 * 60 * 60 * 24 * 30;
+export const SECURE_COOKIES = config.server.secure_cookies;
+export const COOKIE_DOMAIN =
+    "." + new URL(config.app.base_url).hostname.split(".").slice(-2).join(".");
+
+export async function createNewtSession(
+    token: string,
+    newtId: string,
+): Promise<NewtSession> {
+    const sessionId = encodeHexLowerCase(
+        sha256(new TextEncoder().encode(token)),
+    );
+    const session: NewtSession = {
+        sessionId: sessionId,
+        newtId,
+        expiresAt: new Date(Date.now() + SESSION_COOKIE_EXPIRES).getTime(),
+    };
+    await db.insert(newtSessions).values(session);
+    return session;
+}
+
+export async function validateNewtSessionToken(
+    token: string,
+): Promise<SessionValidationResult> {
+    const sessionId = encodeHexLowerCase(
+        sha256(new TextEncoder().encode(token)),
+    );
+    const result = await db
+        .select({ newt: newts, session: newtSessions })
+        .from(newtSessions)
+        .innerJoin(newts, eq(newtSessions.newtId, newts.newtId))
+        .where(eq(newtSessions.sessionId, sessionId));
+    if (result.length < 1) {
+        return { session: null, newt: null };
+    }
+    const { newt, session } = result[0];
+    if (Date.now() >= session.expiresAt) {
+        await db
+            .delete(newtSessions)
+            .where(eq(newtSessions.sessionId, session.sessionId));
+        return { session: null, newt: null };
+    }
+    if (Date.now() >= session.expiresAt - (SESSION_COOKIE_EXPIRES / 2)) {
+        session.expiresAt = new Date(
+            Date.now() + SESSION_COOKIE_EXPIRES,
+        ).getTime();
+        await db
+            .update(newtSessions)
+            .set({
+                expiresAt: session.expiresAt,
+            })
+            .where(eq(newtSessions.sessionId, session.sessionId));
+    }
+    return { session, newt };
+}
+
+export async function invalidateNewtSession(sessionId: string): Promise<void> {
+    await db.delete(newtSessions).where(eq(newtSessions.sessionId, sessionId));
+}
+
+export async function invalidateAllNewtSessions(newtId: string): Promise<void> {
+    await db.delete(newtSessions).where(eq(newtSessions.newtId, newtId));
+}
+
+export type SessionValidationResult =
+    | { session: NewtSession; newt: Newt }
+    | { session: null; newt: null };

server/db/schema.ts

@@ -73,6 +73,12 @@ export const users = sqliteTable("user", {
     dateCreated: text("dateCreated").notNull(),
 });
 
+export const newts = sqliteTable("newt", {
+    newtId: text("id").primaryKey(),
+    secretHash: text("secretHash").notNull(),
+    dateCreated: text("dateCreated").notNull(),
+});
+
 export const twoFactorBackupCodes = sqliteTable("twoFactorBackupCodes", {
     codeId: integer("id").primaryKey({ autoIncrement: true }),
     userId: text("userId")
@@ -89,6 +95,14 @@ export const sessions = sqliteTable("session", {
     expiresAt: integer("expiresAt").notNull(),
 });
 
+export const newtSessions = sqliteTable("newtSession", {
+    sessionId: text("id").primaryKey(),
+    newtId: text("newtId")
+        .notNull()
+        .references(() => newts.newtId, { onDelete: "cascade" }),
+    expiresAt: integer("expiresAt").notNull(),
+});
+
 export const userOrgs = sqliteTable("userOrgs", {
     userId: text("userId")
         .notNull()
@@ -227,6 +241,8 @@ export type Resource = InferSelectModel<typeof resources>;
 export type ExitNode = InferSelectModel<typeof exitNodes>;
 export type Target = InferSelectModel<typeof targets>;
 export type Session = InferSelectModel<typeof sessions>;
+export type Newt = InferSelectModel<typeof newts>;
+export type NewtSession = InferSelectModel<typeof newtSessions>;
 export type EmailVerificationCode = InferSelectModel<
     typeof emailVerificationCodes
 >;

server/routers/auth/index.ts

@@ -12,6 +12,7 @@ export * from "./verifyTargetAccess";
 export * from "./verifyRoleAccess";
 export * from "./verifyUserAccess";
 export * from "./verifyAdmin";
+// export * from "./verifySuperUser";
 export * from "./verifyEmail";
 export * from "./requestEmailVerificationCode";
 export * from "./changePassword";

server/routers/auth/newtGetToken.ts

@@ -0,0 +1,115 @@
+import { verify } from "@node-rs/argon2";
+import {
+    createSession,
+    generateSessionToken,
+    verifySession,
+} from "@server/auth";
+import db from "@server/db";
+import { newts } from "@server/db/schema";
+import HttpCode from "@server/types/HttpCode";
+import response from "@server/utils/response";
+import { eq } from "drizzle-orm";
+import { NextFunction, Request, Response } from "express";
+import createHttpError from "http-errors";
+import { z } from "zod";
+import { fromError } from "zod-validation-error";
+import config from "@server/config";
+import { validateNewtSessionToken } from "@server/auth/newt";
+
+export const newtGetTokenBodySchema = z.object({
+    newtId: z.string().email(),
+    secret: z.string(),
+    token: z.string().optional(),
+});
+
+export type NewtGetTokenBody = z.infer<typeof newtGetTokenBodySchema>;
+
+export async function newtGetToken(
+    req: Request,
+    res: Response,
+    next: NextFunction
+): Promise<any> {
+    const parsedBody = newtGetTokenBodySchema.safeParse(req.body);
+
+    if (!parsedBody.success) {
+        return next(
+            createHttpError(
+                HttpCode.BAD_REQUEST,
+                fromError(parsedBody.error).toString()
+            )
+        );
+    }
+
+    const { newtId, secret, token } = parsedBody.data;
+
+    try {
+        if (token) {
+            const { session, newt } = await validateNewtSessionToken(
+                token
+            );
+            if (session) {
+                return response<null>(res, {
+                    data: null,
+                    success: true,
+                    error: false,
+                    message: "Token session already valid",
+                    status: HttpCode.OK,
+                });
+            }
+        }
+
+        const existingNewtRes = await db
+            .select()
+            .from(newts)
+            .where(eq(newts.newtId, newtId));
+        if (!existingNewtRes || !existingNewtRes.length) {
+            return next(
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    "No newt found with that newtId"
+                )
+            );
+        }
+
+        const existingNewt = existingNewtRes[0];
+
+        const validSecret = await verify(
+            existingNewt.secretHash,
+            secret,
+            {
+                memoryCost: 19456,
+                timeCost: 2,
+                outputLen: 32,
+                parallelism: 1,
+            }
+        );
+        if (!validSecret) {
+            return next(
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    "Secret is incorrect"
+                )
+            );
+        }
+
+        const resToken = generateSessionToken();
+        await createSession(resToken, existingNewt.newtId);
+
+        return response<{ token: string }>(res, {
+            data: {
+                token: resToken
+            },
+            success: true,
+            error: false,
+            message: "Token created successfully",
+            status: HttpCode.OK,
+        });
+    } catch (e) {
+        return next(
+            createHttpError(
+                HttpCode.INTERNAL_SERVER_ERROR,
+                "Failed to authenticate newt"
+            )
+        );
+    }
+}

server/routers/auth/sendEmailVerificationCode.ts

@@ -4,8 +4,8 @@ import db from "@server/db";
 import { users, emailVerificationCodes } from "@server/db/schema";
 import { eq } from "drizzle-orm";
 import { sendEmail } from "@server/emails";
-import VerifyEmail from "@server/emails/templates/VerifyEmailCode";
 import config from "@server/config";
+import VerifyEmail from "@server/emails/templates/verifyEmailCode";
 
 export async function sendEmailVerificationCode(
     email: string,

server/routers/external.ts

@@ -19,7 +19,6 @@ import {
     verifyResourceAccess,
     verifyTargetAccess,
     verifyRoleAccess,
-    verifyUserInRole,
     verifyUserAccess,
 } from "./auth";
 import { verifyUserHasAction } from "./auth/verifyUserHasAction";

server/routers/ws.ts

@@ -3,6 +3,11 @@ import { Server as HttpServer } from 'http';
 import { WebSocket, WebSocketServer } from 'ws';
 import { IncomingMessage } from 'http';
 import { Socket } from 'net';
+import { Newt, newts, NewtSession } from '@server/db/schema';
+import { eq } from 'drizzle-orm';
+import db from '@server/db';
+import { newtGetToken } from './auth';
+import { validateNewtSessionToken } from '@server/auth/newt';
 
 // Custom interfaces
 interface WebSocketRequest extends IncomingMessage {
@@ -10,13 +15,13 @@ interface WebSocketRequest extends IncomingMessage {
 }
 
 interface AuthenticatedWebSocket extends WebSocket {
-  userId?: string;
+    newt?: Newt;
     isAlive?: boolean;
 }
 
 interface TokenPayload {
-  userId: string;
-  // Add other token payload properties as needed
+    newt: Newt;
+    session: NewtSession;
 }
 
 const router: Router = Router();
@@ -25,10 +30,25 @@ const wss: WebSocketServer = new WebSocketServer({ noServer: true });
 // Token verification middleware
 const verifyToken = async (token: string): Promise<TokenPayload | null> => {
     try {
-    // This is where you'd implement your token verification logic
-    // For example, verify JWT, check against database, etc.
-    // Return the token payload if valid, null if invalid
-    return { userId: 'dummy-user-id' }; // Placeholder return
+
+        const { session, newt } = await validateNewtSessionToken(
+            token
+        );
+
+        if (!session || !newt) {
+            return null;
+        }
+
+        const existingNewt = await db
+            .select()
+            .from(newts)
+            .where(eq(newts.newtId, newt.newtId));
+
+        if (!existingNewt || !existingNewt[0]) {
+            return null;
+        }
+
+        return { newt: existingNewt[0], session };
     } catch (error) {
         console.error('Token verification failed:', error);
         return null;
@@ -41,6 +61,8 @@ router.get('/ws', (req: Request, res: Response) => {
     res.status(200).send('WebSocket endpoint');
 });
 
+router.get('/ws/auth/newtGetToken', newtGetToken);
+
 // Set up WebSocket server handling
 const handleWSUpgrade = (server: HttpServer): void => {
     server.on('upgrade', async (request: WebSocketRequest, socket: Socket, head: Buffer) => {
@@ -68,8 +90,8 @@ const handleWSUpgrade = (server: HttpServer): void => {
             request.token = token;
 
             wss.handleUpgrade(request, socket, head, (ws: AuthenticatedWebSocket) => {
-        // Attach user data to the WebSocket instance
-        ws.userId = tokenPayload.userId;
+                // Attach newt data to the WebSocket instance
+                ws.newt = tokenPayload.newt;
                 ws.isAlive = true;
                 wss.emit('connection', ws, request);
             });
@@ -89,7 +111,7 @@ interface WSMessage {
 
 // WebSocket connection handler
 wss.on('connection', (ws: AuthenticatedWebSocket, request: WebSocketRequest) => {
-  console.log(`Client connected - User ID: ${ws.userId}`);
+    console.log(`Client connected - Newt ID: ${ws.newt?.newtId}`);
 
     // Set up ping-pong for connection health check
     const pingInterval = setInterval(() => {
@@ -129,7 +151,7 @@ wss.on('connection', (ws: AuthenticatedWebSocket, request: WebSocketRequest) =>
     // Handle client disconnect
     ws.on('close', () => {
         clearInterval(pingInterval);
-    console.log(`Client disconnected - User ID: ${ws.userId}`);
+        console.log(`Client disconnected - Newt ID: ${ws.newt?.newtId}`);
     });
 
     // Handle errors