import { setImmediate } from "timers/promises" ; import { RequestContext , RequestContextOptions , getRequestContext , } from "@miniflare/shared" ; import { TestInputGate , noop , triggerPromise , useServer , waitsForOutputGate , } from "@miniflare/shared-test" ; import { CloseEvent , MessageEvent , WebSocket , WebSocketPair , coupleWebSocket , } from "@miniflare/web-sockets" ; import test , { ExecutionContext } from "ava" ; import StandardWebSocket from "ws" ; test ( "WebSocket: can accept multiple times" , ( t ) => { const webSocket = new WebSocket ( ) ; webSocket . accept ( ) ; webSocket . accept ( ) ; t . pass ( ) ; } ) ; test ( "WebSocket: cannot accept if already coupled" , async ( t ) => { const server = await useServer ( t , noop , ( ws ) => ws . send ( "test" ) ) ; const ws = new StandardWebSocket ( server . ws ) ; const [ webSocket1 ] = Object . values ( new WebSocketPair ( ) ) ; await coupleWebSocket ( ws , webSocket1 ) ; t . throws ( ( ) => webSocket1 . accept ( ) , { instanceOf : TypeError , message : "Can't accept() WebSocket that was already used in a response." , } ) ; } ) ; test ( "WebSocket: sends message to pair" , async ( t ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; webSocket2 . accept ( ) ; const messages1 = [ ] ; const messages2 = [ ] ; webSocket1 . addEventListener ( "message" , ( e ) => messages1 . push ( e . data ) ) ; webSocket2 . addEventListener ( "message" , ( e ) => messages2 . push ( e . data ) ) ; webSocket1 . send ( "from1" ) ; await setImmediate ( ) ; t . deepEqual ( messages1 , [ ] ) ; t . deepEqual ( messages2 , [ "from1" ] ) ; webSocket2 . send ( "from2" ) ; await setImmediate ( ) ; t . deepEqual ( messages1 , [ "from2" ] ) ; t . deepEqual ( messages2 , [ "from1" ] ) ; } ) ; test ( "WebSocket: must accept before sending" , ( t ) => { const [ webSocket1 ] = Object . values ( new WebSocketPair ( ) ) ; t . throws ( ( ) => webSocket1 . send ( "test" ) , { instanceOf : TypeError , message : "You must call accept() on this WebSocket before sending messages." , } ) ; } ) ; test ( "WebSocket: queues messages if pair not accepted" , async ( t ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; const messages1 = [ ] ; const messages2 = [ ] ; webSocket1 . addEventListener ( "message" , ( e ) => messages1 . push ( e . data ) ) ; webSocket2 . addEventListener ( "message" , ( e ) => messages2 . push ( e . data ) ) ; webSocket1 . accept ( ) ; webSocket1 . send ( "from1_1" ) ; await setImmediate ( ) ; t . deepEqual ( messages1 , [ ] ) ; t . deepEqual ( messages2 , [ ] ) ; webSocket2 . accept ( ) ; webSocket2 . send ( "from2_1" ) ; await setImmediate ( ) ; t . deepEqual ( messages1 , [ "from2_1" ] ) ; t . deepEqual ( messages2 , [ "from1_1" ] ) ; webSocket1 . send ( "from1_2" ) ; webSocket2 . send ( "from2_2" ) ; await setImmediate ( ) ; t . deepEqual ( messages1 , [ "from2_1" , "from2_2" ] ) ; t . deepEqual ( messages2 , [ "from1_1" , "from1_2" ] ) ; } ) ; test ( "WebSocket: fails to send message to pair if either side closed" , ( t ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; webSocket2 . accept ( ) ; webSocket1 . close ( ) ; t . throws ( ( ) => webSocket1 . send ( "from1" ) , { instanceOf : Error , message : "Can't call WebSocket send() after close()." , } ) ; t . throws ( ( ) => webSocket2 . send ( "from2" ) , { instanceOf : Error , message : "Can't call WebSocket send() after close()." , } ) ; } ) ; test ( "WebSocket: closes both sides of pair" , async ( t ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; webSocket2 . accept ( ) ; const closes = [ ] ; webSocket1 . addEventListener ( "close" , ( ) => closes . push ( 1 ) ) ; webSocket2 . addEventListener ( "close" , ( ) => closes . push ( 2 ) ) ; webSocket1 . close ( ) ; await setImmediate ( ) ; t . deepEqual ( closes , [ 1 , 2 ] ) ; } ) ; test ( "WebSocket: waits for output gate to open before sending message" , async ( t ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; webSocket2 . accept ( ) ; let event ; webSocket2 . addEventListener ( "message" , ( e ) => ( event = e ) ) ; await waitsForOutputGate ( t , ( ) => webSocket1 . send ( "test" ) , ( ) => event ) ; } ) ; test ( "WebSocket: waits for output gate to open before closing" , async ( t ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; webSocket2 . accept ( ) ; let event ; webSocket2 . addEventListener ( "close" , ( e ) => ( event = e ) ) ; await waitsForOutputGate ( t , ( ) => webSocket1 . close ( ) , ( ) => event ) ; } ) ; test ( "WebSocket: must accept before closing" , ( t ) => { const [ webSocket1 ] = Object . values ( new WebSocketPair ( ) ) ; t . throws ( ( ) => webSocket1 . close ( ) , { instanceOf : TypeError , message : "You must call accept() on this WebSocket before sending messages." , } ) ; } ) ; test ( "WebSocket: can only call close once" , ( t ) => { const [ webSocket1 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; webSocket1 . close ( 1000 ) ; t . throws ( ( ) => webSocket1 . close ( 1000 ) , { instanceOf : TypeError , message : "WebSocket already closed" , } ) ; } ) ; test ( "WebSocket: validates close code" , ( t ) => { const [ webSocket1 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; t . throws ( ( ) => webSocket1 . close ( 1005 ) , { instanceOf : TypeError , message : "Invalid WebSocket close code." , } ) ; t . throws ( ( ) => webSocket1 . close ( undefined , "Test Closure" ) , { instanceOf : TypeError , message : "If you specify a WebSocket close reason, you must also specify a code." , } ) ; } ) ; async function eventDispatchWaitsForInputGate ( t , addEventListener , dispatchEvent ) { const inputGate = new TestInputGate ( ) ; const [ openTrigger , openPromise ] = triggerPromise < void > ( ) ; const events = [ ] ; const promise = inputGate . runWith ( ( ) => { void inputGate . runWithClosed ( ( ) => openPromise ) ; return new Promise ( ( resolve ) => { addEventListener ( resolve ) ; } ) . then ( ( ) => events . push ( 1 ) ) ; } ) ; await setImmediate ( ) ; dispatchEvent ( ) ; await inputGate . waitedPromise ; events . push ( 2 ) ; openTrigger ( ) ; await promise ; t . deepEqual ( events , [ 2 , 1 ] ) ; } test ( "WebSocket: waits for input gate to open before receiving message" , async ( t ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; webSocket2 . accept ( ) ; await eventDispatchWaitsForInputGate ( t , ( listener ) => webSocket2 . addEventListener ( "message" , listener ) , ( ) => webSocket1 . send ( "test" ) ) ; } ) ; test ( "WebSocket: waits for input gate to open before receiving close event" , async ( t ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; webSocket2 . accept ( ) ; await eventDispatchWaitsForInputGate ( t , ( listener ) => webSocket2 . addEventListener ( "close" , listener ) , ( ) => webSocket1 . close ( ) ) ; } ) ; test ( "WebSocketPair: requires 'new' operator to construct" , ( t ) => { t . throws ( ( ) => WebSocketPair ( ) , { instanceOf : TypeError , message : /^Failed to construct 'WebSocketPair'/ , } ) ; } ) ; function testWebSocketPairTypes ( ) { const pair = new WebSocketPair ( ) ; let [ webSocket1 , webSocket2 ] = Object . values ( pair ) ; [ webSocket1 , webSocket2 ] = pair ; webSocket1 = pair [ 0 ] ; webSocket2 = pair [ 2 ] ; } function useSubrequest ( ) { getRequestContext ( ) ?. incrementSubrequests ( ) ; } const ctxOpts = { requestDepth : 5 , pipelineDepth : 10 } ; function assertSubrequests ( t , expected ) { const ctx = getRequestContext ( ) ; t . is ( ctx ?. subrequests , expected ) ; t . is ( ctx ?. requestDepth , ctxOpts . requestDepth ) ; t . is ( ctx ?. pipelineDepth , ctxOpts . pipelineDepth ) ; } test ( "WebSocket: shares subrequest limit for WebSockets in regular worker handler" , async ( t ) => { await new RequestContext ( ctxOpts ) . runWith ( async ( ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; webSocket2 . accept ( ) ; useSubrequest ( ) ; const [ trigger , promise ] = triggerPromise < void > ( ) ; webSocket1 . addEventListener ( "message" , ( ) => { assertSubrequests ( t , 1 ) ; trigger ( ) ; } ) ; webSocket2 . send ( "test" ) ; await promise ; } ) ; const [ trigger , promise ] = triggerPromise < void > ( ) ; const webSocket2 = await new RequestContext ( ctxOpts ) . runWith ( async ( ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; useSubrequest ( ) ; webSocket1 . addEventListener ( "message" , ( ) => { assertSubrequests ( t , 1 ) ; trigger ( ) ; } ) ; return webSocket2 ; } ) ; webSocket2 . accept ( ) ; webSocket2 . send ( "test" ) ; await promise ; } ) ; test ( "WebSocket: resets subrequest limit for WebSockets in Durable Object" , async ( t ) => { const durableOpts = { ... ctxOpts , durableObject : true , } ; await new RequestContext ( durableOpts ) . runWith ( async ( ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; webSocket2 . accept ( ) ; useSubrequest ( ) ; const [ trigger , promise ] = triggerPromise < void > ( ) ; webSocket1 . addEventListener ( "message" , ( ) => { assertSubrequests ( t , 0 ) ; trigger ( ) ; } ) ; webSocket2 . send ( "test" ) ; await promise ; } ) ; let [ trigger , promise ] = triggerPromise < void > ( ) ; let webSocket2 = await new RequestContext ( durableOpts ) . runWith ( async ( ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; useSubrequest ( ) ; webSocket1 . addEventListener ( "message" , ( ) => { assertSubrequests ( t , 0 ) ; trigger ( ) ; } ) ; return webSocket2 ; } ) ; new RequestContext ( ctxOpts ) . runWith ( ( ) => { webSocket2 . accept ( ) ; webSocket2 . send ( "test" ) ; } ) ; await promise ; [ trigger , promise ] = triggerPromise < void > ( ) ; webSocket2 = await new RequestContext ( durableOpts ) . runWith ( async ( ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; useSubrequest ( ) ; webSocket1 . addEventListener ( "message" , ( ) => { assertSubrequests ( t , 0 ) ; trigger ( ) ; } ) ; return webSocket2 ; } ) ; webSocket2 . accept ( ) ; webSocket2 . send ( "test" ) ; await promise ; } ) ; test ( "WebSocket: resets subrequest limit for WebSockets outside worker" , async ( t ) => { const webSocket2 = new RequestContext ( ctxOpts ) . runWith ( ( ) => { const [ webSocket1 , webSocket2 ] = Object . values ( new WebSocketPair ( ) ) ; webSocket1 . accept ( ) ; useSubrequest ( ) ; webSocket1 . send ( "test" ) ; return webSocket2 ; } ) ; const [ trigger , promise ] = triggerPromise < void > ( ) ; webSocket2 . accept ( ) ; webSocket2 . addEventListener ( "message" , ( ) => { const ctx = getRequestContext ( ) ; t . is ( ctx ?. subrequests , 0 ) ; t . is ( ctx ?. requestDepth , 1 ) ; t . is ( ctx ?. pipelineDepth , 1 ) ; trigger ( ) ; } ) ; await promise ; } ) ;