Handle unrelay and relaying better

This commit is contained in:
Owen
2025-12-02 11:17:08 -05:00
parent a7e32d4013
commit 152fb47ca4
9 changed files with 125 additions and 81 deletions

View File

@@ -1011,76 +1011,18 @@ async function handleMessagesForClientSites(
continue;
}
// Add peer to newt
const isRelayed = true; // Default to relaying for new connections
newtJobs.push(
newtAddPeer(
site.siteId,
{
publicKey: client.pubKey,
allowedIps: [`${client.subnet.split("/")[0]}/32`],
endpoint: isRelayed ? "" : ""
},
newt.newtId
)
await holepunchSiteAdd(
// this will kick off the add peer process for the client
client.clientId,
{
siteId: site.siteId,
exitNode: {
publicKey: exitNode.publicKey,
endpoint: exitNode.endpoint
}
},
olmId
);
// Get all site resources for this site that the client has access to
const accessibleResources = await trx
.select()
.from(siteResources)
.innerJoin(
clientSiteResourcesAssociationsCache,
eq(
siteResources.siteResourceId,
clientSiteResourcesAssociationsCache.siteResourceId
)
)
.where(
and(
eq(siteResources.siteId, site.siteId),
eq(
clientSiteResourcesAssociationsCache.clientId,
client.clientId
)
)
);
try {
// Add peer to olm
olmJobs.push(
olmAddPeer(
client.clientId,
{
siteId: site.siteId,
endpoint:
isRelayed || !site.endpoint
? `${exitNode.endpoint}:21820`
: site.endpoint,
publicKey: site.publicKey,
serverIP: site.address || "",
serverPort: site.listenPort || 0,
remoteSubnets: generateRemoteSubnets(
accessibleResources.map(
({ siteResources }) => siteResources
)
)
},
olmId
)
);
} catch (error) {
// if the error includes not found then its just because the olm does not exist anymore or yet and its fine if we dont send
if (
error instanceof Error &&
error.message.includes("not found")
) {
logger.debug(
`Olm data not found for client ${client.clientId}, skipping removal`
);
} else {
throw error;
}
}
}
// Update exit node destinations

View File

@@ -237,7 +237,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
);
}
let endpoint = site.endpoint;
let relayEndpoint: string | undefined = undefined;
if (relay) {
const [exitNode] = await db
.select()
@@ -248,7 +248,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
logger.warn(`Exit node not found for site ${site.siteId}`);
continue;
}
endpoint = `${exitNode.endpoint}:${config.getRawConfig().gerbil.clients_start_port}`;
relayEndpoint = `${exitNode.endpoint}:${config.getRawConfig().gerbil.clients_start_port}`;
}
const allSiteResources = await db // only get the site resources that this client has access to
@@ -274,7 +274,8 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
// Add site configuration to the array
siteConfigurations.push({
siteId: site.siteId,
endpoint: endpoint,
relayEndpoint: relayEndpoint, // this can be undefined now if not relayed
endpoint: site.endpoint,
publicKey: site.publicKey,
serverIP: site.address,
serverPort: site.listenPort,

View File

@@ -85,12 +85,11 @@ export const handleOlmRelayMessage: MessageHandler = async (context) => {
return {
message: {
type: "olm/wg/peer/relay",
data: {
siteId: siteId,
endpoint: exitNode.endpoint,
publicKey: exitNode.publicKey
}
type: "olm/wg/peer/relay",
data: {
siteId: siteId,
relayEndpoint: exitNode.endpoint
}
},
broadcast: false,
excludeSender: false

View File

@@ -135,6 +135,8 @@ export const handleOlmServerPeerAddMessage: MessageHandler = async (
return;
}
// NOTE: here we are always starting direct to the peer and will relay later
await newtAddPeer(siteId, {
publicKey: client.pubKey,
allowedIps: [`${client.subnet.split("/")[0]}/32`], // we want to only allow from that client

View File

@@ -0,0 +1,96 @@
import { db, exitNodes, sites } from "@server/db";
import { MessageHandler } from "@server/routers/ws";
import { clients, clientSitesAssociationsCache, Olm } from "@server/db";
import { and, eq } from "drizzle-orm";
import { updatePeer as newtUpdatePeer } from "../newt/peers";
import logger from "@server/logger";
export const handleOlmUnRelayMessage: MessageHandler = async (context) => {
const { message, client: c, sendToClient } = context;
const olm = c as Olm;
logger.info("Handling unrelay olm message!");
if (!olm) {
logger.warn("Olm not found");
return;
}
if (!olm.clientId) {
logger.warn("Olm has no site!"); // TODO: Maybe we create the site here?
return;
}
const clientId = olm.clientId;
const [client] = await db
.select()
.from(clients)
.where(eq(clients.clientId, clientId))
.limit(1);
if (!client) {
logger.warn("Client not found");
return;
}
// make sure we hand endpoints for both the site and the client and the lastHolePunch is not too old
if (!client.pubKey) {
logger.warn("Client has no endpoint or listen port");
return;
}
const { siteId } = message.data;
// Get the site
const [site] = await db
.select()
.from(sites)
.where(eq(sites.siteId, siteId))
.limit(1);
if (!site) {
logger.warn("Site not found or has no exit node");
return;
}
const [clientSiteAssociation] = await db
.update(clientSitesAssociationsCache)
.set({
isRelayed: false
})
.where(
and(
eq(clientSitesAssociationsCache.clientId, olm.clientId),
eq(clientSitesAssociationsCache.siteId, siteId)
)
)
.returning();
if (!clientSiteAssociation) {
logger.warn("Client-Site association not found");
return;
}
if (!clientSiteAssociation.endpoint) {
logger.warn("Client-Site association has no endpoint, cannot unrelay");
return;
}
// update the peer on the exit node
await newtUpdatePeer(siteId, client.pubKey, {
endpoint: clientSiteAssociation.endpoint // this is the endpoint of the client to connect directly to the exit node
});
return {
message: {
type: "olm/wg/peer/unrelay",
data: {
siteId: siteId,
endpoint: site.endpoint
}
},
broadcast: false,
excludeSender: false
};
};

View File

@@ -7,4 +7,5 @@ export * from "./deleteUserOlm";
export * from "./listUserOlms";
export * from "./deleteUserOlm";
export * from "./getUserOlm";
export * from "./handleOlmServerPeerAddMessage";
export * from "./handleOlmServerPeerAddMessage";
export * from "./handleOlmUnRelayMessage";

View File

@@ -103,6 +103,7 @@ export async function updatePeer(
siteId: peer.siteId,
publicKey: peer.publicKey,
endpoint: peer.endpoint,
relayEndpoint: peer.serverIP,
serverIP: peer.serverIP,
serverPort: peer.serverPort,
remoteSubnets: peer.remoteSubnets

View File

@@ -193,7 +193,7 @@ export async function deleteOrg(
// Send termination messages outside of transaction to prevent blocking
for (const newtId of deletedNewtIds) {
const payload = {
type: `newt/terminate`,
type: `newt/wg/terminate`,
data: {}
};
// Don't await this to prevent blocking the response

View File

@@ -12,7 +12,8 @@ import {
handleOlmRelayMessage,
handleOlmPingMessage,
startOlmOfflineChecker,
handleOlmServerPeerAddMessage
handleOlmServerPeerAddMessage,
handleOlmUnRelayMessage
} from "../olm";
import { handleHealthcheckStatusMessage } from "../target";
import { MessageHandler } from "./types";
@@ -21,6 +22,7 @@ export const messageHandlers: Record<string, MessageHandler> = {
"olm/wg/server/peer/add": handleOlmServerPeerAddMessage,
"olm/wg/register": handleOlmRegisterMessage,
"olm/wg/relay": handleOlmRelayMessage,
"olm/wg/unrelay": handleOlmUnRelayMessage,
"olm/ping": handleOlmPingMessage,
"newt/wg/register": handleNewtRegisterMessage,
"newt/wg/get-config": handleGetConfigMessage,