diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..f2a207d --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,15 @@ +name: Tests +on: [push, pull_request] +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + node: [20, 22, 24] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-node@v4 + with: + node-version: ${{ matrix.node }} + - run: npm install --ignore-scripts --no-package-lock + - run: npm test diff --git a/package.json b/package.json index a8ced48..81fd5e7 100644 --- a/package.json +++ b/package.json @@ -15,7 +15,8 @@ "scripts": { "build": "tsc", "prepublish": "tsc", - "start": "tsc --watch" + "start": "tsc --watch", + "test": "npm run build && node --test test/*.test.cjs" }, "files": [ "build/" diff --git a/src/agent-message.ts b/src/agent-message.ts index afd24a8..a8705d1 100644 --- a/src/agent-message.ts +++ b/src/agent-message.ts @@ -168,10 +168,34 @@ function longToHex( buffer: Uint8Array ) { * Decode a message. */ export function decode( buffer: Uint8Array ): AgentMessage { + if ( buffer.length < Offset.Payload ) { + throw new Error( 'Truncated SSM message header' ); + } + const messageType = decodeText( buffer, Offset.MessageType, Offset.MessageType + Length.MessageType ).trimEnd() as MessageType; + // MGS publication controls may have zeroed header metadata. AWS skips + // payload validation for these messages; they carry no application data. + if ( messageType === MessageType.StartPublicationMessage || messageType === MessageType.PausePublicationMessage ) { + return { MessageType: messageType, SequenceNumber: 0, MessageId: decodeUuid( buffer, Offset.MessageId, Offset.MessageId + Length.MessageId ), + PayloadType: 0, Payload: '', PayloadBytes: new Uint8Array() }; + } const headerLength = decodeInt( buffer, Offset.HL, Offset.HL + Length.HL ); + if ( headerLength < Offset.PayloadLength || headerLength + Length.PayloadLength > buffer.length ) { + throw new Error( 'Invalid SSM header length' ); + } + const payloadLength = decodeInt( buffer, headerLength, headerLength + Length.PayloadLength ); + if ( payloadLength !== buffer.length - headerLength - Length.PayloadLength ) { + throw new Error( 'Invalid SSM payload length' ); + } const payloadBytes = buffer.slice( headerLength + Length.PayloadLength ); + const digest = sha256.update( payloadBytes ).digest(); + if ( payloadLength !== 0 && ! digest.every( ( byte, index ) => byte === buffer[ Offset.PayloadDigest + index ] ) ) { + throw new Error( 'Invalid SSM payload digest' ); + } + if ( ! Number.isSafeInteger( decodeInt( buffer, Offset.SequenceNumber, Offset.SequenceNumber + Length.SequenceNumber ) ) ) { + throw new Error( 'SSM sequence number exceeds JavaScript integer precision' ); + } return { - MessageType: decodeText( buffer, Offset.MessageType, Offset.MessageType + Length.MessageType ).trimEnd() as MessageType, + MessageType: messageType, SequenceNumber: decodeInt( buffer, Offset.SequenceNumber, Offset.SequenceNumber + Length.SequenceNumber ), MessageId: decodeUuid( buffer, Offset.MessageId, Offset.MessageId + Length.MessageId ), PayloadType: decodeInt( buffer, Offset.PayloadType, Offset.PayloadType + Length.PayloadType ), @@ -180,7 +204,8 @@ export function decode( buffer: Uint8Array ): AgentMessage { }; } -export type EncodeableAgentMessage = Omit & { +export type EncodeableAgentMessage = Omit & { + Payload: string | Uint8Array; // Packed control flags (see header notes above). Defaults to 0, which is // what the reference client sends for all data messages; it uses 3 for // acknowledgements. @@ -193,7 +218,7 @@ export type EncodeableAgentMessage = Omit; + keepaliveTimer?: ReturnType; + private didDisconnect = false; handlers: Handler[] = []; paused: boolean = false; // Terminal output is decoded with a single streaming decoder so that @@ -89,99 +99,107 @@ export default class AWSSSMSession { }; this.webSocket.onmessage = (ev: any) => { - const uint8 = new Uint8Array( ev.data ); - const message = decode( uint8 ); - - switch( message.MessageType ) { - case 'output_stream_data': { - if ( message.SequenceNumber === this.expectedSequenceNumber ) { - this.acknowledgeMessage( message.MessageType, message.MessageId, message.SequenceNumber ); - this.handleOutputMessage( message ); - this.expectedSequenceNumber++; - - // Process any buffered messages which are now in sequence. - let next; - while ( ( next = this.incomingMessageBuffer.get( this.expectedSequenceNumber ) ) !== undefined ) { - this.incomingMessageBuffer.delete( this.expectedSequenceNumber ); - this.handleOutputMessage( next ); + if ( this.connectionClosed ) return; + try { + const uint8 = new Uint8Array( ev.data ); + const message = decode( uint8 ); + + switch( message.MessageType ) { + case 'output_stream_data': { + if ( message.SequenceNumber === this.expectedSequenceNumber ) { + this.acknowledgeMessage( message.MessageType, message.MessageId, message.SequenceNumber ); + this.handleOutputMessage( message ); this.expectedSequenceNumber++; - } - } else if ( message.SequenceNumber > this.expectedSequenceNumber ) { - // Out-of-order message: acknowledge it and hold it until the - // gap before it is filled. - if ( this.incomingMessageBuffer.size < INCOMING_MESSAGE_BUFFER_CAPACITY ) { + + // Process any buffered messages which are now in sequence. + let next; + while ( ( next = this.incomingMessageBuffer.get( this.expectedSequenceNumber ) ) !== undefined ) { + this.incomingMessageBuffer.delete( this.expectedSequenceNumber ); + this.handleOutputMessage( next ); + this.expectedSequenceNumber++; + } + } else if ( message.SequenceNumber > this.expectedSequenceNumber ) { + // Out-of-order message: acknowledge it and hold it until the + // gap before it is filled. + if ( this.incomingMessageBuffer.size < INCOMING_MESSAGE_BUFFER_CAPACITY ) { + this.acknowledgeMessage( message.MessageType, message.MessageId, message.SequenceNumber ); + this.incomingMessageBuffer.set( message.SequenceNumber, message ); + } + } else { + // Already processed (e.g. a resend after our acknowledgement + // was lost): re-acknowledge, but don't re-emit. this.acknowledgeMessage( message.MessageType, message.MessageId, message.SequenceNumber ); - this.incomingMessageBuffer.set( message.SequenceNumber, message ); } - } else { - // Already processed (e.g. a resend after our acknowledgement - // was lost): re-acknowledge, but don't re-emit. - this.acknowledgeMessage( message.MessageType, message.MessageId, message.SequenceNumber ); + break } - break - } - case 'start_publication': { - // Note: the sequence number deliberately continues from where it - // left off. It must never rewind: an acknowledgement for sequence - // number N means the remote agent has received N, so re-using an - // acknowledged number would be discarded as a duplicate. The - // reference client keeps its counter monotonic even across full - // websocket reconnects. - this.paused = false; - this.emit( 'resume', '' ); - break; - } - case 'pause_publication': { - this.paused = true; - this.emit( 'pause', '' ); - break; - } - case 'acknowledge': { - if ( message.Payload ) { - const payloadData = JSON.parse( message.Payload ); - const acknowledged = payloadData.AcknowledgedMessageSequenceNumber; - if ( this.lastAcknowledgedSequenceNumber === undefined || this.lastAcknowledgedSequenceNumber < acknowledged ) { - this.lastAcknowledgedSequenceNumber = acknowledged; + case 'start_publication': { + // Note: the sequence number deliberately continues from where it + // left off. It must never rewind: an acknowledgement for sequence + // number N means the remote agent has received N, so re-using an + // acknowledged number would be discarded as a duplicate. The + // reference client keeps its counter monotonic even across full + // websocket reconnects. + this.paused = false; + this.flushPendingMessages(); + this.emit( 'resume', '' ); + break; + } + case 'pause_publication': { + this.paused = true; + this.emit( 'pause', '' ); + break; + } + case 'acknowledge': { + if ( message.Payload ) { + const payloadData = JSON.parse( message.Payload ); + const acknowledged = payloadData.AcknowledgedMessageSequenceNumber; + if ( this.lastAcknowledgedSequenceNumber === undefined || this.lastAcknowledgedSequenceNumber < acknowledged ) { + this.lastAcknowledgedSequenceNumber = acknowledged; + } + + // The message has been received; no need to keep it around + // for retransmission. + this.outgoingMessageBuffer = this.outgoingMessageBuffer.filter( m => m.sequenceNumber !== acknowledged ); } - - // The message has been received; no need to keep it around - // for retransmission. - this.outgoingMessageBuffer = this.outgoingMessageBuffer.filter( m => m.sequenceNumber !== acknowledged ); + break; } - break; - } - case 'input_stream_data': { - break; - } - case 'channel_closed': { - if ( message.Payload ) { - const payloadData = JSON.parse( message.Payload ); - if ( payloadData.Output ) { - this.emit( 'output', payloadData.Output + '\r\n' ); + case 'input_stream_data': { + break; + } + case 'channel_closed': { + if ( message.Payload ) { + const payloadData = JSON.parse( message.Payload ); + if ( payloadData.Output ) { + this.emit( 'output', payloadData.Output + '\r\n' ); + } } + this.connectionTerminated = true; + this.close(); + break; + } + default: { + console.error( `Unhandled message type ${ message.MessageType }.` ); } - this.connectionTerminated = true; - break; - } - default: { - console.error( `Unhandled message type ${ message.MessageType }.` ); } + } catch { + this.fail( 'Invalid SSM message or session handler failure.' ); } }; - this.webSocket.onclose = (ev: any) => { - this.connectionClosed = true; - if ( this.resendTimer !== undefined ) { - clearInterval( this.resendTimer ); - this.resendTimer = undefined; - } - - this.emit( 'disconnect', ev.reason ); - }; + this.webSocket.onclose = (ev: any) => this.disconnected( ev.reason || '' ); + this.webSocket.onerror = () => this.fail( 'SSM WebSocket connection failed.' ); // Retransmit unacknowledged messages, like the reference client's // ResendStreamDataMessageScheduler. - this.resendTimer = setInterval( () => this.resendUnacknowledgedMessage(), RESEND_CHECK_INTERVAL ); + this.resendTimer = setInterval( () => { + try { this.resendUnacknowledgedMessage(); } + catch { this.fail( 'Could not resend SSM data.' ); } + }, RESEND_CHECK_INTERVAL ); + this.keepaliveTimer = setInterval( () => { + if ( this.webSocket.readyState !== 1 ) return; + try { this.ping(); } + catch { this.fail( 'SSM keepalive failed.' ); } + }, 5 * 60 * 1000 ); } // Resend the oldest unacknowledged message if its acknowledgement is @@ -189,55 +207,61 @@ export default class AWSSSMSession { // messages below its expected sequence number. private resendUnacknowledgedMessage() { const head = this.outgoingMessageBuffer[ 0 ]; - if ( ! head || this.webSocket.readyState !== 1 ) { + if ( ! head || this.paused || this.webSocket.readyState !== 1 ) { return; } if ( Date.now() - head.lastSentTime < RETRANSMISSION_TIMEOUT ) { return; } + if ( head.resendAttempts >= 300 ) { + this.fail( 'SSM acknowledgement timed out.' ); + return; + } + head.resendAttempts++; head.lastSentTime = Date.now(); + head.sent = true; this.webSocket.send( head.data ); } send( message: string | ArrayBuffer | SharedArrayBuffer | ArrayBufferView | Blob ) { if ( this.webSocket.readyState !== 1 ) { - console.warn( 'WebSocket is not yet ready.' ); - return; + throw new Error( 'SSM WebSocket is not open' ); + } + + if ( this.paused && ! ( message instanceof Uint8Array ) ) { + throw new Error( 'Cannot queue an unencoded SSM message while publication is paused' ); } // Keep sequential messages around until they're acknowledged, so they // can be retransmitted if the acknowledgement doesn't arrive. if ( message instanceof Uint8Array ) { + if ( this.outgoingMessageBuffer.length >= OUTGOING_MESSAGE_BUFFER_CAPACITY ) { + throw new Error( 'SSM outgoing buffer is full; refusing to drop unacknowledged data' ); + } this.outgoingMessageBuffer.push( { sequenceNumber: this.outgoingSequenceNumber, data: message, - lastSentTime: Date.now(), + lastSentTime: this.paused ? 0 : Date.now(), + sent: ! this.paused, + resendAttempts: 0, } ); - if ( this.outgoingMessageBuffer.length > OUTGOING_MESSAGE_BUFFER_CAPACITY ) { - this.outgoingMessageBuffer.shift(); - } } this.outgoingSequenceNumber++; - this.webSocket.send( message ) + if ( ! this.paused ) this.webSocket.send( message ); } - on( eventName: 'output' | 'connect' | 'disconnect' | 'pause' | 'resume', callback: Function ) { + on( eventName: Handler['event'], callback: Function ) { this.handlers.push( { event: eventName, callback, } ); + return this; } write( message: string ) { - const data = { - MessageType: MessageType.InputStreamDataMessage, - SequenceNumber: this.outgoingSequenceNumber, - Payload: message, - PayloadType: PayloadType.Output, - }; - this.send( encode( data ) ); + this.writeRaw( new TextEncoder().encode( message ) ); } ping() { @@ -253,7 +277,7 @@ export default class AWSSSMSession { } } - emit( eventName: 'output' | 'connect' | 'disconnect' | 'pause' | 'resume', data: string | ArrayBuffer | SharedArrayBuffer | ArrayBufferView | Blob ) { + emit( eventName: Handler['event'], data: unknown ) { this.handlers.filter( h => h.event === eventName ).map( h => h.callback( data ) ); } @@ -268,7 +292,79 @@ export default class AWSSSMSession { } close() { - this.webSocket.close(); + this.clearTimers(); + if ( this.webSocket.readyState === 0 || this.webSocket.readyState === 1 ) this.webSocket.close(); + } + + private clearTimers() { + if ( this.resendTimer !== undefined ) clearInterval( this.resendTimer ); + if ( this.keepaliveTimer !== undefined ) clearInterval( this.keepaliveTimer ); + this.resendTimer = undefined; + this.keepaliveTimer = undefined; + } + + private disconnected( reason: string ) { + this.connectionClosed = true; + this.clearTimers(); + this.incomingMessageBuffer.clear(); + this.outgoingMessageBuffer = []; + if ( ! this.didDisconnect ) { + this.didDisconnect = true; + this.emit( 'disconnect', reason ); + } + } + + private fail( reason: string ) { + this.disconnected( reason ); + this.close(); + } + + private flushPendingMessages() { + for ( const message of this.outgoingMessageBuffer ) { + if ( this.paused || this.webSocket.readyState !== 1 ) return; + if ( ! message.sent ) { + message.sent = true; + message.lastSentTime = Date.now(); + this.webSocket.send( message.data ); + } + } + } + + /** Write raw bytes without UTF-8 conversion or newline rewriting. */ + writeRaw( bytes: Uint8Array ) { + const count = Math.ceil( bytes.byteLength / 1024 ); + if ( this.outgoingMessageBuffer.length + count > OUTGOING_MESSAGE_BUFFER_CAPACITY ) { + throw new Error( 'SSM outgoing buffer is full; refusing to partially write binary data' ); + } + for ( let start = 0; start < bytes.byteLength; start += 1024 ) { + this.send( encode( { + MessageType: MessageType.InputStreamDataMessage, + SequenceNumber: this.outgoingSequenceNumber, + Payload: bytes.subarray( start, start + 1024 ), + PayloadType: PayloadType.Output, + } ) ); + } + } + + sendHandshakeResponse( response: HandshakeResponsePayload ) { + this.send( encode( { + MessageType: MessageType.InputStreamDataMessage, + SequenceNumber: this.outgoingSequenceNumber, + Payload: JSON.stringify( response ), + PayloadType: PayloadType.HandshakeResponse, + } ) ); + } + + /** Basic-port reconnect and termination controls, not database bytes. */ + sendFlag( flag: PayloadTypeFlag ) { + const bytes = new Uint8Array( 4 ); + new DataView( bytes.buffer ).setUint32( 0, flag, false ); + this.send( encode( { + MessageType: MessageType.InputStreamDataMessage, + SequenceNumber: this.outgoingSequenceNumber, + Payload: bytes, + PayloadType: PayloadType.Flag, + } ) ); } // Process an output_stream_data message according to its payload type. @@ -276,7 +372,11 @@ export default class AWSSSMSession { // session setup and must not be shown to the user. private handleOutputMessage( message: AgentMessage ) { switch ( message.PayloadType ) { - case PayloadType.Output: + case PayloadType.Output: { + this.emit( 'rawOutput', message.PayloadBytes ); + this.emit( 'output', this.outputDecoder.decode( message.PayloadBytes, { stream: true } ) ); + break; + } case PayloadType.Error: case PayloadType.StdErr: { this.emit( 'output', this.outputDecoder.decode( message.PayloadBytes, { stream: true } ) ); @@ -288,11 +388,17 @@ export default class AWSSSMSession { } case PayloadType.HandshakeComplete: { const payloadData = JSON.parse( message.Payload ); - if ( payloadData.CustomerMessage ) { + this.emit( 'handshakeComplete', payloadData ); + if ( payloadData.CustomerMessage && ! this.handlers.some( h => h.event === 'handshakeComplete' ) ) { this.emit( 'output', payloadData.CustomerMessage + '\r\n' ); } break; } + case PayloadType.Flag: { + if ( message.PayloadBytes.length !== 4 ) throw new Error( 'Invalid SSM flag' ); + this.emit( 'flag', new DataView( message.PayloadBytes.buffer, message.PayloadBytes.byteOffset, 4 ).getUint32( 0, false ) ); + break; + } default: { console.error( `Unhandled payload type ${ message.PayloadType }.` ); } @@ -305,10 +411,15 @@ export default class AWSSSMSession { // KMSEncryption) is unsupported, which the agent treats as a failed // handshake, and it will cancel the session. private handleHandshakeRequest( message: AgentMessage ) { - const request = JSON.parse( message.Payload ); + const request = JSON.parse( message.Payload ) as HandshakeRequestPayload; + if ( ! request || ! Array.isArray( request.RequestedClientActions ) ) throw new Error( 'Invalid SSM handshake' ); + if ( this.handlers.some( h => h.event === 'handshake' ) ) { + this.emit( 'handshake', request ); + return; + } const errors: string[] = []; - const actions = ( request.RequestedClientActions || [] ).map( ( action: { ActionType: string } ) => { - if ( action.ActionType === 'SessionType' ) { + const actions = ( request.RequestedClientActions || [] ).map( action => { + if ( action.ActionType === 'SessionType' && [ 'Standard_Stream', 'InteractiveCommands', 'NonInteractiveCommands' ].includes( action.ActionParameters?.SessionType || '' ) ) { return { ActionType: action.ActionType, ActionStatus: ActionStatus.Success, @@ -329,12 +440,7 @@ export default class AWSSSMSession { ProcessedClientActions: actions, Errors: errors, }; - this.send( encode( { - MessageType: MessageType.InputStreamDataMessage, - SequenceNumber: this.outgoingSequenceNumber, - Payload: JSON.stringify( response ), - PayloadType: PayloadType.HandshakeResponse, - } ) ); + this.sendHandshakeResponse( response ); } private acknowledgeMessage( messageType: MessageType, messageId: string, messageSequenceNumber: number ) { diff --git a/src/channel-messages.ts b/src/channel-messages.ts index ec87d44..db4f8d5 100644 --- a/src/channel-messages.ts +++ b/src/channel-messages.ts @@ -45,3 +45,10 @@ export enum ActionStatus { Failed = 2, Unsupported = 3, } + +// Four-byte, big-endian values carried in payload type Flag. +export enum PayloadTypeFlag { + DisconnectToPort = 1, + TerminateSession = 2, + ConnectToPortError = 3, +} diff --git a/src/handshake.ts b/src/handshake.ts new file mode 100644 index 0000000..3ade7c8 --- /dev/null +++ b/src/handshake.ts @@ -0,0 +1,20 @@ +import { ActionStatus } from './channel-messages'; + +export interface HandshakeRequestPayload { + AgentVersion?: string; + RequestedClientActions: { + ActionType: string; + ActionParameters?: { SessionType?: string; Properties?: Record; [key: string]: unknown }; + }[]; +} + +export interface HandshakeResponsePayload { + ClientVersion: string; + ProcessedClientActions: { + ActionType: string; + ActionStatus: ActionStatus; + ActionResult?: unknown; + Error?: string; + }[]; + Errors: string[] | null; +} diff --git a/test/codec.test.cjs b/test/codec.test.cjs new file mode 100644 index 0000000..743748e --- /dev/null +++ b/test/codec.test.cjs @@ -0,0 +1,70 @@ +const assert = require('node:assert/strict'); +const crypto = require('node:crypto'); +const test = require('node:test'); +const { encode, decode } = require('../build/agent-message'); +const { MessageType, PayloadType } = require('../build/channel-messages'); + +function frame(payload, sequence = 0) { + // Independent construction from AWS's ClientMessage offsets, including the + // least-significant UUID half first. No use of this package's encoder. + const bytes = Buffer.alloc(120 + payload.length); + bytes.writeUInt32BE(116, 0); + bytes.write('output_stream_data', 4, 'ascii'); + bytes.writeUInt32BE(1, 36); + bytes.writeBigUInt64BE(1700000000000n, 40); + bytes.writeBigUInt64BE(BigInt(sequence), 48); + Buffer.from('8899aabbccddeeff0011223344556677', 'hex').copy(bytes, 64); + crypto.createHash('sha256').update(payload).digest().copy(bytes, 80); + bytes.writeUInt32BE(PayloadType.Output, 112); + bytes.writeUInt32BE(payload.length, 116); + payload.copy(bytes, 120); + return bytes; +} + +test('decodes the AWS wire layout with binary payloads and 64-bit sequence numbers', () => { + const payload = Buffer.from([0, 255, 128, 10, 0]); + const message = decode(frame(payload, 4294967297)); + assert.equal(message.MessageId, '00112233-4455-6677-8899-aabbccddeeff'); + assert.equal(message.SequenceNumber, 4294967297); + assert.deepEqual(Buffer.from(message.PayloadBytes), payload); +}); + +test('encodes correct byte lengths and SHA-256 for binary and Unicode text', () => { + for (const payload of ['é€\0', Uint8Array.from([0, 255, 128, 10])]) { + const bytes = Buffer.from(encode({ MessageType: MessageType.InputStreamDataMessage, + SequenceNumber: 4294967297, PayloadType: PayloadType.Output, Payload: payload })); + const expected = typeof payload === 'string' ? Buffer.from(payload) : Buffer.from(payload); + assert.equal(bytes.readUInt32BE(0), 116); + assert.equal(bytes.readUInt32BE(116), expected.length); + assert.equal(bytes.readBigUInt64BE(48), 4294967297n); + assert.deepEqual(bytes.subarray(120), expected); + assert.deepEqual(bytes.subarray(80, 112), crypto.createHash('sha256').update(expected).digest()); + assert.ok(bytes.readBigUInt64BE(40) > 2n ** 32n); + } +}); + +test('rejects truncated headers, invalid lengths, corrupted payloads, and imprecise sequences', () => { + const valid = frame(Buffer.from('secret')); + const shortHeader = valid.subarray(0, 119); + const badHeader = Buffer.from(valid); badHeader.writeUInt32BE(999, 0); + const badLength = Buffer.from(valid); badLength.writeUInt32BE(1, 116); + const corrupt = Buffer.from(valid); corrupt[120] ^= 1; + const unsafeSequence = Buffer.from(valid); unsafeSequence.writeBigUInt64BE(2n ** 53n, 48); + for (const bytes of [shortHeader, badHeader, badLength, corrupt, unsafeSequence]) { + assert.throws(() => decode(bytes), error => !error.message.includes('secret')); + } +}); + +test('accepts empty-payload frames without a digest, as the AWS validator does', () => { + const bytes = frame(Buffer.alloc(0)); bytes.fill(0, 80, 112); + assert.equal(decode(bytes).PayloadBytes.length, 0); +}); + +test('accepts zeroed publication-control metadata without treating it as application data', () => { + for (const type of ['start_publication', 'pause_publication']) { + const bytes = Buffer.alloc(120); + bytes.write(type, 4, 'ascii'); + assert.equal(decode(bytes).MessageType, type); + assert.equal(decode(bytes).PayloadBytes.length, 0); + } +}); diff --git a/test/session.test.cjs b/test/session.test.cjs new file mode 100644 index 0000000..3f1cbea --- /dev/null +++ b/test/session.test.cjs @@ -0,0 +1,264 @@ +const assert = require('node:assert/strict'); +const test = require('node:test'); + +class WebSocket { + constructor() { this.readyState = 0; this.sent = []; this.pings = []; } + send(data) { this.sent.push(typeof data === 'string' ? data : Uint8Array.from(data)); } + ping(data) { this.pings.push(data); } + open() { this.readyState = 1; this.onopen({}); } + receive(data) { this.onmessage({ data: Uint8Array.from(data).buffer }); } + close() { this.readyState = 3; this.onclose({ reason: '' }); } +} +const wsPath = require.resolve('isomorphic-ws'); +require.cache[wsPath] = { id: wsPath, filename: wsPath, loaded: true, exports: WebSocket }; +const Session = require('../build/aws-ssm-session').default; +const { encode, decode } = require('../build/agent-message'); +const { MessageType, PayloadType, PayloadTypeFlag, ActionStatus } = require('../build/channel-messages'); +function create(t, opened = true) { + const session = new Session('wss://example.test', 'test-session', 'test-token'); + t.after(() => session.close()); + if (opened) { session.webSocket.open(); session.webSocket.sent = []; } + return session; +} +function incoming(session, payload, type = PayloadType.Output, seq = session.expectedSequenceNumber, messageType = MessageType.OutputStreamDataMessage) { + session.webSocket.receive(encode({ MessageType: messageType, PayloadType: type, SequenceNumber: seq, + Payload: typeof payload === 'object' && !(payload instanceof Uint8Array) ? JSON.stringify(payload) : payload })); +} +function control(session, type) { + const bytes = new Uint8Array(120); bytes.set(Buffer.from(type), 4); session.webSocket.receive(bytes); +} +function sent(session, type = MessageType.InputStreamDataMessage) { + return session.webSocket.sent.filter(data => typeof data !== 'string').map(decode).filter(message => message.MessageType === type); +} +function ack(session, sequence) { + incoming(session, { AcknowledgedMessageSequenceNumber: sequence }, PayloadType.Output, 0, MessageType.AcknowledgeMessage); +} + +test('large binary writes are split into 1024-byte frames and preserve all bytes including lone LF', t => { + const session = create(t); + const data = Uint8Array.from({ length: 65539 }, (_, i) => i % 256); + session.writeRaw(data); + const messages = sent(session); + assert.equal(messages.length, 65); + assert.ok(messages.every(message => message.PayloadBytes.length <= 1024)); + assert.deepEqual(Buffer.concat(messages.map(message => Buffer.from(message.PayloadBytes))), Buffer.from(data)); + assert.deepEqual(messages.map(message => message.SequenceNumber), Array.from({ length: 65 }, (_, i) => i)); + session.writeRaw(Uint8Array.of(10)); + assert.deepEqual(Array.from(sent(session).at(-1).PayloadBytes), [10]); +}); + +test('binary output remains ordered and duplicate frames are acknowledged without redelivery', t => { + const session = create(t); + const output = []; + session.on('rawOutput', bytes => output.push(...bytes)); + incoming(session, Uint8Array.of(255), PayloadType.Output, 2); + incoming(session, Uint8Array.of(0), PayloadType.Output, 0); + incoming(session, Uint8Array.of(128), PayloadType.Output, 1); + incoming(session, Uint8Array.of(128), PayloadType.Output, 1); + assert.deepEqual(output, [0, 128, 255]); + assert.equal(sent(session, MessageType.AcknowledgeMessage).length, 4); +}); + +test('terminal output still preserves split UTF-8 and NUL bytes', t => { + const session = create(t); + let output = ''; + session.on('output', text => { output += text; }); + incoming(session, Uint8Array.of(0xe2)); + incoming(session, Uint8Array.of(0x82)); + incoming(session, Uint8Array.of(0xac, 0)); + assert.equal(output, '€\0'); +}); + +test('manual port handshake sends one response, exposes completion, and does not emit control JSON as raw data', t => { + const session = create(t); + const raw = []; + let complete; + session.on('rawOutput', data => raw.push(data)); + session.on('handshake', request => { + assert.equal(request.RequestedClientActions[0].ActionParameters.SessionType, 'Port'); + session.sendHandshakeResponse({ ClientVersion: '1.1.69', Errors: [], + ProcessedClientActions: [{ ActionType: 'SessionType', ActionStatus: ActionStatus.Success }] }); + }); + session.on('handshakeComplete', data => { complete = data; }); + incoming(session, { RequestedClientActions: [{ ActionType: 'SessionType', ActionParameters: { SessionType: 'Port' } }] }, PayloadType.HandshakeRequest); + incoming(session, { CustomerMessage: 'Ready' }, PayloadType.HandshakeComplete); + assert.equal(sent(session).filter(message => message.PayloadType === PayloadType.HandshakeResponse).length, 1); + assert.deepEqual(complete, { CustomerMessage: 'Ready' }); + assert.deepEqual(raw, []); +}); + +test('automatic terminal handshake remains available and rejects KMS encryption and unknown session types', t => { + const session = create(t); + incoming(session, { RequestedClientActions: [ + { ActionType: 'SessionType', ActionParameters: { SessionType: 'Standard_Stream' } }, + { ActionType: 'KMSEncryption', ActionParameters: {} }, + { ActionType: 'SessionType', ActionParameters: { SessionType: 'Unknown' } }, + ] }, PayloadType.HandshakeRequest); + const response = JSON.parse(sent(session)[0].Payload); + assert.deepEqual(response.ProcessedClientActions.map(action => action.ActionStatus), [1, 3, 3]); + assert.equal(response.Errors.length, 2); +}); + +test('a Port handshake without an explicit consumer is not silently accepted', t => { + const session = create(t); + incoming(session, { RequestedClientActions: [{ ActionType: 'SessionType', ActionParameters: { SessionType: 'Port' } }] }, PayloadType.HandshakeRequest); + assert.equal(JSON.parse(sent(session)[0].Payload).ProcessedClientActions[0].ActionStatus, ActionStatus.Unsupported); +}); + +test('publication pause queues data, resume flushes it in order, and sequence numbers never rewind', t => { + const session = create(t); + session.writeRaw(Uint8Array.of(1)); + control(session, 'pause_publication'); + session.writeRaw(Uint8Array.of(2)); session.writeRaw(Uint8Array.of(3)); + assert.equal(sent(session).length, 1); + control(session, 'start_publication'); + session.writeRaw(Uint8Array.of(4)); + assert.deepEqual(sent(session).map(message => message.SequenceNumber), [0, 1, 2, 3]); + assert.deepEqual(sent(session).map(message => message.PayloadBytes[0]), [1, 2, 3, 4]); +}); + +test('acknowledgements use sequence 0 and flags 3 without advancing the outgoing stream', t => { + const session = create(t); + incoming(session, 'x'); + const bytes = Buffer.from(session.webSocket.sent[0]); + assert.equal(bytes.readBigUInt64BE(48), 0n); + assert.equal(bytes.readBigUInt64BE(56), 3n); + assert.equal(session.outgoingSequenceNumber, 0); +}); + +test('unacknowledged binary data is retransmitted unchanged and stops after acknowledgement', t => { + const session = create(t); + session.writeRaw(Uint8Array.of(0, 255)); + session.outgoingMessageBuffer[0].lastSentTime = Date.now() - 2000; + session.resendUnacknowledgedMessage(); + assert.deepEqual(session.webSocket.sent[0], session.webSocket.sent[1]); + ack(session, 0); + session.resendUnacknowledgedMessage(); + assert.equal(session.webSocket.sent.length, 2); + assert.equal(session.outgoingMessageBuffer.length, 0); +}); + +test('publication pause also suspends retransmission', t => { + const session = create(t); + session.writeRaw(Uint8Array.of(1)); + session.outgoingMessageBuffer[0].lastSentTime = Date.now() - 2000; + control(session, 'pause_publication'); + session.resendUnacknowledgedMessage(); + assert.equal(sent(session).length, 1); +}); + +test('full outgoing buffers reject an entire binary write without dropping queued frames', t => { + const session = create(t); + const pending = { sequenceNumber: 0, data: Uint8Array.of(1), lastSentTime: 0, sent: true, resendAttempts: 0 }; + session.outgoingMessageBuffer = Array(9999).fill(pending); + assert.throws(() => session.writeRaw(new Uint8Array(1025)), /partially write/); + assert.equal(session.outgoingMessageBuffer.length, 9999); + assert.equal(session.webSocket.sent.length, 0); + assert.equal(session.outgoingSequenceNumber, 0); +}); + +test('write before WebSocket open fails instead of silently discarding data', t => { + const session = create(t, false); + assert.throws(() => session.writeRaw(Uint8Array.of(1)), /not open/); + assert.equal(session.outgoingSequenceNumber, 0); +}); + +test('WebSocket keepalives never enter the SSM application stream', t => { + const session = create(t); + session.ping(); + assert.deepEqual(session.webSocket.pings, ['keepalive']); + assert.equal(session.webSocket.sent.length, 0); + assert.equal(session.outgoingSequenceNumber, 0); +}); + +test('control flags use their own payload type and a 32-bit big-endian value', t => { + const session = create(t); + session.sendFlag(PayloadTypeFlag.DisconnectToPort); + session.sendFlag(PayloadTypeFlag.TerminateSession); + assert.deepEqual(sent(session).map(message => message.PayloadType), [10, 10]); + assert.deepEqual(sent(session).map(message => Array.from(message.PayloadBytes)), [[0, 0, 0, 1], [0, 0, 0, 2]]); + let flag; + session.on('flag', value => { flag = value; }); + incoming(session, Uint8Array.of(0, 0, 0, 3), PayloadType.Flag); + assert.equal(flag, PayloadTypeFlag.ConnectToPortError); +}); + +test('corrupted frames are not acknowledged or delivered and disconnect without leaking payloads', t => { + const session = create(t); + let reason; + session.on('disconnect', value => { reason = value; }); + session.on('rawOutput', () => assert.fail('Corrupted data must not be delivered')); + const frame = encode({ MessageType: MessageType.OutputStreamDataMessage, SequenceNumber: 0, + PayloadType: PayloadType.Output, Payload: 'secret' }); + frame[120] ^= 1; + session.webSocket.receive(frame); + assert.equal(session.webSocket.sent.length, 0); + assert.equal(session.connectionClosed, true); + assert.ok(!reason.includes('secret')); +}); + +test('WebSocket errors notify disconnect once and clear both timers', t => { + const session = create(t); + let disconnects = 0; + session.on('disconnect', () => { disconnects++; }); + session.webSocket.onerror({ message: 'test-token' }); + session.webSocket.onclose({ reason: '' }); + assert.equal(disconnects, 1); + assert.equal(session.resendTimer, undefined); + assert.equal(session.keepaliveTimer, undefined); +}); + +test('retransmission exhaustion fails the connection rather than retrying forever', t => { + const session = create(t); + let reason; + session.on('disconnect', value => { reason = value; }); + session.writeRaw(Uint8Array.of(1)); + session.outgoingMessageBuffer[0].lastSentTime = Date.now() - 2000; + session.outgoingMessageBuffer[0].resendAttempts = 300; + session.resendUnacknowledgedMessage(); + assert.match(reason, /timed out/); + assert.equal(session.connectionClosed, true); +}); + +test('channel_closed closes the WebSocket and clears the retransmission state', t => { + const session = create(t); + incoming(session, {}, PayloadType.Output, 0, MessageType.ChannelClosedMessage); + assert.equal(session.connectionTerminated, true); + assert.equal(session.connectionClosed, true); + assert.equal(session.keepaliveTimer, undefined); +}); + +test('automatic keepalive matches AWS five-minute interval and is cancelled on close', () => { + const originalSet = global.setInterval; + const originalClear = global.clearInterval; + const timers = []; + let session; + try { + global.setInterval = (callback, interval) => { + const timer = { callback, interval, cleared: false }; + timers.push(timer); + return timer; + }; + global.clearInterval = timer => { timer.cleared = true; }; + session = new Session('wss://example.test', 'test-session', 'test-token'); + session.webSocket.open(); + const keepalive = timers.find(timer => timer.interval === 300000); + assert.ok(keepalive); + keepalive.callback(); + assert.deepEqual(session.webSocket.pings, ['keepalive']); + session.close(); + assert.ok(timers.every(timer => timer.cleared)); + } finally { + session?.close(); + global.setInterval = originalSet; + global.clearInterval = originalClear; + } +}); + +test('large terminal input shares binary-safe framing and retains UTF-8 bytes', t => { + const session = create(t); + const text = '€\0'.repeat(1000); + session.write(text); + assert.deepEqual(Buffer.concat(sent(session).map(message => Buffer.from(message.PayloadBytes))), Buffer.from(text)); + assert.ok(sent(session).every(message => message.PayloadBytes.length <= 1024)); +}); diff --git a/test/socket.test.cjs b/test/socket.test.cjs new file mode 100644 index 0000000..f40dad1 --- /dev/null +++ b/test/socket.test.cjs @@ -0,0 +1,67 @@ +const assert = require('node:assert/strict'); +const { once } = require('node:events'); +const test = require('node:test'); +const { Server } = require('ws'); +const Session = require('../build/aws-ssm-session').default; +const { encode, decode } = require('../build/agent-message'); +const { MessageType, PayloadType } = require('../build/channel-messages'); + +test('real WebSocket handshake and binary round trip preserve bytes across frame boundaries', { timeout: 10000 }, async t => { + const server = new Server({ host: '127.0.0.1', port: 0 }); + let session; + t.after(async () => { + session?.close(); + for (const client of server.clients) client.terminate(); + await new Promise(resolve => server.close(resolve)); + }); + await once(server, 'listening'); + const expected = Buffer.from(Array.from({ length: 8195 }, (_, i) => i % 256)); + let responseCount = 0; + const complete = new Promise((resolve, reject) => { + server.on('connection', socket => { + let authenticated = false; + let sequence = 0; + socket.on('message', data => { + try { + if (!authenticated) { + assert.equal(JSON.parse(data.toString()).TokenValue, 'synthetic-token'); + authenticated = true; + socket.send(encode({ MessageType: MessageType.OutputStreamDataMessage, SequenceNumber: sequence++, PayloadType: PayloadType.HandshakeRequest, + Payload: JSON.stringify({ RequestedClientActions: [{ ActionType: 'SessionType', ActionParameters: { SessionType: 'Port', Properties: { type: 'LocalPortForwarding' } } }] }) })); + return; + } + const message = decode(new Uint8Array(data)); + if (message.MessageType === MessageType.AcknowledgeMessage) return; + socket.send(encode({ MessageType: MessageType.AcknowledgeMessage, SequenceNumber: 0, Flags: 3, + PayloadType: PayloadType.Output, Payload: JSON.stringify({ AcknowledgedMessageSequenceNumber: message.SequenceNumber }) })); + if (message.PayloadType === PayloadType.HandshakeResponse) { + responseCount++; + assert.equal(JSON.parse(message.Payload).ClientVersion, '1.1.69'); + socket.send(encode({ MessageType: MessageType.OutputStreamDataMessage, SequenceNumber: sequence++, PayloadType: PayloadType.HandshakeComplete, Payload: '{}' })); + } else { + assert.equal(message.PayloadType, PayloadType.Output); + assert.ok(message.PayloadBytes.length <= 1024); + socket.send(encode({ MessageType: MessageType.OutputStreamDataMessage, SequenceNumber: sequence++, PayloadType: PayloadType.Output, Payload: message.PayloadBytes })); + } + } catch (error) { reject(error); } + }); + }); + session = new Session(`ws://127.0.0.1:${server.address().port}`, 'synthetic-session', 'synthetic-token'); + const received = []; + let count = 0; + session.on('disconnect', reason => { if (count !== expected.length) reject(new Error(reason || 'Premature disconnect')); }); + session.on('handshake', () => session.sendHandshakeResponse({ ClientVersion: '1.1.69', Errors: [], ProcessedClientActions: [{ ActionType: 'SessionType', ActionStatus: 1 }] })); + session.on('handshakeComplete', () => session.writeRaw(expected)); + session.on('rawOutput', data => { + received.push(Buffer.from(data)); count += data.length; + if (count === expected.length) { + try { + assert.deepEqual(Buffer.concat(received), expected); + assert.equal(responseCount, 1); + resolve(); + } catch (error) { reject(error); } + } + }); + }); + await complete; +});