@kortix/agent-tunnel 0.1.0 → 0.1.2

package/src/index.ts

@@ -28,6 +28,8 @@ export type {
   TunnelRelayConfig,
   HeartbeatConfig,
   TunnelServerConfig,
+  TunnelAuthMessage,
+  AuthResult,
   TunnelCapability,
   TunnelMethod,
   TunnelErrorCodeValue,

package/src/server/relay.ts

@@ -1,5 +1,5 @@
 import { EventEmitter } from 'events';
-import { signMessage } from '../shared/crypto';
+import { signMessage, verifyMessageSignature } from '../shared/crypto';
 import {
   type JsonRpcRequest,
   type JsonRpcResponse,
@@ -16,6 +16,7 @@ interface AgentConnection {
   ws: WebSocket;
   signingKey: string;
   nonce: number;
+  lastResponseNonce: number;
   connectedAt: number;
   metadata?: Record<string, unknown>;
 }
@@ -27,6 +28,8 @@ export class TunnelRelay extends EventEmitter {
   private pendingRPCs = new Map<string, PendingRPC>();
   private config: Required<TunnelRelayConfig>;
 
+  onAuthorizeRPC?: (tunnelId: string, method: string, params: Record<string, unknown>) => Promise<boolean>;
+
   constructor(config?: TunnelRelayConfig) {
     super();
     this.config = {
@@ -47,11 +50,21 @@ export class TunnelRelay extends EventEmitter {
   ): void {
     const existing = this.agents.get(tunnelId);
     if (existing) {
+      for (const [requestId, pending] of this.pendingRPCs) {
+        if (pending.tunnelId === tunnelId) {
+          clearTimeout(pending.timer);
+          pending.reject(new TunnelRelayError(
+            TunnelErrorCode.NOT_CONNECTED,
+            'Agent connection replaced',
+          ));
+          this.pendingRPCs.delete(requestId);
+        }
+      }
       try { existing.ws.close(1000, 'replaced by new connection'); } catch {}
       this.emitEvent('connection:replaced', { tunnelId });
     }
 
-    this.agents.set(tunnelId, { ws, signingKey, nonce: 0, connectedAt: Date.now(), metadata });
+    this.agents.set(tunnelId, { ws, signingKey, nonce: 0, lastResponseNonce: 0, connectedAt: Date.now(), metadata });
     this.emitEvent('agent:connect', { tunnelId, metadata });
     console.log(`[tunnel-relay] Agent registered: ${tunnelId} (total: ${this.agents.size})`);
   }
@@ -87,7 +100,6 @@ export class TunnelRelay extends EventEmitter {
     for (const [tunnelId, conn] of this.agents) {
       result.set(tunnelId, {
         tunnelId,
-        signingKey: conn.signingKey,
         connectedAt: conn.connectedAt,
         metadata: conn.metadata,
       });
@@ -100,7 +112,7 @@ export class TunnelRelay extends EventEmitter {
   }
 
   handleAgentMessage(tunnelId: string, raw: string | Buffer): void {
-    let msg: JsonRpcResponse;
+    let msg: any;
     try {
       msg = JSON.parse(typeof raw === 'string' ? raw : raw.toString('utf-8'));
     } catch {
@@ -108,8 +120,30 @@ export class TunnelRelay extends EventEmitter {
       return;
     }
 
-    if ('method' in msg && (msg as any).method === 'tunnel.pong') {
-      this.emitEvent('message:pong', { tunnelId, params: (msg as any).params });
+    // Verify HMAC signature on ALL messages from agent (including pong)
+    const agent = this.agents.get(tunnelId);
+    if (agent && msg._sig !== undefined && msg._nonce !== undefined) {
+      if (msg._nonce <= agent.lastResponseNonce) {
+        console.warn(`[tunnel-relay] Replay detected from agent ${tunnelId}: nonce ${msg._nonce} <= ${agent.lastResponseNonce}`);
+        return;
+      }
+
+      const { _sig, _nonce, ...payloadObj } = msg;
+      const payload = JSON.stringify(payloadObj);
+
+      if (!verifyMessageSignature(agent.signingKey, payload, _nonce, _sig)) {
+        console.warn(`[tunnel-relay] Invalid signature from agent ${tunnelId}`);
+        return;
+      }
+
+      agent.lastResponseNonce = _nonce;
+    } else if (agent) {
+      console.warn(`[tunnel-relay] Unsigned message from agent ${tunnelId}, discarding`);
+      return;
+    }
+   
+    if ('method' in msg && msg.method === 'tunnel.pong') {
+      this.emitEvent('message:pong', { tunnelId, params: msg.params });
       return;
     }
 
@@ -162,6 +196,16 @@ export class TunnelRelay extends EventEmitter {
       );
     }
 
+    if (this.onAuthorizeRPC) {
+      const allowed = await this.onAuthorizeRPC(tunnelId, method, params);
+      if (!allowed) {
+        throw new TunnelRelayError(
+          TunnelErrorCode.PERMISSION_DENIED,
+          `RPC ${method} denied for tunnel ${tunnelId}`,
+        );
+      }
+    }
+
     const requestId = crypto.randomUUID();
     const timeoutMs = options?.timeoutMs ?? this.config.rpcTimeoutMs;
 

package/src/server/routes.ts

@@ -3,9 +3,22 @@ import type { TunnelRelay } from './relay';
 import { TunnelRelayError } from './relay';
 import { TunnelErrorCode } from '../shared/types';
 
-export function createTunnelRouter(relay: TunnelRelay): Hono {
+export function createTunnelRouter(
+  relay: TunnelRelay,
+  onAuthorizeHTTP?: (req: Request) => Promise<boolean>,
+): Hono {
   const router = new Hono();
 
+  if (onAuthorizeHTTP) {
+    router.use('*', async (c, next) => {
+      const authorized = await onAuthorizeHTTP(c.req.raw);
+      if (!authorized) {
+        return c.json({ error: 'Unauthorized' }, 403);
+      }
+      await next();
+    });
+  }
+
   router.get('/connections', (c) => {
     const agents = relay.getConnectedAgents();
     const list = Array.from(agents.values()).map((a) => ({
@@ -51,6 +64,7 @@ export function createTunnelRouter(relay: TunnelRelay): Hono {
 
       const httpStatus = errorCode === TunnelErrorCode.NOT_CONNECTED ? 502
         : errorCode === TunnelErrorCode.TIMEOUT ? 504
+        : errorCode === TunnelErrorCode.PERMISSION_DENIED ? 403
         : 500;
 
       return c.json({ error: errorMessage, code: errorCode }, httpStatus);

package/src/server/server.ts

@@ -3,9 +3,10 @@
  *
  * Usage:
  *   import { startTunnelServer } from 'agent-tunnel';
- *   const server = startTunnelServer({ port: 8080 });
+ *   const server = startTunnelServer({ port: 8080, onAuthenticate: ... });
  *
- * That's it. Agents connect via: ws://localhost:8080/ws?token=xxx&tunnelId=yyy
+ * Agents connect via: ws://localhost:8080/ws?tunnelId=yyy
+ * Then send { type: "auth", token: "tnl_xxx" } as the first message.
  * Your app calls relay.relayRPC(tunnelId, method, params) to reach the agent.
  */
 
@@ -14,17 +15,13 @@ import { TunnelRelay } from './relay';
 import { HeartbeatManager } from './heartbeat';
 import { createTunnelRouter } from './routes';
 import { createWsHandlers } from './ws-handler';
-import { deriveSigningKey } from '../shared/crypto';
 import type { TunnelServerConfig } from '../shared/types';
 
 export interface TunnelServer {
-  /** Hono app — mount additional routes if needed. */
   app: Hono;
-  /** The relay instance — call relayRPC() to reach connected agents. */
   relay: TunnelRelay;
   heartbeat: HeartbeatManager;
   wsHandlers: ReturnType<typeof createWsHandlers>;
-  /** Stop the server, heartbeat, and close all connections. */
   stop: () => void;
 }
 
@@ -33,7 +30,7 @@ export interface TunnelServer {
  * - HTTP routes for connections list + RPC
  * - WebSocket upgrades on /ws for agent connections
  * - Heartbeat ping/pong
- * - HMAC signing key derivation from tokens
+ * - HMAC signing key derivation via onAuthenticate hook
  *
  * Requires Bun runtime for WebSocket server support.
  */
@@ -41,57 +38,60 @@ export function startTunnelServer(config?: TunnelServerConfig): TunnelServer {
   const port = config?.port ?? parseInt(process.env.PORT || '8080', 10);
 
   const relay = new TunnelRelay(config?.relay);
+
+  if (config?.onAuthorizeRPC) {
+    relay.onAuthorizeRPC = config.onAuthorizeRPC;
+  }
+
   const heartbeat = new HeartbeatManager(relay, config?.heartbeat);
-  const wsHandlers = createWsHandlers(relay, { heartbeat });
+  const wsHandlers = createWsHandlers(relay, {
+    heartbeat,
+    onAuthenticate: config?.onAuthenticate,
+  });
 
   const app = new Hono();
 
-  // Mount tunnel routes (GET /connections, POST /rpc/:tunnelId, etc.)
-  const tunnelRouter = createTunnelRouter(relay);
+  const tunnelRouter = createTunnelRouter(relay, config?.onAuthorizeHTTP);
   app.route('/', tunnelRouter);
 
-  // Health check
   app.get('/health', (c) => c.json({ status: 'ok', connections: relay.getConnectedCount() }));
 
-  // Start heartbeat
+  relay.on('message:pong', (data: { tunnelId: string }) => {
+    heartbeat.recordPong(data.tunnelId);
+  });
+
   heartbeat.start();
 
-  // Start Bun server with WS support
   const bunServer = Bun.serve({
     port,
     fetch(req, server) {
       const url = new URL(req.url);
 
-      // WS upgrade on /ws
       if (url.pathname === '/ws') {
-        const token = url.searchParams.get('token');
         const tunnelId = url.searchParams.get('tunnelId');
 
-        if (!token || !tunnelId) {
-          return new Response(JSON.stringify({ error: 'Missing token or tunnelId' }), {
+        if (!tunnelId) {
+          return new Response(JSON.stringify({ error: 'Missing tunnelId' }), {
             status: 400,
             headers: { 'Content-Type': 'application/json' },
           });
         }
 
-        const signingKey = deriveSigningKey(token);
-
         const success = server.upgrade(req, {
-          data: { tunnelId, signingKey } as any,
+          data: { tunnelId } as any,
         });
         if (success) return undefined;
 
         return new Response('WebSocket upgrade failed', { status: 500 });
       }
 
-      // All other requests → Hono
       return app.fetch(req);
     },
 
     websocket: {
       idleTimeout: 0,
       open(ws: any) {
-        wsHandlers.onOpen(ws.data.tunnelId, ws, ws.data.signingKey);
+        wsHandlers.onOpen(ws.data.tunnelId, ws);
       },
       message(ws: any, message: string | Buffer) {
         wsHandlers.onMessage(ws.data.tunnelId, message);

package/src/server/ws-handler.ts

@@ -1,13 +1,16 @@
 import type { TunnelRelay } from './relay';
 import type { HeartbeatManager } from './heartbeat';
+import type { TunnelAuthMessage, AuthResult } from '../shared/types';
 
 export interface WsHandlerOptions {
   heartbeat?: HeartbeatManager;
   maxMessageSize?: number;
+  onAuthenticate?: (tunnelId: string, token: string) => Promise<AuthResult | null>;
+  authTimeoutMs?: number;
 }
 
 export interface WsHandlers {
-  onOpen(tunnelId: string, ws: WebSocket, signingKey: string, metadata?: Record<string, unknown>): void;
+  onOpen(tunnelId: string, ws: WebSocket): void;
   onMessage(tunnelId: string, message: string | Buffer): void;
   onClose(tunnelId: string): void;
 }
@@ -15,36 +18,97 @@ export interface WsHandlers {
 export function createWsHandlers(relay: TunnelRelay, opts?: WsHandlerOptions): WsHandlers {
   const maxMessageSize = opts?.maxMessageSize ?? 5 * 1024 * 1024;
   const heartbeat = opts?.heartbeat;
+  const onAuthenticate = opts?.onAuthenticate;
+  const authTimeoutMs = opts?.authTimeoutMs ?? 10_000;
+
+  const pendingConnections = new Map<string, { ws: WebSocket; timer: ReturnType<typeof setTimeout> }>();
 
   return {
-    onOpen(tunnelId: string, ws: WebSocket, signingKey: string, metadata?: Record<string, unknown>) {
-      relay.registerAgent(tunnelId, ws, signingKey, metadata);
-      if (heartbeat) {
-        heartbeat.register(tunnelId);
-      }
+    onOpen(tunnelId: string, ws: WebSocket) {
+      const timer = setTimeout(() => {
+        pendingConnections.delete(tunnelId);
+        try { ws.close(4001, 'auth timeout'); } catch {}
+      }, authTimeoutMs);
+
+      pendingConnections.set(tunnelId, { ws, timer });
     },
 
-    onMessage(tunnelId: string, message: string | Buffer) {
+    async onMessage(tunnelId: string, message: string | Buffer) {
+      const msgStr = typeof message === 'string' ? message : message.toString('utf-8');
       const msgSize = typeof message === 'string' ? message.length : (message as Buffer).byteLength;
+
       if (msgSize > maxMessageSize) {
         console.warn(`[tunnel-ws] Oversized message from ${tunnelId}: ${msgSize} bytes (limit: ${maxMessageSize})`);
+        const pending = pendingConnections.get(tunnelId);
+        const ws = pending?.ws;
+        if (ws) {
+          try { ws.close(4002, 'message too large'); } catch {}
+        }
         return;
       }
 
-      try {
-        const parsed = JSON.parse(typeof message === 'string' ? message : message.toString('utf-8'));
-        if (parsed.method === 'tunnel.pong') {
+      const pending = pendingConnections.get(tunnelId);
+      if (pending) {
+        let authMsg: TunnelAuthMessage;
+        try {
+          authMsg = JSON.parse(msgStr);
+        } catch {
+          try { pending.ws.close(4001, 'invalid auth message'); } catch {}
+          clearTimeout(pending.timer);
+          pendingConnections.delete(tunnelId);
+          return;
+        }
+
+        if (authMsg.type !== 'auth' || !authMsg.token) {
+          try { pending.ws.close(4001, 'expected auth message'); } catch {}
+          clearTimeout(pending.timer);
+          pendingConnections.delete(tunnelId);
+          return;
+        }
+
+        clearTimeout(pending.timer);
+        pendingConnections.delete(tunnelId);
+
+        if (!onAuthenticate) {
+          try { pending.ws.close(4001, 'no authenticator configured'); } catch {}
+          return;
+        }
+
+        try {
+          const result = await onAuthenticate(tunnelId, authMsg.token);
+          if (!result) {
+            try { pending.ws.close(4001, 'authentication failed'); } catch {}
+            return;
+          }
+
+          // Send signing key to agent so it never needs the server secret
+          try {
+            pending.ws.send(JSON.stringify({ type: 'auth_ok', signingKey: result.signingKey }));
+          } catch {}
+
+          relay.registerAgent(tunnelId, pending.ws, result.signingKey, result.metadata);
           if (heartbeat) {
-            heartbeat.recordPong(tunnelId);
+            heartbeat.register(tunnelId);
           }
+        } catch (err) {
+          console.error(`[tunnel-ws] Auth error for ${tunnelId}:`, err);
+          try { pending.ws.close(4001, 'authentication error'); } catch {}
         }
-      } catch {
+
+        return;
       }
 
       relay.handleAgentMessage(tunnelId, message);
     },
 
     onClose(tunnelId: string) {
+      const pending = pendingConnections.get(tunnelId);
+      if (pending) {
+        clearTimeout(pending.timer);
+        pendingConnections.delete(tunnelId);
+        return;
+      }
+
       relay.unregisterAgent(tunnelId);
       if (heartbeat) {
         heartbeat.unregister(tunnelId);

package/src/shared/crypto.ts

@@ -1,6 +1,5 @@
 import { createHash, createHmac, timingSafeEqual, randomBytes } from 'crypto';
 
-const SIGNING_KEY_CONTEXT = 'kortix-tunnel-signing-v1';
 const CHARS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
 
 function randomAlphanumeric(length: number): string {
@@ -12,8 +11,8 @@ function randomAlphanumeric(length: number): string {
   return result;
 }
 
-export function deriveSigningKey(token: string): string {
-  return createHmac('sha256', SIGNING_KEY_CONTEXT)
+export function deriveSigningKey(token: string, secret: string): string {
+  return createHmac('sha256', secret)
     .update(token)
     .digest('hex');
 }

package/src/shared/index.ts

@@ -25,6 +25,8 @@ export type {
   TunnelRelayConfig,
   HeartbeatConfig,
   TunnelServerConfig,
+  TunnelAuthMessage,
+  AuthResult,
   TunnelCapability,
   TunnelMethod,
   TunnelErrorCodeValue,

package/src/shared/types.ts

@@ -41,6 +41,7 @@ export const TunnelErrorCode = {
   NOT_CONNECTED: -32004,
   EXPIRED: -32005,
   RATE_LIMITED: -32006,
+  AUTH_FAILED: -32007,
 } as const;
 
 export type TunnelErrorCodeValue = (typeof TunnelErrorCode)[keyof typeof TunnelErrorCode];
@@ -128,9 +129,9 @@ export interface RelayRpcOptions {
   timeoutMs?: number;
 }
 
+/** Public agent info — does NOT expose signing key. */
 export interface AgentInfo {
   tunnelId: string;
-  signingKey: string;
   connectedAt: number;
   metadata?: Record<string, unknown>;
 }
@@ -157,8 +158,36 @@ export interface HeartbeatConfig {
   maxMissed?: number;
 }
 
+/** Auth handshake message sent by agent as first WS message. */
+export interface TunnelAuthMessage {
+  type: 'auth';
+  token: string;
+}
+
+/** Result returned by onAuthenticate hook on success. */
+export interface AuthResult {
+  signingKey: string;
+  metadata?: Record<string, unknown>;
+}
+
 export interface TunnelServerConfig {
   port?: number;
   relay?: TunnelRelayConfig;
   heartbeat?: HeartbeatConfig;
+  /**
+   * Called when an agent sends its auth handshake.
+   * Return { signingKey, metadata } to accept, or null to reject.
+   * If not provided, all connections are rejected.
+   */
+  onAuthenticate?: (tunnelId: string, token: string) => Promise<AuthResult | null>;
+  /**
+   * Called before relaying an RPC to the agent.
+   * Return false to deny. If not provided, all RPCs are allowed.
+   */
+  onAuthorizeRPC?: (tunnelId: string, method: string, params: Record<string, unknown>) => Promise<boolean>;
+  /**
+   * Called before handling HTTP requests to relay routes (/connections, /rpc).
+   * Return false to deny. If not provided, routes are open.
+   */
+  onAuthorizeHTTP?: (req: Request) => Promise<boolean>;
 }