Merge branch 'hp-multi-client' into auth-providers-clients

server/routers/gerbil/getAllRelays.ts

@@ -0,0 +1,160 @@
+import { Request, Response, NextFunction } from "express";
+import { z } from "zod";
+import { clients, exitNodes, newts, olms, Site, sites, clientSites } from "@server/db/schema";
+import { db } from "@server/db";
+import { eq } from "drizzle-orm";
+import HttpCode from "@server/types/HttpCode";
+import createHttpError from "http-errors";
+import logger from "@server/logger";
+import { fromError } from "zod-validation-error";
+
+// Define Zod schema for request validation
+const getAllRelaysSchema = z.object({
+    publicKey: z.string().optional(),
+});
+
+// Type for peer destination
+interface PeerDestination {
+    destinationIP: string;
+    destinationPort: number;
+}
+
+// Updated mappings type to support multiple destinations per endpoint
+interface ProxyMapping {
+    destinations: PeerDestination[];
+}
+
+export async function getAllRelays(
+    req: Request,
+    res: Response,
+    next: NextFunction
+): Promise<any> {
+    try {
+        // Validate request parameters
+        const parsedParams = getAllRelaysSchema.safeParse(req.body);
+        if (!parsedParams.success) {
+            return next(
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    fromError(parsedParams.error).toString()
+                )
+            );
+        }
+
+        const { publicKey } = parsedParams.data;
+
+        if (!publicKey) {
+            return next(createHttpError(HttpCode.BAD_REQUEST, 'publicKey is required'));
+        }
+
+        // Fetch exit node
+        let [exitNode] = await db.select().from(exitNodes).where(eq(exitNodes.publicKey, publicKey));
+        if (!exitNode) {
+            return next(createHttpError(HttpCode.NOT_FOUND, "Exit node not found"));
+        }
+
+        // Fetch sites for this exit node
+        const sitesRes = await db.select().from(sites).where(eq(sites.exitNodeId, exitNode.exitNodeId));
+
+        if (sitesRes.length === 0) {
+            return res.status(HttpCode.OK).send({
+                mappings: {}
+            });
+        }
+        
+        // Initialize mappings object for multi-peer support
+        let mappings: { [key: string]: ProxyMapping } = {};
+
+        // Process each site
+        for (const site of sitesRes) {
+            if (!site.endpoint || !site.subnet || !site.listenPort) {
+                continue;
+            }
+
+            // Find all clients associated with this site through clientSites
+            const clientSitesRes = await db
+                .select()
+                .from(clientSites)
+                .where(eq(clientSites.siteId, site.siteId));
+            
+            for (const clientSite of clientSitesRes) {
+                // Get client information
+                const [client] = await db
+                    .select()
+                    .from(clients)
+                    .where(eq(clients.clientId, clientSite.clientId));
+
+                if (!client || !client.endpoint) {
+                    continue;
+                }
+
+                // Add this site as a destination for the client
+                if (!mappings[client.endpoint]) {
+                    mappings[client.endpoint] = { destinations: [] };
+                }
+
+                // Add site as a destination for this client
+                const destination: PeerDestination = {
+                    destinationIP: site.subnet.split("/")[0],
+                    destinationPort: site.listenPort
+                };
+
+                // Check if this destination is already in the array to avoid duplicates
+                const isDuplicate = mappings[client.endpoint].destinations.some(
+                    dest => dest.destinationIP === destination.destinationIP && 
+                            dest.destinationPort === destination.destinationPort
+                );
+
+                if (!isDuplicate) {
+                    mappings[client.endpoint].destinations.push(destination);
+                }
+            }
+
+            // Also handle site-to-site communication (all sites in the same org)
+            if (site.orgId) {
+                const orgSites = await db
+                    .select()
+                    .from(sites)
+                    .where(eq(sites.orgId, site.orgId));
+                
+                for (const peer of orgSites) {
+                    // Skip self
+                    if (peer.siteId === site.siteId || !peer.endpoint || !peer.subnet || !peer.listenPort) {
+                        continue;
+                    }
+
+                    // Add peer site as a destination for this site
+                    if (!mappings[site.endpoint]) {
+                        mappings[site.endpoint] = { destinations: [] };
+                    }
+
+                    const destination: PeerDestination = {
+                        destinationIP: peer.subnet.split("/")[0],
+                        destinationPort: peer.listenPort
+                    };
+
+                    // Check for duplicates
+                    const isDuplicate = mappings[site.endpoint].destinations.some(
+                        dest => dest.destinationIP === destination.destinationIP && 
+                                dest.destinationPort === destination.destinationPort
+                    );
+
+                    if (!isDuplicate) {
+                        mappings[site.endpoint].destinations.push(destination);
+                    }
+                }
+            }
+        }
+
+        logger.debug(`Returning mappings for ${Object.keys(mappings).length} endpoints`);
+        return res.status(HttpCode.OK).send({ mappings });
+    } catch (error) {
+        logger.error(error);
+        return next(
+            createHttpError(
+                HttpCode.INTERNAL_SERVER_ERROR,
+                "An error occurred..."
+            )
+        );
+    }
+}

server/routers/gerbil/getConfig.ts

@@ -79,14 +79,12 @@ export async function getConfig(req: Request, res: Response, next: NextFunction)
         }
 
         // Fetch sites for this exit node
-        const sitesRes = await db.query.sites.findMany({
-            where: eq(sites.exitNodeId, exitNode[0].exitNodeId),
-        });
+        const sitesRes = await db.select().from(sites).where(eq(sites.exitNodeId, exitNode[0].exitNodeId));
 
         const peers = await Promise.all(sitesRes.map(async (site) => {
             return {
                 publicKey: site.pubKey,
-                allowedIps: await getAllowedIps(site.siteId)
+                allowedIps: await getAllowedIps(site.siteId) // put 0.0.0.0/0 for now
             };
         }));
 

server/routers/gerbil/index.ts

@@ -1,2 +1,4 @@
 export * from "./getConfig";
 export * from "./receiveBandwidth";
+export * from "./updateHolePunch";
+export * from "./getAllRelays";

server/routers/gerbil/receiveBandwidth.ts

@@ -30,12 +30,13 @@ export const receiveBandwidth = async (
                 const { publicKey, bytesIn, bytesOut } = peer;
 
                 // Find the site by public key
-                const site = await trx.query.sites.findFirst({
-                    where: eq(sites.pubKey, publicKey)
-                });
+                const [site] = await trx
+                    .select()
+                    .from(sites)
+                    .where(eq(sites.pubKey, publicKey))
+                    .limit(1);
 
                 if (!site) {
-                    logger.warn(`Site not found for public key: ${publicKey}`);
                     continue;
                 }
                 let online = site.online;

server/routers/gerbil/updateHolePunch.ts

@@ -0,0 +1,242 @@
+import { Request, Response, NextFunction } from "express";
+import { z } from "zod";
+import { clients, newts, olms, Site, sites, clientSites } from "@server/db/schema";
+import { db } from "@server/db";
+import { eq } from "drizzle-orm";
+import HttpCode from "@server/types/HttpCode";
+import createHttpError from "http-errors";
+import logger from "@server/logger";
+import { fromError } from "zod-validation-error";
+import { validateNewtSessionToken } from "@server/auth/sessions/newt";
+import { validateOlmSessionToken } from "@server/auth/sessions/olm";
+
+// Define Zod schema for request validation
+const updateHolePunchSchema = z.object({
+    olmId: z.string().optional(),
+    newtId: z.string().optional(),
+    token: z.string(),
+    ip: z.string(),
+    port: z.number(),
+    timestamp: z.number()
+});
+
+// New response type with multi-peer destination support
+interface PeerDestination {
+    destinationIP: string;
+    destinationPort: number;
+}
+
+export async function updateHolePunch(
+    req: Request,
+    res: Response,
+    next: NextFunction
+): Promise<any> {
+    try {
+        // Validate request parameters
+        const parsedParams = updateHolePunchSchema.safeParse(req.body);
+        if (!parsedParams.success) {
+            return next(
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    fromError(parsedParams.error).toString()
+                )
+            );
+        }
+
+        const { olmId, newtId, ip, port, timestamp, token } = parsedParams.data;
+
+        
+        let currentSiteId: number | undefined;
+        let destinations: PeerDestination[] = [];
+        
+        if (olmId) {
+            logger.debug(`Got hole punch with ip: ${ip}, port: ${port} for olmId: ${olmId}`);
+
+            const { session, olm: olmSession } =
+                await validateOlmSessionToken(token);
+            if (!session || !olmSession) {
+                return next(
+                    createHttpError(HttpCode.UNAUTHORIZED, "Unauthorized")
+                );
+            }
+
+            if (olmId !== olmSession.olmId) {
+                logger.warn(`Olm ID mismatch: ${olmId} !== ${olmSession.olmId}`);
+                return next(
+                    createHttpError(HttpCode.UNAUTHORIZED, "Unauthorized")
+                );
+            }
+
+            const [olm] = await db
+                .select()
+                .from(olms)
+                .where(eq(olms.olmId, olmId));
+
+            if (!olm || !olm.clientId) {
+                logger.warn(`Olm not found: ${olmId}`);
+                return next(
+                    createHttpError(HttpCode.NOT_FOUND, "Olm not found")
+                );
+            }
+
+            const [client] = await db
+                .update(clients)
+                .set({
+                    endpoint: `${ip}:${port}`,
+                    lastHolePunch: timestamp
+                })
+                .where(eq(clients.clientId, olm.clientId))
+                .returning();
+            
+            if (!client) {
+                logger.warn(`Client not found for olm: ${olmId}`);
+                return next(
+                    createHttpError(HttpCode.NOT_FOUND, "Client not found")
+                );
+            }
+
+            // Get all sites that this client is connected to
+            const clientSitePairs = await db
+                .select()
+                .from(clientSites)
+                .where(eq(clientSites.clientId, client.clientId));
+            
+            if (clientSitePairs.length === 0) {
+                logger.warn(`No sites found for client: ${client.clientId}`);
+                return next(
+                    createHttpError(HttpCode.NOT_FOUND, "No sites found for client")
+                );
+            }
+            
+            // Get all sites details
+            const siteIds = clientSitePairs.map(pair => pair.siteId);
+            
+            for (const siteId of siteIds) {
+                const [site] = await db
+                    .select()
+                    .from(sites)
+                    .where(eq(sites.siteId, siteId));
+                
+                if (site && site.subnet && site.listenPort) {
+                    destinations.push({
+                        destinationIP: site.subnet.split("/")[0],
+                        destinationPort: site.listenPort
+                    });
+                }
+            }
+
+        } else if (newtId) {
+            const { session, newt: newtSession } =
+                await validateNewtSessionToken(token);
+
+            if (!session || !newtSession) {
+                return next(
+                    createHttpError(HttpCode.UNAUTHORIZED, "Unauthorized")
+                );
+            }
+
+            if (newtId !== newtSession.newtId) {
+                logger.warn(`Newt ID mismatch: ${newtId} !== ${newtSession.newtId}`);
+                return next(
+                    createHttpError(HttpCode.UNAUTHORIZED, "Unauthorized")
+                );
+            }
+
+            const [newt] = await db
+                .select()
+                .from(newts)
+                .where(eq(newts.newtId, newtId));
+
+            if (!newt || !newt.siteId) {
+                logger.warn(`Newt not found: ${newtId}`);
+                return next(
+                    createHttpError(HttpCode.NOT_FOUND, "New not found")
+                );
+            }
+
+            currentSiteId = newt.siteId;
+
+            // Update the current site with the new endpoint
+            const [updatedSite] = await db
+                .update(sites)
+                .set({
+                    endpoint: `${ip}:${port}`,
+                    lastHolePunch: timestamp
+                })
+                .where(eq(sites.siteId, newt.siteId))
+                .returning();
+            
+            if (!updatedSite || !updatedSite.subnet) {
+                logger.warn(`Site not found: ${newt.siteId}`);
+                return next(
+                    createHttpError(HttpCode.NOT_FOUND, "Site not found")
+                );
+            }
+
+            // Find all clients that connect to this site
+            const sitesClientPairs = await db
+                .select()
+                .from(clientSites)
+                .where(eq(clientSites.siteId, newt.siteId));
+            
+            // Get client details for each client
+            for (const pair of sitesClientPairs) {
+                const [client] = await db
+                    .select()
+                    .from(clients)
+                    .where(eq(clients.clientId, pair.clientId));
+                
+                if (client && client.endpoint) {
+                    const [host, portStr] = client.endpoint.split(':');
+                    if (host && portStr) {
+                        destinations.push({
+                            destinationIP: host,
+                            destinationPort: parseInt(portStr, 10)
+                        });
+                    }
+                }
+            }
+            
+            // If this is a newt/site, also add other sites in the same org
+        //     if (updatedSite.orgId) {
+        //         const orgSites = await db
+        //             .select()
+        //             .from(sites)
+        //             .where(eq(sites.orgId, updatedSite.orgId));
+                
+        //         for (const site of orgSites) {
+        //             // Don't add the current site to the destinations
+        //             if (site.siteId !== currentSiteId && site.subnet && site.endpoint && site.listenPort) {
+        //                 const [host, portStr] = site.endpoint.split(':');
+        //                 if (host && portStr) {
+        //                     destinations.push({
+        //                         destinationIP: host,
+        //                         destinationPort: site.listenPort
+        //                     });
+        //                 }
+        //             }
+        //         }
+        //     }
+        }
+
+        // if (destinations.length === 0) {
+        //     logger.warn(
+        //         `No peer destinations found for olmId: ${olmId} or newtId: ${newtId}`
+        //     );
+        //     return next(createHttpError(HttpCode.NOT_FOUND, "No peer destinations found"));
+        // }
+
+        // Return the new multi-peer structure
+        return res.status(HttpCode.OK).send({
+            destinations: destinations
+        });
+    } catch (error) {
+        logger.error(error);
+        return next(
+            createHttpError(
+                HttpCode.INTERNAL_SERVER_ERROR,
+                "An error occurred..."
+            )
+        );
+    }
+}