@dainprotocol/tunnel 1.0.4 → 1.0.7

package/dist/client/index.d.ts

@@ -7,14 +7,19 @@ declare class DainTunnel extends EventEmitter {
     private reconnectAttempts;
     private maxReconnectAttempts;
     private reconnectDelay;
-    private auth;
+    private apiKey;
     private tunnelId;
+    private webSocketClients;
+    private sseClients;
     constructor(serverUrl: string, apiKey: string);
     start(port: number): Promise<string>;
     private connect;
-    private requestChallenge;
     private handleMessage;
     private handleRequest;
+    private handleWebSocketConnection;
+    private handleWebSocketMessage;
+    private handleSSEConnection;
+    private handleSSEClose;
     private forwardRequest;
     private sendMessage;
     private attemptReconnect;

package/dist/client/index.js

@@ -7,8 +7,7 @@ exports.DainTunnel = void 0;
 const ws_1 = __importDefault(require("ws"));
 const http_1 = __importDefault(require("http"));
 const events_1 = require("events");
-const bs58_1 = __importDefault(require("bs58"));
-const client_1 = require("@dainprotocol/service-sdk/client");
+const crypto_1 = require("crypto");
 class DainTunnel extends events_1.EventEmitter {
     constructor(serverUrl, apiKey) {
         super();
@@ -19,8 +18,12 @@ class DainTunnel extends events_1.EventEmitter {
         this.reconnectAttempts = 0;
         this.maxReconnectAttempts = 5;
         this.reconnectDelay = 5000;
-        this.auth = new client_1.DainClientAuth({ apiKey });
-        this.tunnelId = bs58_1.default.encode(this.auth.getPublicKey());
+        this.webSocketClients = new Map();
+        this.sseClients = new Map();
+        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);
     }
     async start(port) {
         this.port = port;
@@ -28,65 +31,59 @@ class DainTunnel extends events_1.EventEmitter {
     }
     async connect() {
         return new Promise((resolve, reject) => {
-            this.ws = new ws_1.default(this.serverUrl);
-            this.ws.on("open", async () => {
-                this.reconnectAttempts = 0;
-                const challenge = await this.requestChallenge();
-                const signature = this.auth.signMessage(challenge);
-                this.sendMessage({
-                    type: "start",
-                    port: this.port,
-                    challenge,
-                    signature,
-                    tunnelId: this.tunnelId
+            try {
+                console.log(`Connecting to WebSocket server: ${this.serverUrl}`);
+                this.ws = new ws_1.default(this.serverUrl);
+                this.ws.on("open", () => {
+                    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");
                 });
-                this.emit("connected");
-            });
-            this.ws.on("message", (data) => {
-                const message = JSON.parse(data);
-                this.handleMessage(message, resolve);
-            });
-            this.ws.on("close", () => {
-                this.emit("disconnected");
-                this.attemptReconnect();
-            });
-            this.ws.on("error", (error) => {
-                this.emit("error", error);
-                reject(error);
-            });
-            // Add a timeout to reject the promise if connection takes too long
-            setTimeout(() => {
-                var _a;
-                if (((_a = this.ws) === null || _a === void 0 ? void 0 : _a.readyState) !== ws_1.default.OPEN) {
-                    reject(new Error("Connection timeout"));
-                }
-            }, 5000);
-        });
-    }
-    async requestChallenge() {
-        return new Promise((resolve, reject) => {
-            if (!this.ws) {
-                reject(new Error("WebSocket is not connected"));
-                return;
+                this.ws.on("message", (data) => {
+                    try {
+                        const message = JSON.parse(data);
+                        this.handleMessage(message, resolve);
+                    }
+                    catch (err) {
+                        console.error('Error handling message:', err);
+                        reject(err);
+                    }
+                });
+                this.ws.on("close", () => {
+                    console.log('WebSocket connection closed');
+                    if (this.tunnelUrl) {
+                        this.emit("disconnected");
+                        this.attemptReconnect();
+                    }
+                    else {
+                        // If we haven't received a tunnelUrl yet, just reject
+                        reject(new Error("Connection closed before tunnel was established"));
+                    }
+                });
+                this.ws.on("error", (error) => {
+                    // Prevent unhandled error events
+                    // Don't reject here, the close handler will be called after an error
+                    this.emit("error", error);
+                });
+                // Add a timeout to reject the promise if connection takes too long
+                setTimeout(() => {
+                    if (!this.ws || this.ws.readyState !== ws_1.default.OPEN) {
+                        const timeoutError = new Error("Connection timeout");
+                        reject(timeoutError);
+                    }
+                }, 10000);
+            }
+            catch (err) {
+                console.error('Error creating WebSocket:', err);
+                reject(err);
             }
-            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) {
@@ -99,6 +96,18 @@ class DainTunnel extends events_1.EventEmitter {
             case "request":
                 this.handleRequest(message);
                 break;
+            case "websocket_connection":
+                this.handleWebSocketConnection(message);
+                break;
+            case "sse_connection":
+                this.handleSSEConnection(message);
+                break;
+            case "sse_close":
+                this.handleSSEClose(message);
+                break;
+            case "websocket":
+                this.handleWebSocketMessage(message);
+                break;
         }
     }
     async handleRequest(request) {
@@ -111,6 +120,191 @@ class DainTunnel extends events_1.EventEmitter {
             this.emit("request_error", { request, error });
         }
     }
+    handleWebSocketConnection(message) {
+        try {
+            console.log(`Creating local WebSocket connection to: ws://localhost:${this.port}${message.path}`);
+            // Create a WebSocket connection to the local server
+            const client = new ws_1.default(`ws://localhost:${this.port}${message.path}`, {
+                headers: message.headers
+            });
+            // Store the client
+            this.webSocketClients.set(message.id, client);
+            // Handle connection open
+            client.on('open', () => {
+                console.log(`Local WebSocket connection opened: ${message.id}`);
+            });
+            // Handle messages from the local server
+            client.on('message', (data) => {
+                if (this.ws && this.ws.readyState === ws_1.default.OPEN) {
+                    this.ws.send(JSON.stringify({
+                        type: 'websocket',
+                        id: message.id,
+                        event: 'message',
+                        data: data.toString('base64')
+                    }));
+                }
+            });
+            // Handle close
+            client.on('close', () => {
+                console.log(`Local WebSocket connection closed: ${message.id}`);
+                if (this.ws && this.ws.readyState === ws_1.default.OPEN) {
+                    this.ws.send(JSON.stringify({
+                        type: 'websocket',
+                        id: message.id,
+                        event: 'close'
+                    }));
+                }
+                this.webSocketClients.delete(message.id);
+            });
+            // Handle errors
+            client.on('error', (error) => {
+                console.error(`Local WebSocket connection error: ${message.id}`, error.message);
+                if (this.ws && this.ws.readyState === ws_1.default.OPEN) {
+                    this.ws.send(JSON.stringify({
+                        type: 'websocket',
+                        id: message.id,
+                        event: 'error',
+                        data: error.message
+                    }));
+                }
+                this.webSocketClients.delete(message.id);
+            });
+            this.emit("websocket_connection", { id: message.id, path: message.path });
+        }
+        catch (error) {
+            console.error('Error establishing WebSocket connection:', error);
+            if (this.ws && this.ws.readyState === ws_1.default.OPEN) {
+                this.ws.send(JSON.stringify({
+                    type: 'websocket',
+                    id: message.id,
+                    event: 'error',
+                    data: error.message
+                }));
+            }
+        }
+    }
+    handleWebSocketMessage(message) {
+        const client = this.webSocketClients.get(message.id);
+        if (!client)
+            return;
+        if (message.event === 'message' && message.data) {
+            const data = Buffer.from(message.data, 'base64');
+            if (client.readyState === ws_1.default.OPEN) {
+                client.send(data);
+            }
+        }
+        else if (message.event === 'close') {
+            if (client.readyState === ws_1.default.OPEN) {
+                client.close();
+            }
+            this.webSocketClients.delete(message.id);
+        }
+    }
+    handleSSEConnection(message) {
+        try {
+            // Create an EventSource-like stream to the local server
+            // Since Node.js doesn't have a built-in EventSource, we'll use HTTP
+            const options = {
+                hostname: 'localhost',
+                port: this.port,
+                path: message.path,
+                method: 'GET',
+                headers: message.headers,
+            };
+            const req = http_1.default.request(options, (res) => {
+                if (res.statusCode !== 200) {
+                    if (this.ws && this.ws.readyState === ws_1.default.OPEN) {
+                        this.ws.send(JSON.stringify({
+                            type: 'sse',
+                            id: message.id,
+                            event: 'error',
+                            data: `Server responded with status code ${res.statusCode}`
+                        }));
+                    }
+                    return;
+                }
+                // Process SSE stream
+                let buffer = '';
+                res.on('data', (chunk) => {
+                    buffer += chunk.toString();
+                    // Process complete SSE messages
+                    while (buffer.includes('\n\n')) {
+                        const messageEndIndex = buffer.indexOf('\n\n');
+                        const messageData = buffer.substring(0, messageEndIndex);
+                        buffer = buffer.substring(messageEndIndex + 2);
+                        // Parse the SSE message
+                        const lines = messageData.split('\n');
+                        let event = 'message';
+                        let data = '';
+                        for (const line of lines) {
+                            if (line.startsWith('event:')) {
+                                event = line.substring(6).trim();
+                            }
+                            else if (line.startsWith('data:')) {
+                                data += line.substring(5).trim() + '\n';
+                            }
+                        }
+                        // Remove trailing newline
+                        if (data.endsWith('\n')) {
+                            data = data.substring(0, data.length - 1);
+                        }
+                        // Forward to server
+                        if (this.ws && this.ws.readyState === ws_1.default.OPEN) {
+                            this.ws.send(JSON.stringify({
+                                type: 'sse',
+                                id: message.id,
+                                event,
+                                data
+                            }));
+                        }
+                    }
+                });
+                res.on('end', () => {
+                    if (this.ws && this.ws.readyState === ws_1.default.OPEN) {
+                        this.ws.send(JSON.stringify({
+                            type: 'sse',
+                            id: message.id,
+                            event: 'close',
+                            data: 'Connection closed'
+                        }));
+                    }
+                });
+            });
+            req.on('error', (error) => {
+                if (this.ws && this.ws.readyState === ws_1.default.OPEN) {
+                    this.ws.send(JSON.stringify({
+                        type: 'sse',
+                        id: message.id,
+                        event: 'error',
+                        data: error.message
+                    }));
+                }
+            });
+            req.end();
+            // Store a reference to abort the connection later if needed
+            this.sseClients.set(message.id, req);
+            this.emit("sse_connection", { id: message.id, path: message.path });
+        }
+        catch (error) {
+            console.error('Error establishing SSE connection:', error);
+            if (this.ws && this.ws.readyState === ws_1.default.OPEN) {
+                this.ws.send(JSON.stringify({
+                    type: 'sse',
+                    id: message.id,
+                    event: 'error',
+                    data: error.message
+                }));
+            }
+        }
+    }
+    handleSSEClose(message) {
+        const client = this.sseClients.get(message.id);
+        if (client) {
+            // Abort the request if it's still active
+            client.destroy();
+            this.sseClients.delete(message.id);
+        }
+    }
     forwardRequest(request) {
         return new Promise((resolve, reject) => {
             const options = {
@@ -153,19 +347,73 @@ class DainTunnel extends events_1.EventEmitter {
     attemptReconnect() {
         if (this.reconnectAttempts < this.maxReconnectAttempts) {
             this.reconnectAttempts++;
-            setTimeout(() => this.connect(), this.reconnectDelay);
+            console.log(`Attempting to reconnect (${this.reconnectAttempts}/${this.maxReconnectAttempts})`);
+            setTimeout(async () => {
+                try {
+                    await this.connect();
+                    // If the connect succeeds but we don't have a tunnelUrl yet, we're still initializing
+                    if (!this.tunnelUrl) {
+                        console.log('Reconnected but tunnelUrl not set. Waiting for tunnel URL...');
+                    }
+                    else {
+                        console.log(`Reconnected successfully. Tunnel URL: ${this.tunnelUrl}`);
+                    }
+                }
+                catch (error) {
+                    console.error(`Reconnection attempt ${this.reconnectAttempts} failed:`, error);
+                }
+            }, this.reconnectDelay);
         }
         else {
+            console.log('Maximum reconnection attempts reached');
             this.emit("max_reconnect_attempts");
         }
     }
     async stop() {
-        return new Promise((resolve) => {
-            if (this.ws) {
-                this.ws.close();
-                this.ws = null;
+        return new Promise(async (resolve) => {
+            try {
+                // Close all WebSocket clients
+                for (const [id, client] of this.webSocketClients.entries()) {
+                    try {
+                        if (client.readyState === ws_1.default.OPEN) {
+                            client.close();
+                        }
+                    }
+                    catch (error) {
+                        console.error(`Error closing WebSocket client ${id}:`, error);
+                    }
+                }
+                this.webSocketClients.clear();
+                // Close all SSE clients
+                for (const [id, client] of this.sseClients.entries()) {
+                    try {
+                        client.destroy();
+                    }
+                    catch (error) {
+                        console.error(`Error destroying SSE client ${id}:`, error);
+                    }
+                }
+                this.sseClients.clear();
+                // Close main WebSocket connection
+                if (this.ws) {
+                    try {
+                        if (this.ws.readyState === ws_1.default.OPEN) {
+                            this.ws.close();
+                        }
+                    }
+                    catch (error) {
+                        console.error('Error closing main WebSocket:', error);
+                    }
+                    this.ws = null;
+                }
+                // Wait for all connections to close properly
+                await new Promise(resolve => setTimeout(resolve, 500));
+                resolve();
+            }
+            catch (error) {
+                console.error('Error in stop method:', error);
+                resolve(); // Resolve anyway to not block cleaning up
             }
-            resolve();
         });
     }
 }

package/dist/server/index.d.ts

@@ -6,14 +6,19 @@ declare class DainTunnelServer {
     private wss;
     private tunnels;
     private pendingRequests;
-    private challenges;
+    private sseConnections;
+    private wsConnections;
     constructor(hostname: string, port: number);
     private setupExpressRoutes;
     private setupWebSocketServer;
-    private handleChallengeRequest;
+    private handleTunnelClientConnection;
+    private handleProxiedWebSocketConnection;
     private handleStartMessage;
     private handleResponseMessage;
-    private handleHttpRequest;
+    private handleSSEMessage;
+    private handleWebSocketMessage;
+    private handleRequest;
+    private handleSSERequest;
     private removeTunnel;
     start(): Promise<void>;
     stop(): Promise<void>;

package/dist/server/index.js

@@ -9,20 +9,27 @@ 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 client_1 = require("@dainprotocol/service-sdk/client");
-const bs58_1 = __importDefault(require("bs58"));
+const url_1 = require("url");
 class DainTunnelServer {
     constructor(hostname, port) {
         this.hostname = hostname;
         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)();
         this.server = http_1.default.createServer(this.app);
-        this.wss = new ws_1.default.Server({ server: this.server });
+        this.wss = new ws_1.default.Server({
+            server: this.server,
+            path: undefined // Allow connections on any path
+        });
         // Use cors middleware
-        this.app.use((0, cors_1.default)());
+        this.app.use((0, cors_1.default)({
+            origin: '*',
+            methods: ['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'],
+            credentials: true
+        }));
         // Update CORS middleware
         this.app.use((req, res, next) => {
             res.header("Access-Control-Allow-Origin", "*");
@@ -39,81 +46,200 @@ class DainTunnelServer {
         this.setupWebSocketServer();
     }
     setupExpressRoutes() {
-        this.app.use("/:tunnelId", this.handleHttpRequest.bind(this));
+        // Generic route handler for all tunnel requests
+        this.app.use("/:tunnelId", this.handleRequest.bind(this));
     }
     setupWebSocketServer() {
-        this.wss.on("connection", (ws) => {
+        // Handle WebSocket connections from tunnel clients
+        this.wss.on("connection", (ws, req) => {
+            var _a;
             console.log("New WebSocket connection");
-            ws.on("message", (message) => {
-                try {
-                    const data = JSON.parse(message);
-                    console.log(`Received WebSocket message: ${data.type}`);
-                    if (data.type === "challenge_request") {
-                        this.handleChallengeRequest(ws);
-                    }
-                    else if (data.type === "start") {
-                        this.handleStartMessage(ws, data);
-                    }
-                    else if (data.type === "response") {
-                        this.handleResponseMessage(data);
-                    }
+            try {
+                // Check if this is a tunnel client connection or a user connection to be proxied
+                const url = (0, url_1.parse)(req.url || '', true);
+                const pathParts = ((_a = url.pathname) === null || _a === void 0 ? void 0 : _a.split('/')) || [];
+                console.log(`WebSocket connection path: ${url.pathname}`);
+                // Client tunnel connection has no tunnelId in the path (or is at root)
+                if (!url.pathname || url.pathname === '/' || pathParts.length <= 1 || !pathParts[1]) {
+                    console.log("Handling as tunnel client connection");
+                    this.handleTunnelClientConnection(ws, req);
                 }
-                catch (error) {
-                    console.error("Error processing message:", error);
-                    ws.close(1008, "Invalid message");
+                else {
+                    // This is a WebSocket connection to be proxied through the tunnel
+                    console.log(`Handling as proxied WebSocket connection for path: ${url.pathname}`);
+                    this.handleProxiedWebSocketConnection(ws, req);
                 }
-            });
-            ws.on("close", () => {
-                console.log("WebSocket connection closed");
-                this.removeTunnel(ws);
-            });
-            ws.on("error", (error) => {
-                console.error("WebSocket error:", error);
-            });
+            }
+            catch (error) {
+                console.error("Error in WebSocket connection setup:", error);
+                ws.close(1011, "Internal server error");
+            }
         });
     }
-    handleChallengeRequest(ws) {
-        const challenge = (0, uuid_1.v4)();
-        const challengeObj = { ws, challenge };
-        this.challenges.set(challenge, challengeObj);
-        ws.send(JSON.stringify({ type: "challenge", challenge }));
+    // Handle WebSocket connections from tunnel clients
+    handleTunnelClientConnection(ws, req) {
+        ws.on("message", (message) => {
+            try {
+                const data = JSON.parse(message);
+                console.log(`Received WebSocket message: ${data.type}`);
+                if (data.type === "start") {
+                    this.handleStartMessage(ws, data);
+                }
+                else if (data.type === "response") {
+                    this.handleResponseMessage(data);
+                }
+                else if (data.type === "sse") {
+                    this.handleSSEMessage(data);
+                }
+                else if (data.type === "websocket") {
+                    this.handleWebSocketMessage(data);
+                }
+            }
+            catch (error) {
+                console.error("Error processing message:", error);
+                ws.close(1008, "Invalid message");
+            }
+        });
+        ws.on("close", () => {
+            console.log("WebSocket connection closed");
+            this.removeTunnel(ws);
+        });
+        ws.on("error", (error) => {
+            console.error("WebSocket error:", error);
+        });
     }
-    handleStartMessage(ws, data) {
-        const { challenge, signature, tunnelId } = data;
-        const challengeObj = this.challenges.get(challenge);
-        if (!challengeObj || challengeObj.ws !== ws) {
-            ws.close(1008, "Invalid challenge");
+    // Handle incoming WebSocket connections to be proxied through the tunnel
+    handleProxiedWebSocketConnection(ws, req) {
+        var _a;
+        const url = (0, url_1.parse)(req.url || '', true);
+        const pathParts = ((_a = url.pathname) === null || _a === void 0 ? void 0 : _a.split('/')) || [];
+        if (pathParts.length < 2) {
+            ws.close(1008, "Invalid tunnel ID");
             return;
         }
-        this.challenges.delete(challenge);
-        const publicKey = bs58_1.default.decode(tunnelId);
-        if (!client_1.DainClientAuth.verifyMessage(challenge, signature, publicKey)) {
-            ws.close(1008, "Invalid signature");
+        const tunnelId = pathParts[1];
+        const remainingPath = '/' + pathParts.slice(2).join('/');
+        const tunnel = this.tunnels.get(tunnelId);
+        console.log(`Handling WebSocket connection for tunnel: ${tunnelId}, path: ${remainingPath}`);
+        console.log(`Available tunnels: ${Array.from(this.tunnels.keys()).join(', ')}`);
+        if (!tunnel) {
+            console.log(`Tunnel not found for WebSocket connection: ${tunnelId}`);
+            ws.close(1008, "Tunnel not found");
             return;
         }
-        this.tunnels.set(tunnelId, { id: tunnelId, ws });
-        console.log(`Tunnel added: ${tunnelId}`);
-        console.log(`Current tunnels: ${Array.from(this.tunnels.keys()).join(', ')}`);
-        // Add a periodic check to ensure the tunnel is still in the map
-        const intervalId = setInterval(() => {
+        // Create a WebSocket connection ID
+        const wsConnectionId = (0, uuid_1.v4)();
+        // Store this connection
+        this.wsConnections.set(wsConnectionId, {
+            clientSocket: ws,
+            id: wsConnectionId,
+            path: remainingPath,
+            headers: req.headers
+        });
+        // Notify tunnel client about the new WebSocket connection
+        tunnel.ws.send(JSON.stringify({
+            type: "websocket_connection",
+            id: wsConnectionId,
+            path: remainingPath,
+            headers: req.headers
+        }));
+        // Handle messages from the client
+        ws.on("message", (data) => {
+            tunnel.ws.send(JSON.stringify({
+                type: "websocket",
+                id: wsConnectionId,
+                event: "message",
+                data: data.toString('base64')
+            }));
+        });
+        // Handle client disconnection
+        ws.on("close", () => {
+            tunnel.ws.send(JSON.stringify({
+                type: "websocket",
+                id: wsConnectionId,
+                event: "close"
+            }));
+            this.wsConnections.delete(wsConnectionId);
+        });
+        // Handle errors
+        ws.on("error", (error) => {
+            tunnel.ws.send(JSON.stringify({
+                type: "websocket",
+                id: wsConnectionId,
+                event: "error",
+                data: error.message
+            }));
+        });
+    }
+    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");
+            return;
+        }
+        // TODO: Validate API key with platform in production
+        // For now, accept any properly formatted API key (dev mode)
+        try {
+            // If tunnel already exists, remove old one
             if (this.tunnels.has(tunnelId)) {
-                console.log(`Tunnel ${tunnelId} still active`);
+                console.log(`Tunnel ${tunnelId} already exists, replacing it`);
+                const oldTunnel = this.tunnels.get(tunnelId);
+                if (oldTunnel && oldTunnel.ws !== ws) {
+                    oldTunnel.ws.close(1000, "Replaced by new connection");
+                }
             }
-            else {
-                console.log(`Tunnel ${tunnelId} not found in periodic check`);
+            this.tunnels.set(tunnelId, { id: tunnelId, ws });
+            console.log(`Tunnel added: ${tunnelId}`);
+            console.log(`Current tunnels: ${Array.from(this.tunnels.keys()).join(', ')}`);
+            // Save tunnel ID on the WebSocket object for easy lookup on close
+            ws.tunnelId = tunnelId;
+            // Add a periodic check to ensure the tunnel is still in the map
+            const intervalId = setInterval(() => {
+                if (this.tunnels.has(tunnelId)) {
+                    const tunnel = this.tunnels.get(tunnelId);
+                    if (tunnel && tunnel.ws === ws && ws.readyState === ws_1.default.OPEN) {
+                        console.log(`Tunnel ${tunnelId} still active`);
+                        // Send a ping to keep the connection alive
+                        try {
+                            ws.ping();
+                        }
+                        catch (error) {
+                            console.error(`Error sending ping to tunnel ${tunnelId}:`, error);
+                            clearInterval(intervalId);
+                        }
+                    }
+                    else {
+                        console.log(`Tunnel ${tunnelId} exists but WebSocket state is invalid, cleaning up`);
+                        clearInterval(intervalId);
+                        if (tunnel && tunnel.ws === ws) {
+                            this.tunnels.delete(tunnelId);
+                        }
+                    }
+                }
+                else {
+                    console.log(`Tunnel ${tunnelId} not found in periodic check`);
+                    clearInterval(intervalId);
+                }
+            }, 5000); // Check every 5 seconds
+            // Store the interval ID on the WebSocket object so we can clean it up
+            ws.keepAliveInterval = intervalId;
+            ws.on("close", () => {
+                console.log(`WebSocket for tunnel ${tunnelId} closed, clearing interval`);
                 clearInterval(intervalId);
+            });
+            let tunnelUrl = `${this.hostname}`;
+            if (process.env.SKIP_PORT !== "true") {
+                tunnelUrl += `:${this.port}`;
             }
-        }, 5000); // Check every 5 seconds
-        ws.on("close", () => {
-            clearInterval(intervalId);
-        });
-        let tunnelUrl = `${this.hostname}`;
-        if (process.env.SKIP_PORT !== "true") {
-            tunnelUrl += `:${this.port}`;
+            tunnelUrl += `/${tunnelId}`;
+            ws.send(JSON.stringify({ type: "tunnelUrl", url: tunnelUrl }));
+            console.log(`New tunnel created: ${tunnelUrl}`);
+        }
+        catch (error) {
+            console.error(`Error in handleStartMessage for tunnel ${tunnelId}:`, error);
+            ws.close(1011, "Internal server error");
         }
-        tunnelUrl += `/${tunnelId}`;
-        ws.send(JSON.stringify({ type: "tunnelUrl", url: tunnelUrl }));
-        console.log(`New tunnel created: ${tunnelUrl}`);
     }
     handleResponseMessage(data) {
         const pendingRequest = this.pendingRequests.get(data.requestId);
@@ -133,8 +259,61 @@ class DainTunnelServer {
             this.pendingRequests.delete(data.requestId);
         }
     }
-    async handleHttpRequest(req, res) {
+    handleSSEMessage(data) {
+        const connection = this.sseConnections.get(data.id);
+        if (!connection) {
+            console.log(`SSE connection not found: ${data.id}`);
+            return;
+        }
+        const { res } = connection;
+        if (data.event) {
+            res.write(`event: ${data.event}\n`);
+        }
+        // Split data by newlines and send each line
+        const dataLines = data.data.split('\n');
+        for (const line of dataLines) {
+            res.write(`data: ${line}\n`);
+        }
+        res.write('\n');
+        // Check if this is a "close" event
+        if (data.event === 'close') {
+            res.end();
+            this.sseConnections.delete(data.id);
+        }
+    }
+    handleWebSocketMessage(data) {
+        const connection = this.wsConnections.get(data.id);
+        if (!connection) {
+            console.log(`WebSocket connection not found: ${data.id}`);
+            return;
+        }
+        const { clientSocket } = connection;
+        if (data.event === 'message' && data.data) {
+            const messageData = Buffer.from(data.data, 'base64');
+            if (clientSocket.readyState === ws_1.default.OPEN) {
+                clientSocket.send(messageData);
+            }
+        }
+        else if (data.event === 'close') {
+            if (clientSocket.readyState === ws_1.default.OPEN) {
+                clientSocket.close();
+            }
+            this.wsConnections.delete(data.id);
+        }
+        else if (data.event === 'error' && data.data) {
+            if (clientSocket.readyState === ws_1.default.OPEN) {
+                clientSocket.close(1011, data.data);
+            }
+            this.wsConnections.delete(data.id);
+        }
+    }
+    async handleRequest(req, res) {
         const tunnelId = req.params.tunnelId;
+        // Check for upgraded connections (WebSockets) - these are handled by the WebSocket server
+        if (req.headers.upgrade && req.headers.upgrade.toLowerCase() === 'websocket') {
+            // This is handled by the WebSocket server now
+            return;
+        }
         let tunnel;
         let retries = 3;
         while (retries > 0 && !tunnel) {
@@ -150,6 +329,11 @@ class DainTunnelServer {
             console.log(`Tunnel not found after retries: ${tunnelId}`);
             return res.status(404).send("Tunnel not found");
         }
+        // Check for SSE request
+        if (req.headers.accept && req.headers.accept.includes('text/event-stream')) {
+            return this.handleSSERequest(req, res, tunnelId, tunnel);
+        }
+        // Handle regular HTTP request
         const requestId = (0, uuid_1.v4)();
         const startTime = Date.now();
         this.pendingRequests.set(requestId, { res, startTime });
@@ -166,29 +350,102 @@ class DainTunnelServer {
         tunnel.ws.send(JSON.stringify(requestMessage));
         console.log(`Request forwarded: ${requestId}, Method: ${req.method}, Path: ${req.url}`);
     }
+    handleSSERequest(req, res, tunnelId, tunnel) {
+        // Setup SSE connection
+        const sseId = (0, uuid_1.v4)();
+        res.writeHead(200, {
+            'Content-Type': 'text/event-stream',
+            'Cache-Control': 'no-cache',
+            'Connection': 'keep-alive'
+        });
+        // Send initial connection event
+        res.write('\n');
+        // Store the SSE connection
+        this.sseConnections.set(sseId, {
+            req,
+            res,
+            id: sseId,
+            tunnelId
+        });
+        // Notify the tunnel client about the new SSE connection
+        tunnel.ws.send(JSON.stringify({
+            type: "sse_connection",
+            id: sseId,
+            path: req.url,
+            headers: req.headers
+        }));
+        // Handle client disconnect
+        req.on('close', () => {
+            tunnel.ws.send(JSON.stringify({
+                type: "sse_close",
+                id: sseId
+            }));
+            this.sseConnections.delete(sseId);
+        });
+    }
     removeTunnel(ws) {
-        let removedTunnelId;
-        for (const [id, tunnel] of this.tunnels.entries()) {
-            if (tunnel.ws === ws) {
-                this.tunnels.delete(id);
-                removedTunnelId = id;
-                console.log(`Tunnel removed: ${id}`);
-                break;
+        try {
+            // Clear any interval timer
+            if (ws.keepAliveInterval) {
+                clearInterval(ws.keepAliveInterval);
             }
-        }
-        if (removedTunnelId) {
-            console.log(`Tunnel ${removedTunnelId} removed. Current tunnels: ${Array.from(this.tunnels.keys()).join(', ')}`);
-        }
-        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);
-                break;
+            // If we saved tunnelId on the WebSocket object, use it for faster lookup
+            const tunnelId = ws.tunnelId;
+            let removedTunnelId = tunnelId;
+            if (tunnelId && this.tunnels.has(tunnelId)) {
+                const tunnel = this.tunnels.get(tunnelId);
+                if (tunnel && tunnel.ws === ws) {
+                    this.tunnels.delete(tunnelId);
+                    console.log(`Tunnel removed using stored ID: ${tunnelId}`);
+                }
+            }
+            else {
+                // Fall back to iterating through all tunnels
+                removedTunnelId = undefined;
+                for (const [id, tunnel] of this.tunnels.entries()) {
+                    if (tunnel.ws === ws) {
+                        this.tunnels.delete(id);
+                        removedTunnelId = id;
+                        console.log(`Tunnel removed: ${id}`);
+                        break;
+                    }
+                }
+            }
+            if (removedTunnelId) {
+                console.log(`Tunnel ${removedTunnelId} removed. Current tunnels: ${Array.from(this.tunnels.keys()).join(', ')}`);
+                // Close all SSE connections associated with this tunnel
+                for (const [sseId, sseConnection] of this.sseConnections.entries()) {
+                    if (sseConnection.tunnelId === removedTunnelId) {
+                        try {
+                            sseConnection.res.end();
+                        }
+                        catch (error) {
+                            console.error(`Error closing SSE connection ${sseId}:`, error);
+                        }
+                        this.sseConnections.delete(sseId);
+                    }
+                }
+                // Close all WebSocket connections associated with this tunnel
+                for (const [wsId, wsConnection] of this.wsConnections.entries()) {
+                    const wsPath = wsConnection.path;
+                    if (wsPath.startsWith(`/${removedTunnelId}/`)) {
+                        try {
+                            wsConnection.clientSocket.close(1001, "Tunnel closed");
+                        }
+                        catch (error) {
+                            console.error(`Error closing WebSocket connection ${wsId}:`, error);
+                        }
+                        this.wsConnections.delete(wsId);
+                    }
+                }
+            }
+            else {
+                console.log(`No tunnel found to remove for the closed WebSocket connection`);
             }
         }
+        catch (error) {
+            console.error("Error in removeTunnel:", error);
+        }
     }
     async start() {
         return new Promise((resolve) => {
@@ -200,11 +457,39 @@ class DainTunnelServer {
     }
     async stop() {
         return new Promise((resolve) => {
-            this.wss.close(() => {
-                this.server.close(() => {
-                    resolve();
+            try {
+                // Close all SSE connections
+                for (const [sseId, sseConnection] of this.sseConnections.entries()) {
+                    try {
+                        sseConnection.res.end();
+                    }
+                    catch (error) {
+                        console.error(`Error closing SSE connection ${sseId}:`, error);
+                    }
+                    this.sseConnections.delete(sseId);
+                }
+                // Close all WebSocket connections
+                for (const [wsId, wsConnection] of this.wsConnections.entries()) {
+                    try {
+                        wsConnection.clientSocket.close(1001, "Server shutting down");
+                    }
+                    catch (error) {
+                        console.error(`Error closing WebSocket connection ${wsId}:`, error);
+                    }
+                    this.wsConnections.delete(wsId);
+                }
+                // Close the WebSocket server
+                this.wss.close(() => {
+                    // Close the HTTP server
+                    this.server.close(() => {
+                        resolve();
+                    });
                 });
-            });
+            }
+            catch (error) {
+                console.error('Error during server shutdown:', error);
+                resolve(); // Resolve anyway to prevent hanging
+            }
         });
     }
 }

package/dist/server/start.js

@@ -13,6 +13,7 @@ const server = new index_1.default(hostname, port);
 server.start()
     .then(() => {
     console.log(`DainTunnel Server started on http://${hostname}:${port}`);
+    console.log(`Supports HTTP, WebSockets, and Server-Sent Events (SSE)`);
 })
     .catch((error) => {
     console.error('Failed to start DainTunnel Server:', error);

package/package.json

@@ -1,13 +1,12 @@
 {
   "name": "@dainprotocol/tunnel",
-  "version": "1.0.4",
+  "version": "1.0.7",
   "description": "",
   "main": "dist/index.js",
   "private": false,
   "publishConfig": {
     "access": "public"
   },
-
   "scripts": {
     "build": "tsc",
     "build:types": "tsc --emitDeclarationOnly",
@@ -20,11 +19,11 @@
   "author": "Ryan",
   "license": "ISC",
   "dependencies": {
-    "@dainprotocol/service-sdk": "^1.0.11",
+    "@dainprotocol/service-sdk": "^1.2.5",
     "@types/body-parser": "^1.19.5",
     "@types/cors": "^2.8.17",
+    "@types/eventsource": "^3.0.0",
     "@types/express": "^4.17.21",
-    "@types/jest": "^29.5.12",
     "@types/node": "^22.5.4",
     "@types/uuid": "^10.0.0",
     "@types/ws": "^8.5.12",
@@ -32,21 +31,26 @@
     "bs58": "^6.0.0",
     "cors": "^2.8.5",
     "dotenv": "^16.4.5",
+    "eventsource": "^3.0.6",
     "express": "^4.19.2",
     "fetch-mock": "^11.1.3",
-    "jest": "^29.7.0",
-    "ts-jest": "^29.2.5",
     "ts-node": "^10.9.2",
-    "typescript": "^5.5.4",
     "uuid": "^10.0.0",
     "ws": "^8.18.0"
   },
+  "devDependencies": {
+    "@types/express": "^4.17.21",
+    "@types/jest": "^29.5.12",
+    "jest": "^29.7.0",
+    "ts-jest": "^29.1.5",
+    "typescript": "^5.0.0"
+  },
   "files": [
     "dist",
     "README.md"
   ],
   "engines": {
-    "node": "20.7.0"
+    "node": ">=20.7.0"
   },
   "exports": {
     ".": "./dist/index.js",