@dainprotocol/tunnel 1.0.7 → 1.1.0

package/dist/client/index.d.ts

@@ -9,11 +9,18 @@ declare class DainTunnel extends EventEmitter {
     private reconnectDelay;
     private apiKey;
     private tunnelId;
+    private secret;
     private webSocketClients;
     private sseClients;
     constructor(serverUrl: string, apiKey: string);
+    /**
+     * Sign a challenge using HMAC-SHA256
+     * @private
+     */
+    private signChallenge;
     start(port: number): Promise<string>;
     private connect;
+    private requestChallenge;
     private handleMessage;
     private handleRequest;
     private handleWebSocketConnection;

package/dist/client/index.js

@@ -8,6 +8,7 @@ const ws_1 = __importDefault(require("ws"));
 const http_1 = __importDefault(require("http"));
 const events_1 = require("events");
 const crypto_1 = require("crypto");
+const auth_1 = require("@dainprotocol/service-sdk/service/auth");
 class DainTunnel extends events_1.EventEmitter {
     constructor(serverUrl, apiKey) {
         super();
@@ -20,10 +21,23 @@ class DainTunnel extends events_1.EventEmitter {
         this.reconnectDelay = 5000;
         this.webSocketClients = new Map();
         this.sseClients = new Map();
+        // Parse API key to extract agentId and secret
+        const parsed = (0, auth_1.parseAPIKey)(apiKey);
+        if (!parsed) {
+            throw new Error('Invalid API key format. Expected: sk_agent_{agentId}_{orgId}_{secret}');
+        }
         this.apiKey = apiKey;
-        // Create tunnel ID from API key hash
-        const hash = (0, crypto_1.createHash)('sha256').update(apiKey).digest();
-        this.tunnelId = hash.toString('base64url').substring(0, 16);
+        this.tunnelId = parsed.agentId; // agentId is the tunnel identifier
+        this.secret = parsed.secret; // secret for HMAC signatures
+    }
+    /**
+     * Sign a challenge using HMAC-SHA256
+     * @private
+     */
+    signChallenge(challenge) {
+        return (0, crypto_1.createHmac)('sha256', this.secret)
+            .update(challenge)
+            .digest('hex');
     }
     async start(port) {
         this.port = port;
@@ -34,17 +48,26 @@ class DainTunnel extends events_1.EventEmitter {
             try {
                 console.log(`Connecting to WebSocket server: ${this.serverUrl}`);
                 this.ws = new ws_1.default(this.serverUrl);
-                this.ws.on("open", () => {
+                this.ws.on("open", async () => {
                     console.log('WebSocket connection opened');
                     this.reconnectAttempts = 0;
-                    // Send API key authentication directly (no challenge-response)
-                    this.sendMessage({
-                        type: "start",
-                        port: this.port,
-                        apiKey: this.apiKey,
-                        tunnelId: this.tunnelId
-                    });
-                    this.emit("connected");
+                    try {
+                        const challenge = await this.requestChallenge();
+                        const signature = this.signChallenge(challenge);
+                        this.sendMessage({
+                            type: "start",
+                            port: this.port,
+                            challenge,
+                            signature,
+                            tunnelId: this.tunnelId,
+                            apiKey: this.apiKey // Send API key for server validation
+                        });
+                        this.emit("connected");
+                    }
+                    catch (err) {
+                        console.error('Error during challenge-response:', err);
+                        reject(err);
+                    }
                 });
                 this.ws.on("message", (data) => {
                     try {
@@ -86,6 +109,32 @@ class DainTunnel extends events_1.EventEmitter {
             }
         });
     }
+    async requestChallenge() {
+        return new Promise((resolve, reject) => {
+            if (!this.ws) {
+                reject(new Error("WebSocket is not connected"));
+                return;
+            }
+            this.ws.send(JSON.stringify({ type: "challenge_request" }));
+            const challengeHandler = (message) => {
+                const data = JSON.parse(message);
+                if (data.type === "challenge") {
+                    if (this.ws) {
+                        this.ws.removeListener("message", challengeHandler);
+                    }
+                    resolve(data.challenge);
+                }
+            };
+            this.ws.on("message", challengeHandler);
+            // Add a timeout for the challenge request
+            setTimeout(() => {
+                if (this.ws) {
+                    this.ws.removeListener("message", challengeHandler);
+                }
+                reject(new Error("Challenge request timeout"));
+            }, 5000);
+        });
+    }
     handleMessage(message, resolve) {
         switch (message.type) {
             case "tunnelUrl":

package/dist/server/index.d.ts

@@ -6,6 +6,7 @@ declare class DainTunnelServer {
     private wss;
     private tunnels;
     private pendingRequests;
+    private challenges;
     private sseConnections;
     private wsConnections;
     constructor(hostname: string, port: number);
@@ -13,6 +14,7 @@ declare class DainTunnelServer {
     private setupWebSocketServer;
     private handleTunnelClientConnection;
     private handleProxiedWebSocketConnection;
+    private handleChallengeRequest;
     private handleStartMessage;
     private handleResponseMessage;
     private handleSSEMessage;

package/dist/server/index.js

@@ -9,6 +9,8 @@ const ws_1 = __importDefault(require("ws"));
 const uuid_1 = require("uuid");
 const body_parser_1 = __importDefault(require("body-parser"));
 const cors_1 = __importDefault(require("cors"));
+const auth_1 = require("@dainprotocol/service-sdk/service/auth");
+const crypto_1 = require("crypto");
 const url_1 = require("url");
 class DainTunnelServer {
     constructor(hostname, port) {
@@ -16,6 +18,7 @@ class DainTunnelServer {
         this.port = port;
         this.tunnels = new Map();
         this.pendingRequests = new Map();
+        this.challenges = new Map();
         this.sseConnections = new Map();
         this.wsConnections = new Map();
         this.app = (0, express_1.default)();
@@ -82,7 +85,10 @@ class DainTunnelServer {
             try {
                 const data = JSON.parse(message);
                 console.log(`Received WebSocket message: ${data.type}`);
-                if (data.type === "start") {
+                if (data.type === "challenge_request") {
+                    this.handleChallengeRequest(ws);
+                }
+                else if (data.type === "start") {
                     this.handleStartMessage(ws, data);
                 }
                 else if (data.type === "response") {
@@ -171,16 +177,44 @@ class DainTunnelServer {
             }));
         });
     }
+    handleChallengeRequest(ws) {
+        const challenge = (0, uuid_1.v4)();
+        const challengeObj = { ws, challenge };
+        this.challenges.set(challenge, challengeObj);
+        ws.send(JSON.stringify({ type: "challenge", challenge }));
+    }
     handleStartMessage(ws, data) {
-        const { apiKey, tunnelId } = data;
-        // Validate API key format (sk_agent_{agentId}_{orgId}_{secret})
-        if (!apiKey || !apiKey.startsWith('sk_agent_')) {
-            ws.close(1008, "Invalid API key");
+        const { challenge, signature, tunnelId, apiKey } = data;
+        const challengeObj = this.challenges.get(challenge);
+        if (!challengeObj || challengeObj.ws !== ws) {
+            ws.close(1008, "Invalid challenge");
             return;
         }
-        // TODO: Validate API key with platform in production
-        // For now, accept any properly formatted API key (dev mode)
+        this.challenges.delete(challenge);
         try {
+            // Parse API key to get the secret for HMAC validation
+            if (!apiKey) {
+                ws.close(1008, "API key required");
+                return;
+            }
+            const parsed = (0, auth_1.parseAPIKey)(apiKey);
+            if (!parsed) {
+                ws.close(1008, "Invalid API key format");
+                return;
+            }
+            // Validate HMAC signature
+            const expectedSignature = (0, crypto_1.createHmac)('sha256', parsed.secret)
+                .update(challenge)
+                .digest('hex');
+            if (expectedSignature !== signature) {
+                ws.close(1008, "Invalid signature");
+                return;
+            }
+            // Verify that tunnelId matches the agentId from the API key
+            if (tunnelId !== parsed.agentId) {
+                ws.close(1008, "Tunnel ID does not match API key");
+                return;
+            }
             // If tunnel already exists, remove old one
             if (this.tunnels.has(tunnelId)) {
                 console.log(`Tunnel ${tunnelId} already exists, replacing it`);
@@ -442,6 +476,14 @@ class DainTunnelServer {
             else {
                 console.log(`No tunnel found to remove for the closed WebSocket connection`);
             }
+            // Also remove any pending challenges for this WebSocket
+            for (const [challenge, challengeObj] of this.challenges.entries()) {
+                if (challengeObj.ws === ws) {
+                    this.challenges.delete(challenge);
+                    console.log(`Challenge removed for closed WebSocket`);
+                    break;
+                }
+            }
         }
         catch (error) {
             console.error("Error in removeTunnel:", error);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@dainprotocol/tunnel",
-  "version": "1.0.7",
+  "version": "1.1.0",
   "description": "",
   "main": "dist/index.js",
   "private": false,
@@ -19,7 +19,7 @@
   "author": "Ryan",
   "license": "ISC",
   "dependencies": {
-    "@dainprotocol/service-sdk": "^1.2.5",
+    "@dainprotocol/service-sdk": "^1.3.0",
     "@types/body-parser": "^1.19.5",
     "@types/cors": "^2.8.17",
     "@types/eventsource": "^3.0.0",