import { serializeFactory , aadFactory , MessageHeader , Maximum , } from '@aws-crypto/serialize' ; import { Transform as PortableTransform } from 'readable-stream' ; import { Transform } from 'stream' ; import { GetCipher , AwsEsdkJsCipherGCM , needs , NodeAlgorithmSuite , } from '@aws-crypto/material-management-node' ; const fromUtf8 = ( input ) => Buffer . from ( input , 'utf8' ) ; const serialize = serializeFactory ( fromUtf8 ) ; const { finalFrameHeader , frameHeader } = serialize ; const aadUtility = aadFactory ( fromUtf8 ) ; interface AccumulatingFrame { contentLength : number ; content : Buffer [ ] ; sequenceNumber : number ; } interface EncryptFrame { content : Buffer [ ] ; bodyHeader : Buffer ; headerSent ? : boolean ; cipher : AwsEsdkJsCipherGCM ; isFinalFrame : boolean ; } const PortableTransformWithType = PortableTransform as new ( ... args ) => Transform ; const ioTick = async ( ) => new Promise ( ( resolve ) => setImmediate ( resolve ) ) ; const noop = ( ) => { } ; type ErrBack = ( err ? : Error ) => void ; export function getFramedEncryptStream ( getCipher , messageHeader , dispose , { plaintextLength , suite , } ) { let accumulatingFrame = { contentLength : 0 , content : [ ] , sequenceNumber : 1 , } ; let pathologicalDrain = noop ; const { frameLength } = messageHeader ; needs ( ! plaintextLength || ( plaintextLength >= 0 && Maximum . BYTES_PER_MESSAGE >= plaintextLength ) , 'plaintextLength out of bounds.' ) ; return new ( class FramedEncryptStream extends PortableTransformWithType { _transform ( chunk , encoding , callback ) { const contentLeft = frameLength - accumulatingFrame . contentLength ; needs ( ! plaintextLength || ( plaintextLength -= chunk . length ) >= 0 , 'Encrypted data exceeded plaintextLength.' ) ; if ( contentLeft > chunk . length ) { accumulatingFrame . contentLength += chunk . length ; accumulatingFrame . content . push ( chunk ) ; return callback ( ) ; } accumulatingFrame . contentLength += contentLeft ; accumulatingFrame . content . push ( chunk . slice ( 0 , contentLeft ) ) ; const tail = chunk . slice ( contentLeft ) ; const encryptFrame = getEncryptFrame ( { pendingFrame : accumulatingFrame , messageHeader , getCipher , isFinalFrame : false , suite , } ) ; const { sequenceNumber } = accumulatingFrame ; accumulatingFrame = { contentLength : 0 , content : [ ] , sequenceNumber : sequenceNumber + 1 , } ; this . _flushEncryptFrame ( encryptFrame ) . then ( ( ) => this . _transform ( tail , encoding , callback ) ) . catch ( callback ) ; } _flush ( callback ) { const encryptFrame = getEncryptFrame ( { pendingFrame : accumulatingFrame , messageHeader , getCipher , isFinalFrame : true , suite , } ) ; this . _flushEncryptFrame ( encryptFrame ) . then ( ( ) => callback ( ) ) . catch ( callback ) ; } _destroy ( ) { dispose ( ) ; } _read ( size ) { super . _read ( size ) ; pathologicalDrain ( ) ; pathologicalDrain = noop ; } async _flushEncryptFrame ( encryptingFrame ) { const { content , cipher , bodyHeader , isFinalFrame } = encryptingFrame ; this . push ( bodyHeader ) ; let frameSize = 0 ; const cipherContent = [ ] ; for ( const clearChunk of content ) { const cipherText = cipher . update ( clearChunk ) ; frameSize += cipherText . length ; cipherContent . push ( cipherText ) ; await ioTick ( ) ; } const tail = cipher . final ( ) ; frameSize += tail . length ; cipherContent . push ( tail ) ; cipherContent . push ( cipher . getAuthTag ( ) ) ; needs ( frameSize === frameLength || ( isFinalFrame && frameLength >= frameSize ) , 'Malformed frame' ) ; for ( const cipherText of cipherContent ) { if ( ! this . push ( cipherText ) ) { await new Promise ( ( resolve ) => { pathologicalDrain = resolve ; } ) ; } } if ( isFinalFrame ) this . push ( null ) ; } } ) ( ) ; } type EncryptFrameInput = { pendingFrame : AccumulatingFrame ; messageHeader : MessageHeader ; getCipher : GetCipher ; isFinalFrame : boolean ; suite : NodeAlgorithmSuite ; } ; export function getEncryptFrame ( input ) { const { pendingFrame , messageHeader , getCipher , isFinalFrame , suite } = input ; const { sequenceNumber , contentLength , content } = pendingFrame ; const { frameLength , contentType , messageId } = messageHeader ; needs ( frameLength === contentLength || ( isFinalFrame && frameLength >= contentLength ) , `Malformed frame length and content length: ${ JSON . stringify ( { frameLength , contentLength , isFinalFrame , } ) }` ) ; const frameIv = serialize . frameIv ( suite . ivLength , sequenceNumber ) ; const bodyHeader = Buffer . from ( isFinalFrame ? finalFrameHeader ( sequenceNumber , frameIv , contentLength ) : frameHeader ( sequenceNumber , frameIv ) ) ; const contentString = aadUtility . messageAADContentString ( { contentType , isFinalFrame , } ) ; const { buffer , byteOffset , byteLength } = aadUtility . messageAAD ( messageId , contentString , sequenceNumber , contentLength ) ; const cipher = getCipher ( frameIv ) ; cipher . setAAD ( Buffer . from ( buffer , byteOffset , byteLength ) ) ; return { content , cipher , bodyHeader , isFinalFrame } ; }