Newt registration?

bruno/Auth/login.bru

@@ -12,7 +12,7 @@ post {
 
 body:json {
   {
-    "email": "milo@fossorial.io",
+    "email": "owen@fossorial.io",
     "password": "Password123!"
   }
 }

bruno/Auth/signup.bru

@@ -12,7 +12,7 @@ put {
 
 body:json {
   {
-    "email": "milo@fossorial.io",
+    "email": "numbat@fossorial.io",
     "password": "Password123!"
   }
 }

bruno/Sites/Get Site.bru

@@ -5,7 +5,7 @@ meta {
 }
 
 get {
-  url: http://localhost:3000/api/v1/org/theorg/sites/mexican-mole-lizard-windy
+  url: http://localhost:3000/api/v1/org/test/sites/mexican-mole-lizard-windy
   body: none
   auth: none
 }

server/db/schema.ts

@@ -17,7 +17,7 @@ export const sites = sqliteTable("sites", {
         onDelete: "set null",
     }),
     name: text("name").notNull(),
-    pubKey: text("pubKey").notNull(),
+    pubKey: text("pubKey"),
     subnet: text("subnet").notNull(),
     megabytesIn: integer("bytesIn"),
     megabytesOut: integer("bytesOut"),
@@ -77,6 +77,9 @@ export const newts = sqliteTable("newt", {
     newtId: text("id").primaryKey(),
     secretHash: text("secretHash").notNull(),
     dateCreated: text("dateCreated").notNull(),
+    siteId: integer("siteId").references(() => sites.siteId, {
+        onDelete: "cascade",
+    }),
 });
 
 export const twoFactorBackupCodes = sqliteTable("twoFactorBackupCodes", {

server/routers/messageHandlers.ts

@@ -1,6 +1,6 @@
-import { handleNewtMessage } from "./newt";
+import { handleRegisterMessage } from "./newt";
 import { MessageHandler } from "./ws";
 
 export const messageHandlers: Record<string, MessageHandler> = {
-    "newt": handleNewtMessage,
+    "newt/wg/register": handleRegisterMessage,
 };

server/routers/newt/handleNewtMessage.ts

@@ -1,22 +0,0 @@
-// messageHandlers/chat.ts
-import { MessageHandler } from "../ws";
-
-export const handleNewtMessage: MessageHandler = async (context) => {
-    const { message, senderNewtId, sendToClient } = context;
-    
-    // Process chat message
-    // ... your chat logic here ...
-
-    // Example response
-    return {
-        message: {
-            type: 'newt_response',
-            data: {
-                originalMessage: message.data,
-                timestamp: new Date().toISOString()
-            }
-        },
-        broadcast: false,  // Send to all clients
-        excludeSender: false  // Include sender in broadcast
-    };
-};

server/routers/newt/handleRegisterMessage.ts

@@ -0,0 +1,42 @@
+import db from "@server/db";
+import { MessageHandler } from "../ws";
+import { sites } from "@server/db/schema";
+import { eq } from "drizzle-orm";
+
+export const handleRegisterMessage: MessageHandler = async (context) => {
+    const { message, newt, sendToClient } = context;
+    
+    if (!newt) {
+        console.log("Newt not found");
+        return;
+    }
+
+    if (!newt.siteId) {
+        console.log("Newt has no site!"); // TODO: Maybe we create the site here?
+        return;
+    }
+    
+    const siteId = newt.siteId;
+        
+    // get the site
+    const site = await db
+    .select()
+    .from(sites)
+    .where(eq(sites.siteId, siteId))
+    .limit(1);
+
+
+        const { publicKey } = message.data;
+
+    return {
+        message: {
+            type: 'newt/wg/connect',
+            data: {
+                publicKey: 'publicKey',
+
+            }
+        },
+        broadcast: false,  // Send to all clients
+        excludeSender: false  // Include sender in broadcast
+    };
+};

server/routers/newt/index.ts

@@ -1,3 +1,3 @@
 export * from "./createNewt";
 export * from "./getToken";
-export * from "./handleNewtMessage";
+export * from "./handleRegisterMessage";

server/routers/site/createSite.ts

@@ -19,7 +19,7 @@ const createSiteSchema = z.object({
     name: z.string().min(1).max(255),
     exitNodeId: z.number().int().positive(),
     subdomain: z.string().min(1).max(255).optional(),
-    pubKey: z.string(),
+    pubKey: z.string().optional(),
     subnet: z.string(),
 });
 
@@ -68,16 +68,24 @@ export async function createSite(
 
         const niceId = await getUniqueSiteName(orgId);
 
+        let payload: any = {
+            orgId,
+            exitNodeId,
+            name,
+            niceId,
+            subnet,
+        };
+
+        if (pubKey) {
+            payload = {
+                ...payload,
+                pubKey,
+            };
+        }
+
         const [newSite] = await db
             .insert(sites)
-            .values({
-                orgId,
-                exitNodeId,
-                name,
-                niceId,
-                pubKey,
-                subnet,
-            })
+            .values(payload)
             .returning();
 
         const adminRole = await db
@@ -105,11 +113,13 @@ export async function createSite(
             });
         }
 
-        // add the peer to the exit node
-        await addPeer(exitNodeId, {
-            publicKey: pubKey,
-            allowedIps: [],
-        });
+        if (pubKey) {
+            // add the peer to the exit node
+            await addPeer(exitNodeId, {
+                publicKey: pubKey,
+                allowedIps: [],
+            });
+        }
 
         return response(res, {
             data: {

server/routers/site/deleteSite.ts

@@ -48,7 +48,9 @@ export async function deleteSite(
             );
         }
 
-        await deletePeer(deletedSite.exitNodeId!, deletedSite.pubKey);
+        if (deletedSite.pubKey) {
+            await deletePeer(deletedSite.exitNodeId!, deletedSite.pubKey);
+        }
 
         return response(res, {
             data: null,
@@ -79,7 +81,7 @@ async function removePeer(publicKey: string) {
         }
 
         const data = await response.json();
-        console.log("Peer removed successfully:", data.status);
+        logger.info("Peer removed successfully:", data.status);
         return data;
     } catch (error: any) {
         console.error("Error removing peer:", error.message);

server/routers/ws.ts

@@ -8,6 +8,7 @@ import { eq } from "drizzle-orm";
 import db from "@server/db";
 import { validateNewtSessionToken } from "@server/auth/newt";
 import { messageHandlers } from "./messageHandlers";
+import logger from "@server/logger";
 
 // Custom interfaces
 interface WebSocketRequest extends IncomingMessage {
@@ -39,7 +40,7 @@ interface HandlerResponse {
 interface HandlerContext {
     message: WSMessage;
     senderWs: WebSocket;
-    senderNewtId: string;
+    newt: Newt | undefined;
     sendToClient: (newtId: string, message: WSMessage) => boolean;
     broadcastToAllExcept: (message: WSMessage, excludeNewtId?: string) => void;
     connectedClients: Map<string, WebSocket[]>;
@@ -58,7 +59,7 @@ const addClient = (newtId: string, ws: AuthenticatedWebSocket): void => {
     const existingClients = connectedClients.get(newtId) || [];
     existingClients.push(ws);
     connectedClients.set(newtId, existingClients);
-    console.log(`Client added to tracking - Newt ID: ${newtId}, Total connections: ${existingClients.length}`);
+    logger.info(`Client added to tracking - Newt ID: ${newtId}, Total connections: ${existingClients.length}`);
 };
 
 const removeClient = (newtId: string, ws: AuthenticatedWebSocket): void => {
@@ -67,10 +68,10 @@ const removeClient = (newtId: string, ws: AuthenticatedWebSocket): void => {
     
     if (updatedClients.length === 0) {
         connectedClients.delete(newtId);
-        console.log(`All connections removed for Newt ID: ${newtId}`);
+        logger.info(`All connections removed for Newt ID: ${newtId}`);
     } else {
         connectedClients.set(newtId, updatedClients);
-        console.log(`Connection removed - Newt ID: ${newtId}, Remaining connections: ${updatedClients.length}`);
+        logger.info(`Connection removed - Newt ID: ${newtId}, Remaining connections: ${updatedClients.length}`);
     }
 };
 
@@ -78,7 +79,7 @@ const removeClient = (newtId: string, ws: AuthenticatedWebSocket): void => {
 const sendToClient = (newtId: string, message: WSMessage): boolean => {
     const clients = connectedClients.get(newtId);
     if (!clients || clients.length === 0) {
-        console.log(`No active connections found for Newt ID: ${newtId}`);
+        logger.info(`No active connections found for Newt ID: ${newtId}`);
         return false;
     }
 
@@ -198,7 +199,7 @@ wss.on("connection", (ws: AuthenticatedWebSocket, request: WebSocketRequest) =>
     ws.on("message", async (data) => {
         try {
             const message: WSMessage = JSON.parse(data.toString());
-            // console.log(`Message received from Newt ID ${newtId}:`, message);
+            // logger.info(`Message received from Newt ID ${newtId}:`, message);
     
             // Validate message format
             if (!message.type || typeof message.type !== "string") {
@@ -215,7 +216,7 @@ wss.on("connection", (ws: AuthenticatedWebSocket, request: WebSocketRequest) =>
             const response = await handler({
                 message,
                 senderWs: ws,
-                senderNewtId: newtId,
+                newt: ws.newt,
                 sendToClient,
                 broadcastToAllExcept,
                 connectedClients
@@ -250,7 +251,7 @@ wss.on("connection", (ws: AuthenticatedWebSocket, request: WebSocketRequest) =>
     ws.on("close", () => {
         clearInterval(pingInterval);
         removeClient(newtId, ws);
-        console.log(`Client disconnected - Newt ID: ${newtId}`);
+        logger.info(`Client disconnected - Newt ID: ${newtId}`);
     });
 
     ws.on("error", (error: Error) => {