broadlayer impl

1 files changed, 157 insertions, 22 deletions

diff --git a/layer/broadlayer.ts b/layer/broadlayer.ts
index 1e0d415..fa39dc2 100644
--- a/layer/broadlayer.ts
+++ b/layer/broadlayer.ts
@@ -29,18 +29,30 @@ type connection = {
     queue: Buffer[]
     remoteId?: string
     lastSeen: number
+    clientIp: string
+    ws?: Bun.ServerWebSocket<{ sessionId: string }>
 } & socket
 
+function decrementRatelimit(clientIp: string) {
+    let count = connectionsPerIp.get(clientIp) || 0
+    if (count > 0) {
+        count--
+        if (count === 0) connectionsPerIp.delete(clientIp)
+        else connectionsPerIp.set(clientIp, count)
+    }
+}
+
 const udp = await Bun.udpSocket({
     socket: {
         data(socket, data, port, ip) {
-            const id = Object.values(sessions).findIndex(s =>
+            const index = Object.values(sessions).findIndex(s =>
                 s.mode === 'udp' &&
                 s.ip === ip &&
                 s.port === port
             )
+            if (index === -1) return
 
-            if (id === -1) return
+            const id = Object.keys(sessions)[index]!
 
             // send to handler
             if (sessions[id] && sessions[id].mode === 'udp') {
@@ -49,7 +61,10 @@ const udp = await Bun.udpSocket({
         },
         error(socket, error) {
             // dead.
-            for (const sess of Object.values(sessions).filter(e => e.mode === 'udp')) {
+            for (const key of Object.keys(sessions)) {
+                const sess = sessions[key]!
+                if (sess.mode === 'tcp') continue
+                delete sessions[key]
                 sess.die()
             }
         }
@@ -59,7 +74,7 @@ const udp = await Bun.udpSocket({
 const connectionsPerIp = new Map<string, number>()
 const sessions: Record<string, connection> = {}
 
-async function establishTCP(ip: string, port: number): Promise<string | false> {
+async function establishTCP(ip: string, port: number, clientIp: string): Promise<string | false> {
     const id = randomUUIDv7()
 
     const conn: connection = {
@@ -68,7 +83,8 @@ async function establishTCP(ip: string, port: number): Promise<string | false> {
         mode: 'tcp',
         socket: undefined,
         queue: [],
-        lastSeen: Date.now()
+        lastSeen: Date.now(),
+        clientIp
     }
 
     const sockPromise = Bun.connect({
@@ -77,10 +93,26 @@ async function establishTCP(ip: string, port: number): Promise<string | false> {
         reusePort: true,
         tls: false,
         socket: {
-            data(socket, data) { conn.queue.push(data) },
-            close(socket, error) { delete sessions[id] },
-            end(socket) { delete sessions[id] },
-            error(socket, error) { delete sessions[id] }
+            data(socket, data) {
+                // if ws => send there. else, push to queue.
+                if (conn.ws) conn.ws.send(data)
+                else conn.queue.push(data)
+            },
+            close(socket, error) {
+                const sess = sessions[id]
+                if (sess) decrementRatelimit(sess.clientIp)
+                delete sessions[id]
+            },
+            end(socket) {
+                const sess = sessions[id]
+                if (sess) decrementRatelimit(sess.clientIp)
+                delete sessions[id]
+            },
+            error(socket, error) {
+                const sess = sessions[id]
+                if (sess) decrementRatelimit(sess.clientIp)
+                delete sessions[id]
+            }
         }
     })
 
@@ -98,7 +130,7 @@ async function establishTCP(ip: string, port: number): Promise<string | false> {
     return id
 }
 
-async function establishUDP(ip: string, port: number): Promise<string> {
+async function establishUDP(ip: string, port: number, clientIp: string): Promise<string> {
     const id = randomUUIDv7()
 
     const conn: connection = {
@@ -109,7 +141,8 @@ async function establishUDP(ip: string, port: number): Promise<string> {
         die: undefined as any,
         heartbeat: undefined,
         queue: [],
-        lastSeen: Date.now()
+        lastSeen: Date.now(),
+        clientIp
     }
 
     conn.handler = (data) => {
@@ -131,11 +164,14 @@ async function establishUDP(ip: string, port: number): Promise<string> {
                 }
             }
         }
-        conn.queue.push(data)
+        // if ws => send there. else, push to queue.
+        if (conn.ws) conn.ws.send(data)
+        else conn.queue.push(data)
     }
     conn.die = () => {
         clearInterval(conn.heartbeat)
         conn.heartbeat = undefined
+        decrementRatelimit(clientIp)
     }
 
     sessions[id] = conn
@@ -167,15 +203,17 @@ Bun.serve({
             if (isNaN(port)) return new Response('err_badport')
             if (port > 65535 || port < 0) return new Response('err_badport')
 
+            // ratelimit
             const clientIp = req.headers.get('x-forwarded-for')?.split(',')[0]?.trim() || 'unknown'
             let count = connectionsPerIp.get(clientIp) || 0
             if (count >= 4) {
                 return new Response('err_toomany', { status: 429 })
             }
-            connectionsPerIp.set(clientIp, count + 1)
 
-            const res = await establishTCP(req.params.ip, port)
+            const res = await establishTCP(req.params.ip, port, clientIp)
             if (res === false) return new Response('err_timeout')
+
+            connectionsPerIp.set(clientIp, count + 1)
             return new Response('ok_' + res)
         },
         '/udp/connect/:ip/:port': async req => {
@@ -183,6 +221,7 @@ Bun.serve({
             if (isNaN(port)) return new Response('err_badport')
             if (port > 65535 || port < 0) return new Response('err_badport')
 
+            // ratelimit
             const clientIp = req.headers.get('x-forwarded-for')?.split(',')[0]?.trim() || 'unknown'
             let count = connectionsPerIp.get(clientIp) || 0
             if (count >= 4) {
@@ -190,7 +229,7 @@ Bun.serve({
             }
             connectionsPerIp.set(clientIp, count + 1)
 
-            return new Response('ok_' + await establishUDP(req.params.ip, port))
+            return new Response('ok_' + await establishUDP(req.params.ip, port, clientIp))
         },
         '/poll/:id': async req => { // read and consume queue
             const sess = sessions[req.params.id]
@@ -211,13 +250,10 @@ Bun.serve({
             const sess = sessions[req.params.id]
             if (!sess) return new Response('err_badid')
             if (sess.mode === 'tcp') sess.socket?.close()
-
-            delete sessions[req.params.id]
-
-            const clientIp = req.headers.get('x-forwarded-for')?.split(',')[0]?.trim() || 'unknown'
-            let count = connectionsPerIp.get(clientIp) || 0
-            if (count > 0) connectionsPerIp.set(clientIp, count - 1)
-            if (count <= 1) connectionsPerIp.delete(clientIp)
+            else {
+                sess.die()
+                delete sessions[req.params.id]
+            }
 
             return new Response('ok')
         },
@@ -233,6 +269,105 @@ Bun.serve({
 
                 return new Response('ok')
             }
+        },
+        // websocket upgrades
+        '/tcp/ws/:ip/:port': async (req, server) => {
+            const port = parseInt(req.params.port)
+            if (isNaN(port)) return new Response('err_badport')
+            if (port > 65535 || port < 0) return new Response('err_badport')
+
+            // ratelimit
+            const clientIp = req.headers.get('x-forwarded-for')?.split(',')[0]?.trim() || 'unknown'
+            let count = connectionsPerIp.get(clientIp) || 0
+            if (count >= 4) {
+                return new Response('err_toomany', { status: 429 })
+            }
+
+            const res = await establishTCP(req.params.ip, port, clientIp)
+            if (res === false) return new Response('err_timeout')
+
+            connectionsPerIp.set(clientIp, count + 1)
+
+            server.upgrade(req, {
+                data: { sessionId: res },
+                headers: {
+                    'Set-Cookie': `SessionID=${res}`
+                }
+            })
+
+            return new Response('ok_' + res)
+        },
+        '/udp/ws/:ip/:port': async (req, server) => {
+            const port = parseInt(req.params.port)
+            if (isNaN(port)) return new Response('err_badport')
+            if (port > 65535 || port < 0) return new Response('err_badport')
+
+            // ratelimit
+            const clientIp = req.headers.get('x-forwarded-for')?.split(',')[0]?.trim() || 'unknown'
+            let count = connectionsPerIp.get(clientIp) || 0
+            if (count >= 4) {
+                return new Response('err_toomany', { status: 429 })
+            }
+            connectionsPerIp.set(clientIp, count + 1)
+
+            const res = await establishUDP(req.params.ip, port, clientIp)
+
+            server.upgrade(req, {
+                data: { sessionId: res },
+                headers: {
+                    'Set-Cookie': `SessionID=${res}`
+                }
+            })
+
+            return new Response('ok_' + res)
+        }
+    },
+    websocket: {
+        data: { sessionId: '' },
+        perMessageDeflate: true,
+        open(ws) {
+            const sessionId = ws.data.sessionId
+            if (!sessionId || !sessions[sessionId]) {
+                ws.close()
+                return
+            }
+            sessions[sessionId].ws = ws
+            sessions[sessionId].lastSeen = Date.now()
+            ws.send('ok_' + sessionId)
+        },
+        message(ws, message) {
+            const sessionId = ws.data.sessionId
+            const sess = sessions[sessionId]
+            if (!sess) {
+                ws.close()
+                return
+            }
+            sess.lastSeen = Date.now()
+
+            if (typeof message === 'string') {
+                // ping
+                if (message === 'ping') {
+                    ws.send('pong')
+                    return
+                }
+                // send
+                const buf = Buffer.from(message)
+                if (sess.mode === 'tcp') sess.socket?.write(buf)
+                else udp.send(buf, sess.port, sess.ip)
+            } else {
+                // binary
+                if (sess.mode === 'tcp') sess.socket?.write(message)
+                else udp.send(message, sess.port, sess.ip)
+            }
+        },
+        close(ws) {
+            const sessionId = ws.data.sessionId
+            if (sessionId && sessions[sessionId]) {
+                // close websocket but keepalive connection, in case user wants to downgrade to polling
+                const sess = sessions[sessionId]
+                sess.ws = undefined
+                sess.lastSeen = Date.now()
+            }
         }
     }
 })
\ No newline at end of file