path: root/layer/broadlayer.ts

import { randomUUIDv7 } from "bun"

if (process.argv.length < 4) {
    console.log(`more args.\n${process.argv[0]} <ip> <port>`)
}

const ip = process.argv[2] ?? '0.0.0.0'
const port = Number(process.argv[3])

if (!Number.isInteger(port) || port < 1 || port > 65535) {
    console.log("bad port")
    process.exit(1)
}

type tcpsocket = {
    mode: 'tcp'
    socket?: Bun.Socket
}
type udpsocket = {
    mode: 'udp'
    handler: (buf: Buffer) => void
    die: () => void
    heartbeat?: Timer
}
type socket = tcpsocket | udpsocket
type connection = {
    ip: string
    port: number
    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 index = Object.values(sessions).findIndex(s =>
                s.mode === 'udp' &&
                s.ip === ip &&
                s.port === port
            )
            if (index === -1) return

            const id = Object.keys(sessions)[index]!

            // send to handler
            if (sessions[id] && sessions[id].mode === 'udp') {
                sessions[id].handler(data)
            }
        },
        error(socket, error) {
            // dead.
            for (const key of Object.keys(sessions)) {
                const sess = sessions[key]!
                if (sess.mode === 'tcp') continue
                delete sessions[key]
                sess.die()
            }
        }
    }
})

const connectionsPerIp = new Map<string, number>()
const sessions: Record<string, connection> = {}

async function establishTCP(ip: string, port: number, clientIp: string): Promise<string | false> {
    const id = randomUUIDv7()

    const conn: connection = {
        port,
        ip,
        mode: 'tcp',
        socket: undefined,
        queue: [],
        lastSeen: Date.now(),
        clientIp
    }

    const sockPromise = Bun.connect({
        port,
        hostname: ip,
        reusePort: true,
        tls: false,
        socket: {
            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]
            }
        }
    })

    const timeoutPromise: Promise<0> = new Promise((r) => setTimeout(() => r(0), 5000))

    const sock: Bun.Socket | 0 = await Promise.race([sockPromise, timeoutPromise]).catch(err => {
        console.log(err)
        return 0
    })

    if (sock === 0) return false

    conn.socket = sock
    sessions[id] = conn
    return id
}

async function establishUDP(ip: string, port: number, clientIp: string): Promise<string> {
    const id = randomUUIDv7()

    const conn: connection = {
        port,
        ip,
        mode: 'udp',
        handler: undefined as any,
        die: undefined as any,
        heartbeat: undefined,
        queue: [],
        lastSeen: Date.now(),
        clientIp
    }

    conn.handler = (data) => {
        if (!conn.heartbeat) {
            const str = data.toString()
            const bnI = str.indexOf('\n')
            if (str.startsWith('T') && bnI !== -1) {
                const tm = Number(str.substring(1, bnI))
                if (!isNaN(tm)) {
                    conn.heartbeat = setInterval(() => {
                        udp.send('N\n', conn.port, conn.ip)
                    }, tm * 1000)
                }
                const temp_len = bnI + 1
                if (str.length > temp_len) {
                    data = Buffer.from(str.substring(temp_len))
                } else {
                    return
                }
            }
        }
        // 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
    udp.send("N\n", port, ip) // handshake

    return id
}

const eviction_timeout = 60_000

setInterval(() => {
    for (const key of Object.keys(sessions)) {
        const sess = sessions[key]!
        if (sess.lastSeen > Date.now() - eviction_timeout) continue
        // evict
        if (sess.mode === 'tcp') sess.socket?.close()
        else sess.die()
        delete sessions[key]
    }
}, eviction_timeout)

Bun.serve({
    port,
    hostname: ip,
    routes: {
        '/': new Response('this is a broadlayer server.'),
        '/tcp/connect/:ip/:port': async req => {
            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)
            return new Response('ok_' + res)
        },
        '/udp/connect/:ip/:port': async req => {
            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)

            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]
            if (!sess) return new Response('err_badid')
            const resp = JSON.stringify(sess.queue)
            sess.queue = []
            sess.lastSeen = Date.now()
            return new Response(resp)
        },
        '/peek/:id': async req => { // read queue without consuming
            const sess = sessions[req.params.id]
            if (!sess) return new Response('err_badid')
            const resp = JSON.stringify(sess.queue)
            sess.lastSeen = Date.now()
            return new Response(resp)
        },
        '/close/:id': async req => { // closes a connection
            const sess = sessions[req.params.id]
            if (!sess) return new Response('err_badid')
            if (sess.mode === 'tcp') sess.socket?.close()
            else {
                sess.die()
                delete sessions[req.params.id]
            }

            return new Response('ok')
        },
        '/send/:id': {
            POST: async req => {
                const sess = sessions[req.params.id]
                if (!sess) return new Response('err_badid')

                if (!req.body) return new Response('err_empty')
                const buf = await (await req.body.blob())!.arrayBuffer()
                if (sess.mode === 'tcp') sess.socket?.write(buf)
                else udp.send(buf, sess.port, sess.ip)

                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()
            }
        }
    }
})