@@ -11,22 +11,24 @@ import { ServerChannelTransport, getServerChannel } from '../get-server-channel'
1111describe ( 'getServerChannel' , ( ) => {
1212 it ( 'should return a channel' , ( ) => {
1313 const server = { on : vi . fn ( ) } as any as Server ;
14- const result = getServerChannel ( server ) ;
14+ const result = getServerChannel ( server , 'test-token-123' ) ;
1515 expect ( result ) . toBeInstanceOf ( Channel ) ;
1616 } ) ;
1717
1818 it ( 'should attach to the http server' , ( ) => {
1919 const server = { on : vi . fn ( ) } as any as Server ;
20- getServerChannel ( server ) ;
20+ getServerChannel ( server , 'test-token-123' ) ;
2121 expect ( server . on ) . toHaveBeenCalledWith ( 'upgrade' , expect . any ( Function ) ) ;
2222 } ) ;
2323} ) ;
2424
2525describe ( 'ServerChannelTransport' , ( ) => {
26+ const mockToken = 'test-token-123' ;
27+
2628 it ( 'parses simple JSON' , ( ) => {
2729 const server = new EventEmitter ( ) as any as Server ;
2830 const socket = new EventEmitter ( ) ;
29- const transport = new ServerChannelTransport ( server ) ;
31+ const transport = new ServerChannelTransport ( server , mockToken ) ;
3032 const handler = vi . fn ( ) ;
3133 transport . setHandler ( handler ) ;
3234
@@ -36,10 +38,11 @@ describe('ServerChannelTransport', () => {
3638
3739 expect ( handler ) . toHaveBeenCalledWith ( 'hello' ) ;
3840 } ) ;
41+
3942 it ( 'parses object JSON' , ( ) => {
4043 const server = new EventEmitter ( ) as any as Server ;
4144 const socket = new EventEmitter ( ) ;
42- const transport = new ServerChannelTransport ( server ) ;
45+ const transport = new ServerChannelTransport ( server , mockToken ) ;
4346 const handler = vi . fn ( ) ;
4447 transport . setHandler ( handler ) ;
4548
@@ -49,10 +52,11 @@ describe('ServerChannelTransport', () => {
4952
5053 expect ( handler ) . toHaveBeenCalledWith ( { type : 'hello' } ) ;
5154 } ) ;
55+
5256 it ( 'supports telejson cyclical data' , ( ) => {
5357 const server = new EventEmitter ( ) as any as Server ;
5458 const socket = new EventEmitter ( ) ;
55- const transport = new ServerChannelTransport ( server ) ;
59+ const transport = new ServerChannelTransport ( server , mockToken ) ;
5660 const handler = vi . fn ( ) ;
5761 transport . setHandler ( handler ) ;
5862
@@ -70,4 +74,52 @@ describe('ServerChannelTransport', () => {
7074 }
7175 ` ) ;
7276 } ) ;
77+
78+ it ( 'rejects connections with invalid token' , ( ) => {
79+ const server = new EventEmitter ( ) as any as Server ;
80+ const socket = new EventEmitter ( ) as any ;
81+ socket . write = vi . fn ( ) ;
82+ socket . destroy = vi . fn ( ) ;
83+ const destroySpy = vi . spyOn ( socket , 'destroy' ) ;
84+ new ServerChannelTransport ( server , mockToken ) ;
85+
86+ // Simulate upgrade request with wrong token
87+ const request = {
88+ url : '/storybook-server-channel?token=wrong-token' ,
89+ } as any ;
90+ const head = Buffer . from ( '' ) ;
91+
92+ server . listeners ( 'upgrade' ) [ 0 ] ( request , socket , head ) ;
93+
94+ expect ( socket . write ) . toHaveBeenCalledWith (
95+ 'HTTP/1.1 403 Forbidden\r\nConnection: close\r\n\r\n'
96+ ) ;
97+ expect ( destroySpy ) . toHaveBeenCalled ( ) ;
98+ } ) ;
99+
100+ it ( 'accepts connections with valid token' , ( ) => {
101+ const server = new EventEmitter ( ) as any as Server ;
102+ const socket = new EventEmitter ( ) as any ;
103+ socket . write = vi . fn ( ) ;
104+ socket . destroy = vi . fn ( ) ;
105+ const destroySpy = vi . spyOn ( socket , 'destroy' ) ;
106+ const handleUpgradeSpy = vi . fn ( ) ;
107+ const transport = new ServerChannelTransport ( server , mockToken ) ;
108+
109+ // Mock handleUpgrade to track if it's called
110+ // @ts -expect-error (accessing private property)
111+ transport . socket . handleUpgrade = handleUpgradeSpy ;
112+
113+ // Simulate upgrade request with correct token
114+ const request = {
115+ url : `/storybook-server-channel?token=${ mockToken } ` ,
116+ } as any ;
117+ const head = Buffer . from ( '' ) ;
118+
119+ server . listeners ( 'upgrade' ) [ 0 ] ( request , socket , head ) ;
120+
121+ expect ( socket . write ) . not . toHaveBeenCalled ( ) ;
122+ expect ( destroySpy ) . not . toHaveBeenCalled ( ) ;
123+ expect ( handleUpgradeSpy ) . toHaveBeenCalled ( ) ;
124+ } ) ;
73125} ) ;
0 commit comments