Keep-alive implementation in TS SDK

README.md

@@ -15,7 +15,7 @@ Dev tunnels allows developers to securely expose local web services to the Inter
 | Reconnection | ✅ | ✅ | ❌ | ❌ | ❌ |
 | SSH-level Reconnection | ✅ | ✅ | ❌ | ❌ | ❌ |
 | Automatic tunnel access token refresh | ✅ | ✅ | ❌ | ❌ | ❌ |
-| Ssh Keep-alive | ✅ | 🗓️ | ❌ | ❌ | ❌ |
+| Ssh Keep-alive | ✅ | ✅ | ❌ | ❌ | ❌ |
 
 ✅ - Supported  
 🚧 - In Progress  

ts/package.json

@@ -21,8 +21,8 @@
 		"build-pack-publish": "npm run build && npm run pack && npm run publish"
 	},
 	"dependencies": {
-		"@microsoft/dev-tunnels-ssh": "^3.12.5",
-		"@microsoft/dev-tunnels-ssh-tcp": "^3.12.5",
+		"@microsoft/dev-tunnels-ssh": "^3.12.12",
+		"@microsoft/dev-tunnels-ssh-tcp": "^3.12.12",
 		"await-semaphore": "^0.1.3",
 		"axios": "^1.8.4",
 		"buffer": "^5.2.1",

ts/src/connections/index.ts

@@ -17,6 +17,8 @@ export * from './tunnelRelayTunnelClient';
 export * from './tunnelRelayTunnelHost';
 export * from './tunnelConnection';
 export * from './tunnelConnectionBase';
+export * from './tunnelConnectionOptions';
+export * from './sshKeepAliveEventArgs';
 export * from './connectionStatus';
 export * from './connectionStatusChangedEventArgs';
 export { maxReconnectDelayMs } from './relayTunnelConnector';

ts/src/connections/package.json

@@ -20,8 +20,8 @@
 		"vscode-jsonrpc": "^4.0.0",
 		"@microsoft/dev-tunnels-contracts": "^1.3.0",
 		"@microsoft/dev-tunnels-management": "^1.3.0",
-		"@microsoft/dev-tunnels-ssh": "^3.12.5",
-		"@microsoft/dev-tunnels-ssh-tcp": "^3.12.5",
+		"@microsoft/dev-tunnels-ssh": "^3.12.12",
+		"@microsoft/dev-tunnels-ssh-tcp": "^3.12.12",
 		"uuid": "^3.3.3",
 		"await-semaphore": "^0.1.3",
 		"websocket": "^1.0.28",

ts/src/connections/sshKeepAliveEventArgs.ts

@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+/**
+ * Event args raised when an SSH keep-alive succeeds or fails.
+ */
+export class SshKeepAliveEventArgs {
+    /**
+     * The number of keep-alive messages that have been sent with the same state.
+     */
+    public readonly count: number;
+
+    public constructor(count: number) {
+        this.count = count;
+    }
+}

ts/src/connections/tunnelConnection.ts

@@ -7,6 +7,7 @@ import { ConnectionStatusChangedEventArgs } from './connectionStatusChangedEvent
 import { RefreshingTunnelAccessTokenEventArgs } from './refreshingTunnelAccessTokenEventArgs';
 import { RetryingTunnelConnectionEventArgs } from './retryingTunnelConnectionEventArgs';
 import { ForwardedPortConnectingEventArgs } from '@microsoft/dev-tunnels-ssh-tcp';
+import { SshKeepAliveEventArgs } from './sshKeepAliveEventArgs';
 
 /**
  * Tunnel connection.
@@ -45,6 +46,16 @@ export interface TunnelConnection {
      */
     readonly forwardedPortConnecting: Event<ForwardedPortConnectingEventArgs>;
 
+    /**
+     * Event raised when a keep-alive message response is not received.
+     */
+    readonly keepAliveFailed: Event<SshKeepAliveEventArgs>;
+
+    /**
+     * Event raised when a keep-alive message response is received.
+     */
+    readonly keepAliveSucceeded: Event<SshKeepAliveEventArgs>;
+
      /**
       * Disposes this tunnel session.
       */

ts/src/connections/tunnelConnectionBase.ts

@@ -11,6 +11,7 @@ import { RetryingTunnelConnectionEventArgs } from './retryingTunnelConnectionEve
 import { TunnelAccessTokenProperties } from '@microsoft/dev-tunnels-management';
 import { ForwardedPortConnectingEventArgs } from '@microsoft/dev-tunnels-ssh-tcp';
 import { TrackingEmitter } from './utils';
+import { SshKeepAliveEventArgs } from './sshKeepAliveEventArgs';
 
 /**
  * Tunnel connection base class.
@@ -27,6 +28,8 @@ export class TunnelConnectionBase implements TunnelConnection {
         new Emitter<RetryingTunnelConnectionEventArgs>();
     private readonly forwardedPortConnectingEmitter =
         new Emitter<ForwardedPortConnectingEventArgs>();
+    private readonly keepAliveFailedEmitter = new Emitter<SshKeepAliveEventArgs>();
+    private readonly keepAliveSucceededEmitter = new Emitter<SshKeepAliveEventArgs>();
 
     protected constructor(
         /**
@@ -117,10 +120,34 @@ export class TunnelConnectionBase implements TunnelConnection {
      */
     public readonly forwardedPortConnecting = this.forwardedPortConnectingEmitter.event;
 
+    /**
+     * Event raised when a keep-alive message response is not received.
+     */
+    public readonly keepAliveFailed = this.keepAliveFailedEmitter.event;
+
+    /**
+     * Event raised when a keep-alive message response is received.
+     */
+    public readonly keepAliveSucceeded = this.keepAliveSucceededEmitter.event;
+
     protected onForwardedPortConnecting(e: ForwardedPortConnectingEventArgs) {
         this.forwardedPortConnectingEmitter.fire(e);
     }
 
+    /**
+     * Raises the keep-alive failed event.
+     */
+    protected onKeepAliveFailed(count: number) {
+        this.keepAliveFailedEmitter.fire(new SshKeepAliveEventArgs(count));
+    }
+
+    /**
+     * Raises the keep-alive succeeded event.
+     */
+    protected onKeepAliveSucceeded(count: number) {
+        this.keepAliveSucceededEmitter.fire(new SshKeepAliveEventArgs(count));
+    }
+
     /**
      * Closes and disposes the tunnel session.
      */

ts/src/connections/tunnelConnectionOptions.ts

@@ -56,4 +56,16 @@ export interface TunnelConnectionOptions {
      * (most common). This option applies only to client connections.
      */
     hostId?: string;
+
+    /**
+     * Gets or sets the SSH keep-alive interval in seconds. Default is 0 (disabled).
+     * When set to a positive value, the client/host will send SSH keep-alive messages
+     * and raise keep-alive events with the number of consecutive successes or failures.
+     * 
+     * The keep-alive events are raised at the time of sending the next keep-alive request.
+     * For example, if the interval is set to 10 seconds, the first request is sent after
+     * 10 seconds of inactivity and waits for 10 more seconds to call the keep-alive
+     * callback before sending another keep-alive request.
+     */
+    keepAliveIntervalInSeconds?: number;
 }

ts/src/connections/tunnelConnectionSession.ts

@@ -52,7 +52,7 @@ import { v4 as uuidv4 } from 'uuid';
  * Tunnel connection session.
  */
 export class TunnelConnectionSession extends TunnelConnectionBase implements TunnelSession {
-    private connectionOptions?: TunnelConnectionOptions;
+    protected connectionOptions?: TunnelConnectionOptions;
     private connectedTunnel: Tunnel | null = null;
     private connector?: TunnelConnector;
     private reconnectPromise?: Promise<void>;

ts/src/connections/tunnelRelayTunnelClient.ts

@@ -261,6 +261,12 @@ export class TunnelRelayTunnelClient extends TunnelConnectionSession implements
                     // session is optional since it is already over a TLS websocket.
                     config.keyExchangeAlgorithms.splice(0, 0, SshAlgorithms.keyExchange.none);
                 }
+
+                // Configure keep-alive if requested
+                const keepAliveInterval = this.connectionOptions?.keepAliveIntervalInSeconds;
+                if (keepAliveInterval && keepAliveInterval > 0) {
+                    config.keepAliveTimeoutInSeconds = keepAliveInterval;
+                }
             });
             this.sshSession.trace = this.trace;
             this.sshSession.onReportProgress(
@@ -272,6 +278,9 @@ export class TunnelRelayTunnelClient extends TunnelConnectionSession implements
             this.sshSession.onDisconnected(this.onSshSessionDisconnected, this, this.sshSessionDisposables);
             this.sshSession.onRequest(this.onRequest, this, this.sshSessionDisposables);
 
+            this.sshSession.onKeepAliveFailed((count) => this.onKeepAliveFailed(count));
+            this.sshSession.onKeepAliveSucceeded((count) => this.onKeepAliveSucceeded(count));
+
             const pfs = this.sshSession.activateService(PortForwardingService);
             if (this.connectionProtocol === webSocketSubProtocolv2) {
                 pfs.messageFactory = this;

ts/src/connections/tunnelRelayTunnelHost.ts

@@ -471,6 +471,12 @@ export class TunnelRelayTunnelHost extends TunnelConnectionSession implements Tu
         const session = SshHelpers.createSshServerSession(this.reconnectableSessions, (config) => {
             config.protocolExtensions.push(SshProtocolExtensionNames.sessionReconnect);
             config.addService(PortForwardingService);
+
+            // Configure keep-alive if requested
+            const keepAliveInterval = this.connectionOptions?.keepAliveIntervalInSeconds;
+            if (keepAliveInterval && keepAliveInterval > 0) {
+                config.keepAliveTimeoutInSeconds = keepAliveInterval;
+            }
         });
         session.trace = this.trace;
         session.onReportProgress(
@@ -482,7 +488,7 @@ export class TunnelRelayTunnelHost extends TunnelConnectionSession implements Tu
         };
 
         const tcs = new PromiseCompletionSource<void>();
-        
+
         const authenticatingEventRegistration = session.onAuthenticating((e) => {
             this.onSshClientAuthenticating(e);
         });
@@ -504,6 +510,9 @@ export class TunnelRelayTunnelHost extends TunnelConnectionSession implements Tu
             tcs.resolve();
         });
 
+        session.onKeepAliveFailed((count) => this.onKeepAliveFailed(count));
+        session.onKeepAliveSucceeded((count) => this.onKeepAliveSucceeded(count));
+
         try {
             const nodeStream = new NodeStream(stream);
             await session.connect(nodeStream);

ts/test/tunnels-test/tunnelConnectionKeepAliveTests.ts

@@ -0,0 +1,100 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import * as assert from 'assert';
+import { suite, test, slow, timeout } from '@testdeck/mocha';
+import { MockTunnelManagementClient } from './mocks/mockTunnelManagementClient';
+import {
+    Tunnel,
+    TunnelConnectionMode,
+    TunnelRelayTunnelEndpoint,
+} from '@microsoft/dev-tunnels-contracts';
+import {
+    TunnelRelayTunnelClient,
+    TunnelRelayTunnelHost,
+    TunnelConnectionOptions,
+    SshKeepAliveEventArgs,
+} from '@microsoft/dev-tunnels-connections';
+import {
+    KeyPair,
+    SshAlgorithms,
+} from '@microsoft/dev-tunnels-ssh';
+
+@suite
+export class TunnelConnectionKeepAliveTests {
+    private hostKeys!: KeyPair;
+    private managementClient!: MockTunnelManagementClient;
+    private tunnel!: Tunnel;
+
+    public async before() {
+        this.hostKeys = await SshAlgorithms.publicKey.ecdsaSha2Nistp384!.generateKeyPair();
+        this.managementClient = new MockTunnelManagementClient();
+        this.tunnel = {
+            tunnelId: 'tunnel1',
+            name: 'Test Tunnel',
+            domain: 'localhost',
+            accessTokens: {},
+            endpoints: [
+                {
+                    id: 'endpoint1',
+                    connectionMode: TunnelConnectionMode.TunnelRelay,
+                    hostRelayUri: 'wss://localhost:8080/tunnel',
+                    clientRelayUri: 'wss://localhost:8080/tunnel',
+                } as TunnelRelayTunnelEndpoint,
+            ],
+        };
+    }
+
+    @test
+    public async testClientKeepAlive() {
+        const client = new TunnelRelayTunnelClient(this.managementClient);
+
+        let failedEvent: SshKeepAliveEventArgs | undefined;
+        let succeededEvent: SshKeepAliveEventArgs | undefined;
+
+        const failedSubscription = client.keepAliveFailed((e: SshKeepAliveEventArgs) => {
+            failedEvent = e;
+        });
+
+        const succeededSubscription = client.keepAliveSucceeded((e: SshKeepAliveEventArgs) => {
+            succeededEvent = e;
+        });
+
+        (client as any).onKeepAliveFailed(1);
+        (client as any).onKeepAliveSucceeded(1);
+
+        assert.strictEqual(failedEvent?.count, 1, 'Failed event should fire');
+        assert.strictEqual(succeededEvent?.count, 1, 'Succeeded event should fire');
+
+        failedSubscription.dispose();
+        succeededSubscription.dispose();
+        await client.dispose();
+    }
+
+    @test
+    public async testHostKeepAlive() {
+        const host = new TunnelRelayTunnelHost(this.managementClient);
+
+        let failedEvent: SshKeepAliveEventArgs | undefined;
+        let succeededEvent: SshKeepAliveEventArgs | undefined;
+
+        const failedSubscription = host.keepAliveFailed((e: SshKeepAliveEventArgs) => {
+            failedEvent = e;
+        });
+
+        const succeededSubscription = host.keepAliveSucceeded((e: SshKeepAliveEventArgs) => {
+            succeededEvent = e;
+        });
+
+        (host as any).onKeepAliveFailed(2);
+        (host as any).onKeepAliveSucceeded(2);
+
+        assert.strictEqual(failedEvent?.count, 2, 'Failed event should fire');
+        assert.strictEqual(succeededEvent?.count, 2, 'Succeeded event should fire');
+
+
+        failedSubscription.dispose();
+        succeededSubscription.dispose();
+        await host.dispose();
+    }
+}