Add keepalive for agents

Dockerfile

@@ -4,7 +4,7 @@ RUN mkdir /app
 WORKDIR /app
 
 COPY . /app
-RUN apk add --update python gcc g++ make && \
+RUN apk add --update python3 gcc g++ make && \
     yarn install --frozen-lockfile && \
     yarn dist && \
     yarn install --frozen-lockfile --prod

src/__tests__/server.test.ts

@@ -50,6 +50,7 @@ describe("TunnelServer", () => {
 
   beforeEach(async () => {
     server = new TunnelServer();
+    server.disableKeepAlive();
     await server.start(port, "", idpPublicKey, tunnelAddress);
   });
 
@@ -124,7 +125,7 @@ describe("TunnelServer", () => {
 
       const agents = server.getAgentsForClusterId("a026e50d-f9b4-4aa8-ba02-c9722f7f0663");
 
-      agents.push(new Agent(ws as any, "rsa-public-key", server, "test-id"));
+      agents.push(new Agent({ socket: ws as any, publicKey: "rsa-public-key", server, clusterId: "test-id", keepalive: 0 }));
 
       const res = await get("/client/public-key", { "Authorization": `Bearer ${jwtToken}`});
 

src/agent.ts

@@ -8,6 +8,14 @@ export interface Client {
   userId: string;
 }
 
+export interface AgentOpts {
+  socket: WebSocket;
+  publicKey: string;
+  server: TunnelServer;
+  clusterId: string,
+  keepalive: number
+}
+
 export class Agent {
   public socket: WebSocket;
   public publicKey: string;
@@ -16,7 +24,9 @@ export class Agent {
   private mplex: BoredMplexClient;
   private server: TunnelServer;
 
-  constructor(socket: WebSocket, publicKey: string, server: TunnelServer, clusterId: string) {
+  constructor(opts: AgentOpts) {
+    const { socket, publicKey, server, clusterId, keepalive } = opts;
+
     this.socket = socket;
     this.publicKey = publicKey;
     this.server = server;
@@ -25,9 +35,15 @@ export class Agent {
     const stream = WebSocket.createWebSocketStream(this.socket);
 
     this.mplex = new BoredMplexClient();
+
+    if (keepalive) {
+      this.mplex.enableKeepAlive(keepalive);
+    }
+
     this.mplex.pipe(stream).pipe(this.mplex);
 
     this.socket.on("close", () => {
+      this.mplex.disableKeepAlive();
       this.clients.forEach((client) => this.removeClient(client.socket));
     });
 

src/request-handlers/agent-socket.ts

@@ -50,7 +50,7 @@ export function handleAgentSocket(req: IncomingMessage, socket: WebSocket, serve
 
   console.log(`SERVER: agent connected. Cluster id: ${clusterId}`);
   const publicKey = Buffer.from(req.headers["x-bored-publickey"]?.toString() || "", "base64").toString("utf-8");
-  const agent = new Agent(socket, publicKey, server, clusterId);
+  const agent = new Agent({ socket, publicKey, server, clusterId, keepalive: server.keepAlive });
 
   agents.push(agent);
 

src/server.ts

@@ -20,6 +20,7 @@ export class TunnelServer extends EventEmitter {
   public tunnelAddress?: string;
   public agents: Map<ClusterId, Agent[]> = new Map();
   public presenceSockets: Map<ClusterId, WebSocket[]> = new Map();
+  public keepAlive = 10_000;
 
   constructor() {
     super();
@@ -37,6 +38,14 @@ export class TunnelServer extends EventEmitter {
     });
   }
 
+  enableKeepAlive(keepAliveTimeMs: number) {
+    this.keepAlive = keepAliveTimeMs;
+  }
+
+  disableKeepAlive() {
+    this.keepAlive = 0;
+  }
+
   start(port = 8080, agentToken: string, idpPublicKey: string, tunnelAddress = process.env.TUNNEL_ADDRESS || ""): Promise<void> {
     this.agentToken = agentToken;
     this.idpPublicKey = idpPublicKey;