diff --git a/src/agent-message.ts b/src/agent-message.ts index 0c88aec..afd24a8 100644 --- a/src/agent-message.ts +++ b/src/agent-message.ts @@ -19,7 +19,11 @@ export type AgentMessage = { MessageId: string; // PayloadDigest: string; PayloadType: PayloadType; + // The payload decoded as UTF-8 text; JSON payloads are parsed from this. Payload: string; + // The raw payload bytes. Use these for binary-safe handling, or for + // streaming text decoding across message boundaries. + PayloadBytes: Uint8Array; } // HL - HeaderLength is a 4 byte integer that represents the header length. @@ -66,7 +70,11 @@ enum Offset { } /** - * Decode a UTF-8 byte sequence to a string. + * Decode a padded fixed-width header field to a string. + * + * Header fields may be padded with NULs or spaces depending on the sender, + * so strip NULs here (and trim whitespace at the call site). Not for use on + * the payload, where every byte is meaningful. */ function decodeText( buffer: Uint8Array, start: number, end?: number ) { const decoder = new TextDecoder( 'utf8' ); @@ -92,9 +100,11 @@ function putString( buffer: Uint8Array, start: number, end: number, value: strin * Decode a byte sequence to an integer. */ function decodeInt( buffer: Uint8Array, start: number, end: number ) { + // Multiply rather than shift; shifts operate on 32-bit integers, which + // would wrap when decoding 8-byte fields. Exact for values below 2^53. const bytes = buffer.slice( start, end ); return bytes.reduce( ( acc, byte ) => { - return ( acc << 8 ) + byte; + return ( acc * 0x100 ) + byte; }, 0 ) } @@ -115,11 +125,12 @@ function putInt( buffer: Uint8Array, pos: number, value: number ) { /** * Encode a (64-bit) long into a byte sequence. * - * We cheat here; JS doesn't support 64-bit integers (aka longs), so we can - * just use putInt with a 4 byte offset. + * JS numbers are doubles, so this is exact for values up to 2^53 - 1; that + * comfortably covers epoch milliseconds and sequence numbers. */ function putLong( buffer: Uint8Array, pos: number, value: number ) { - putInt( buffer, pos + 4, value ); + putInt( buffer, pos, Math.floor( value / 0x100000000 ) ); + putInt( buffer, pos + 4, value >>> 0 ); } /** @@ -158,23 +169,33 @@ function longToHex( buffer: Uint8Array ) { */ export function decode( buffer: Uint8Array ): AgentMessage { const headerLength = decodeInt( buffer, Offset.HL, Offset.HL + Length.HL ); + const payloadBytes = buffer.slice( headerLength + Length.PayloadLength ); return { MessageType: decodeText( buffer, Offset.MessageType, Offset.MessageType + Length.MessageType ).trimEnd() as 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 ), - Payload: decodeText( buffer, headerLength + Length.PayloadLength ), + Payload: new TextDecoder( 'utf8' ).decode( payloadBytes ), + PayloadBytes: payloadBytes, }; } -export type EncodeableAgentMessage = Omit; +export type EncodeableAgentMessage = Omit & { + // 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. + Flags?: number; +}; /** * Encode a message. */ export function encode( message: EncodeableAgentMessage ) { - const payloadLength = message.Payload.length; - const totalLength = Offset.PayloadLength + payloadLength + 4; + // Encode the payload to UTF-8 first; the length, digest, and bytes on the + // wire must all be computed from the same byte sequence. + const payload = new TextEncoder().encode( message.Payload ); + const payloadLength = payload.length; + const totalLength = Offset.Payload + payloadLength; // | HL| MessageType |Ver| CD | Seq | Flags | // | MessageId | Digest |PayType| PayLen| @@ -186,18 +207,22 @@ export function encode( message: EncodeableAgentMessage ) { putInt( data, Offset.SchemaVersion, 1 ); putLong( data, Offset.CreatedDate, Date.now() ); putLong( data, Offset.SequenceNumber, message.SequenceNumber ); - putLong( data, Offset.Flags, message.SequenceNumber > 0 ? 0b00 : 0b01 ); + putLong( data, Offset.Flags, message.Flags ?? 0 ); - // Generate a UUID. - uuidv4( {}, data, Offset.MessageId ); + // Generate a UUID. The wire format stores the least-significant 8 bytes + // first, then the most-significant 8 (see decodeUuid), so flip the halves. + const messageId = new Uint8Array( 16 ); + uuidv4( {}, messageId, 0 ); + data.set( messageId.subarray( 8, 16 ), Offset.MessageId ); + data.set( messageId.subarray( 0, 8 ), Offset.MessageId + 8 ); // Generate the digest. - data.set( sha256.update( message.Payload ).digest(), Offset.PayloadDigest ); + data.set( sha256.update( payload ).digest(), Offset.PayloadDigest ); // Set the payload. putInt( data, Offset.PayloadType, message.PayloadType ); putInt( data, Offset.PayloadLength, payloadLength ); - putString( data, Offset.Payload, totalLength, message.Payload ); + data.set( payload, Offset.Payload ); return data; } diff --git a/src/aws-ssm-session.ts b/src/aws-ssm-session.ts index 33bcc81..bbc2bf5 100644 --- a/src/aws-ssm-session.ts +++ b/src/aws-ssm-session.ts @@ -1,14 +1,42 @@ import { v4 as uuidv4 } from 'uuid'; import WebSocket from 'isomorphic-ws'; -import { decode, encode, EncodeableAgentMessage } from './agent-message'; -import { MessageType, PayloadType } from './channel-messages'; +import { AgentMessage, decode, encode, EncodeableAgentMessage } from './agent-message'; +import { ActionStatus, MessageType, PayloadType } from './channel-messages'; + +// Sent to the agent in handshake responses. Keep in sync with package.json. +const CLIENT_VERSION = '0.0.2'; + +// Maximum number of out-of-order inbound messages held while waiting for the +// gap before them to be filled. Matches the reference client's +// IncomingMessageBufferCapacity. +const INCOMING_MESSAGE_BUFFER_CAPACITY = 10000; + +// Maximum number of sent-but-unacknowledged messages kept for +// retransmission. Matches the reference client's +// OutgoingMessageBufferCapacity. +const OUTGOING_MESSAGE_BUFFER_CAPACITY = 10000; + +// How often to check whether the oldest unacknowledged message needs +// resending. Matches the reference client's ResendSleepInterval. +const RESEND_CHECK_INTERVAL = 100; + +// How long to wait for an acknowledgement before resending. The reference +// client starts at 200ms and tunes this from measured round-trip times; we +// don't measure RTT, so use a conservative fixed value. +const RETRANSMISSION_TIMEOUT = 1000; interface Handler { event: 'output' | 'connect' | 'disconnect' | 'pause' | 'resume', callback: Function } +interface OutgoingMessage { + sequenceNumber: number, + data: Uint8Array, + lastSentTime: number, +} + export default class AWSSSMSession { webSocket: WebSocket; sessionId: string; @@ -19,17 +47,29 @@ export default class AWSSSMSession { // an incrementing integer). Also, messages can be ordered back into sequence if // anything were to get ahead of line. outgoingSequenceNumber = 0; - outputMap = new Map(); + // The sequence number we expect from the next inbound stream message. + // Messages arriving ahead of it are held in incomingMessageBuffer until + // the gap before them is filled; messages behind it are duplicates and + // are re-acknowledged but not re-processed. + expectedSequenceNumber = 0; + incomingMessageBuffer: Map = new Map(); connectionClosed = false; connectionTerminated = false; // Each message that is sent with a sequenceNumber is replied to with an // acknowledgement from the SSM remote agent. We keep track of the latest - // sequenceNumber that has been acknowledged, so we know what number to reset - // to if we fall out of sync with the remove server. Currently no buffering or - // replaying of messages exists, but it could be added in the future. + // sequenceNumber that has been acknowledged, so we know which of our sent + // messages the remote agent has received. lastAcknowledgedSequenceNumber: number | undefined; + // Sent messages we haven't received an acknowledgement for yet, oldest + // first, kept so they can be retransmitted. + outgoingMessageBuffer: OutgoingMessage[] = []; + resendTimer?: ReturnType; handlers: Handler[] = []; paused: boolean = false; + // Terminal output is decoded with a single streaming decoder so that + // multi-byte UTF-8 sequences split across two messages survive the + // boundary intact. + private outputDecoder = new TextDecoder( 'utf8' ); constructor( streamUrl: string, sessionId: string, tokenValue: string ) { this.streamUrl = streamUrl; @@ -54,21 +94,39 @@ export default class AWSSSMSession { switch( message.MessageType ) { case 'output_stream_data': { - this.acknowledgeMessage( message.MessageType, message.MessageId, message.SequenceNumber ); + if ( message.SequenceNumber === this.expectedSequenceNumber ) { + this.acknowledgeMessage( message.MessageType, message.MessageId, message.SequenceNumber ); + this.handleOutputMessage( message ); + this.expectedSequenceNumber++; - const text = message.Payload; - if ( ! this.outputMap.has( message.MessageId ) ) { - this.emit( 'output', text ); + // 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.outputMap.set(message.MessageId, text); break } case 'start_publication': { - if ( this.lastAcknowledgedSequenceNumber !== undefined ) { - this.outgoingSequenceNumber = this.lastAcknowledgedSequenceNumber; - } else { - this.outgoingSequenceNumber = 0; - } + // 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; @@ -81,9 +139,14 @@ export default class AWSSSMSession { case 'acknowledge': { if ( message.Payload ) { const payloadData = JSON.parse( message.Payload ); - if ( this.lastAcknowledgedSequenceNumber === undefined || this.lastAcknowledgedSequenceNumber < payloadData.AcknowledgedMessageSequenceNumber ) { - this.lastAcknowledgedSequenceNumber = payloadData.AcknowledgedMessageSequenceNumber; + 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 ); } break; } @@ -93,7 +156,7 @@ export default class AWSSSMSession { case 'channel_closed': { if ( message.Payload ) { const payloadData = JSON.parse( message.Payload ); - if ( payloadData.Output.length > 0 ) { + if ( payloadData.Output ) { this.emit( 'output', payloadData.Output + '\r\n' ); } } @@ -108,9 +171,33 @@ export default class AWSSSMSession { this.webSocket.onclose = (ev: any) => { this.connectionClosed = true; + if ( this.resendTimer !== undefined ) { + clearInterval( this.resendTimer ); + this.resendTimer = undefined; + } this.emit( 'disconnect', ev.reason ); }; + + // Retransmit unacknowledged messages, like the reference client's + // ResendStreamDataMessageScheduler. + this.resendTimer = setInterval( () => this.resendUnacknowledgedMessage(), RESEND_CHECK_INTERVAL ); + } + + // Resend the oldest unacknowledged message if its acknowledgement is + // overdue. Duplicates are safe: the receiver re-acknowledges and drops + // messages below its expected sequence number. + private resendUnacknowledgedMessage() { + const head = this.outgoingMessageBuffer[ 0 ]; + if ( ! head || this.webSocket.readyState !== 1 ) { + return; + } + if ( Date.now() - head.lastSentTime < RETRANSMISSION_TIMEOUT ) { + return; + } + + head.lastSentTime = Date.now(); + this.webSocket.send( head.data ); } send( message: string | ArrayBuffer | SharedArrayBuffer | ArrayBufferView | Blob ) { @@ -118,6 +205,20 @@ export default class AWSSSMSession { console.warn( 'WebSocket is not yet ready.' ); return; } + + // Keep sequential messages around until they're acknowledged, so they + // can be retransmitted if the acknowledgement doesn't arrive. + if ( message instanceof Uint8Array ) { + this.outgoingMessageBuffer.push( { + sequenceNumber: this.outgoingSequenceNumber, + data: message, + lastSentTime: Date.now(), + } ); + if ( this.outgoingMessageBuffer.length > OUTGOING_MESSAGE_BUFFER_CAPACITY ) { + this.outgoingMessageBuffer.shift(); + } + } + this.outgoingSequenceNumber++; this.webSocket.send( message ) } @@ -140,7 +241,16 @@ export default class AWSSSMSession { } ping() { - this.webSocket.send( 'ping' ); + // The reference client's keepalive is a WebSocket ping control frame + // with payload "keepalive", sent every 5 minutes. Node's `ws` exposes + // ping(); browsers can't send control frames, so fall back to a bare + // text frame, which the service tolerates. + const webSocket = this.webSocket as WebSocket & { ping?: ( data: string ) => void }; + if ( typeof webSocket.ping === 'function' ) { + webSocket.ping( 'keepalive' ); + } else { + this.webSocket.send( 'ping' ); + } } emit( eventName: 'output' | 'connect' | 'disconnect' | 'pause' | 'resume', data: string | ArrayBuffer | SharedArrayBuffer | ArrayBufferView | Blob ) { @@ -161,10 +271,79 @@ export default class AWSSSMSession { this.webSocket.close(); } + // Process an output_stream_data message according to its payload type. + // Only data payloads are terminal output; handshake payloads are part of + // session setup and must not be shown to the user. + private handleOutputMessage( message: AgentMessage ) { + switch ( message.PayloadType ) { + case PayloadType.Output: + case PayloadType.Error: + case PayloadType.StdErr: { + this.emit( 'output', this.outputDecoder.decode( message.PayloadBytes, { stream: true } ) ); + break; + } + case PayloadType.HandshakeRequest: { + this.handleHandshakeRequest( message ); + break; + } + case PayloadType.HandshakeComplete: { + const payloadData = JSON.parse( message.Payload ); + if ( payloadData.CustomerMessage ) { + this.emit( 'output', payloadData.CustomerMessage + '\r\n' ); + } + break; + } + default: { + console.error( `Unhandled payload type ${ message.PayloadType }.` ); + } + } + } + + // Respond to a handshake request from the agent. + // + // We can satisfy the SessionType action; anything else (notably + // 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 errors: string[] = []; + const actions = ( request.RequestedClientActions || [] ).map( ( action: { ActionType: string } ) => { + if ( action.ActionType === 'SessionType' ) { + return { + ActionType: action.ActionType, + ActionStatus: ActionStatus.Success, + }; + } + + const error = `Unsupported action ${ action.ActionType }`; + errors.push( error ); + return { + ActionType: action.ActionType, + ActionStatus: ActionStatus.Unsupported, + Error: error, + }; + } ); + + const response = { + ClientVersion: CLIENT_VERSION, + ProcessedClientActions: actions, + Errors: errors, + }; + this.send( encode( { + MessageType: MessageType.InputStreamDataMessage, + SequenceNumber: this.outgoingSequenceNumber, + Payload: JSON.stringify( response ), + PayloadType: PayloadType.HandshakeResponse, + } ) ); + } + private acknowledgeMessage( messageType: MessageType, messageId: string, messageSequenceNumber: number ) { const message: EncodeableAgentMessage = { MessageType: MessageType.AcknowledgeMessage, - SequenceNumber: this.outgoingSequenceNumber, + // Acknowledgements aren't part of the sequential stream; the + // reference client sends them with sequence number 0 and flags 3. + SequenceNumber: 0, + Flags: 3, PayloadType: PayloadType.Output, Payload: JSON.stringify( { AcknowledgedMessageType: messageType, diff --git a/src/channel-messages.ts b/src/channel-messages.ts index c22b963..ec87d44 100644 --- a/src/channel-messages.ts +++ b/src/channel-messages.ts @@ -34,4 +34,14 @@ export enum PayloadType { EncChallengeRequest = 8, EncChallengeResponse = 9, Flag = 10, + StdErr = 11, + ExitCode = 12, +} + +// ActionStatus is the result of processing a client action requested by the +// agent during the handshake. +export enum ActionStatus { + Success = 1, + Failed = 2, + Unsupported = 3, }