Skip to content

Commit 205dae0

Browse files
authoredNov 6, 2022
Added WebSocket support for getUpstream() (#275)
Signed-off-by: Matteo Collina <[email protected]> Signed-off-by: Matteo Collina <[email protected]>
1 parent 72703f2 commit 205dae0

File tree

3 files changed

+75
-2
lines changed

3 files changed

+75
-2
lines changed
 

‎index.js

+12-2
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,6 @@ class WebSocketProxy {
108108

109109
findUpstream (request) {
110110
const source = new URL(request.url, 'ws://127.0.0.1')
111-
112111
for (const { prefix, rewritePrefix, upstream, wsClientOptions } of this.prefixList) {
113112
if (source.pathname.startsWith(prefix)) {
114113
const target = new URL(source.pathname.replace(prefix, rewritePrefix), upstream)
@@ -162,7 +161,18 @@ function setupWebSocketProxy (fastify, options, rewritePrefix) {
162161
})
163162
}
164163

165-
wsProxy.addUpstream(fastify.prefix, rewritePrefix, options.upstream, options.wsClientOptions)
164+
if (options.upstream !== '') {
165+
wsProxy.addUpstream(fastify.prefix, rewritePrefix, options.upstream, options.wsClientOptions)
166+
} else if (typeof options.replyOptions.getUpstream === 'function') {
167+
wsProxy.findUpstream = function (request) {
168+
const source = new URL(request.url, 'ws://127.0.0.1')
169+
const upstream = options.replyOptions.getUpstream(request, '')
170+
const target = new URL(source.pathname, upstream)
171+
target.protocol = upstream.indexOf('http:') === 0 ? 'ws:' : 'wss'
172+
target.search = source.search
173+
return { target, wsClientOptions: options.wsClientOptions }
174+
}
175+
}
166176
}
167177

168178
function generateRewritePrefix (prefix = '', opts) {

‎test/test.js

+3
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,9 @@ async function run () {
6666
const server = Fastify()
6767
server.register(proxy, {
6868
upstream: '',
69+
getWebSocketUpstream () {
70+
t.fail('should never be called')
71+
},
6972
replyOptions: {
7073
getUpstream: function (original, base) {
7174
return `http://localhost:${origin.server.address().port}`

‎test/websocket.js

+60
Original file line numberDiff line numberDiff line change
@@ -125,3 +125,63 @@ test('captures errors on start', async (t) => {
125125
t.teardown(app.close.bind(app))
126126
t.teardown(app2.close.bind(app2))
127127
})
128+
129+
test('getWebSocketStream', async (t) => {
130+
t.plan(7)
131+
132+
const origin = createServer()
133+
const wss = new WebSocket.Server({ server: origin })
134+
t.teardown(wss.close.bind(wss))
135+
t.teardown(origin.close.bind(origin))
136+
137+
const serverMessages = []
138+
wss.on('connection', (ws, request) => {
139+
t.equal(ws.protocol, subprotocolValue)
140+
t.equal(request.headers.cookie, cookieValue)
141+
ws.on('message', (message, binary) => {
142+
serverMessages.push([message.toString(), binary])
143+
// echo
144+
ws.send(message, { binary })
145+
})
146+
})
147+
148+
await promisify(origin.listen.bind(origin))({ port: 0 })
149+
150+
const server = Fastify()
151+
server.register(proxy, {
152+
upstream: '',
153+
replyOptions: {
154+
getUpstream: function (original, base) {
155+
return `http://localhost:${origin.address().port}`
156+
}
157+
},
158+
websocket: true
159+
})
160+
161+
await server.listen({ port: 0 })
162+
t.teardown(server.close.bind(server))
163+
164+
const options = { headers: { cookie: cookieValue } }
165+
const ws = new WebSocket(`ws://localhost:${server.server.address().port}`, [subprotocolValue], options)
166+
await once(ws, 'open')
167+
168+
ws.send('hello', { binary: false })
169+
const [reply0, binary0] = await once(ws, 'message')
170+
t.equal(reply0.toString(), 'hello')
171+
t.equal(binary0, false)
172+
173+
ws.send(Buffer.from('fastify'), { binary: true })
174+
const [reply1, binary1] = await once(ws, 'message')
175+
t.equal(reply1.toString(), 'fastify')
176+
t.equal(binary1, true)
177+
178+
t.strictSame(serverMessages, [
179+
['hello', false],
180+
['fastify', true]
181+
])
182+
183+
await Promise.all([
184+
once(ws, 'close'),
185+
server.close()
186+
])
187+
})

0 commit comments

Comments
 (0)
Please sign in to comment.