/*
-- broadlayer - http/websocket frontend for broadd
-- BEOFRE RUNNING:
------ THIS SERVER OPENS ACCESS TO YOUR PRIVATE IP NETWORK AND LOCALHOST PORTS.
------ IF THIS IS NOT INTENDED BEHAVIOR, PATCH IT OUT OF THE PROGRAM.
------ there is a hardcode_target const variable provided below, which you can modify to override any and all
------ incoming ip:port combos to the one that you set below.
------ it is enabled by default, and configured for the respective tcp and udp broadd ports on localhost.
--
------ http:
------ all error replies start with `err_`.
---------- GET /
-------------- returns "this is a broadlayer server.", guaranteed
---------- GET /motd
-------------- returns the server admin's message of the day. arbitrary size permitted.
---------- GET /timeout
-------------- returns the inactivity time after which you will get evicted.
---------- GET /tcp/connect/<hostname>/<port>
-------------- creates a broadd tcp connection to the given hostname/ip on the given port. supports ipv4+ipv6.
-------------- returns `ok_<id>` where id is your unique connection id. store this for polling.
---------- GET /udp/connect/<hostname>/<port>
-------------- creates a broadd udp connection to the given hostname/ip on the given port. supports only ipv4.
-------------- returns `ok_<id>` where id is your unique connection id. store this for polling.
---------- GET /alive/<id>
-------------- for telling the server that you are alive. if all is well, server replies with `me too`
---------- GET /poll/<id>
-------------- fetch and consume the messages that the server has sent you up to this point. returns an array of buffers
-------------- format is an array which contains buffers that are strings formatted with each character being an integer.
-------------- `[ { "type": "Buffer", "data": [ 98, 114, 111, 97, 100, 100 ] } ]` would be `[ "broadd" ]`.
---------- GET /peek/<id>
-------------- same as poll, but does not consume the data, allowing you to re-fetch it later.
---------- GET /close/<id>
-------------- terminates your connection and deletes your session.
-------------- replies with "ok" on success.
---------- POST /send/<id>
-------------- sends the POST body data to the open socket with the server.
-------------- replies with "ok" on success.
---------- GET /tcp/ws/<ip>/<port>
-------------- connects via tcp to the given server:port and upgrades your connection to websockets.
---------- GET /udp/ws/<ip>/<port>
-------------- connects via udp to the given server:port and upgrades your connection to websockets.
------ websocket:
---------- the websocket connection behaves exactly like a plain connection to the server would.
---------- only difference is a special command "ping", which returns "pong" and nothing is relayed to the target server.
---------- websocket returns `ok_<uuid>` as your client id. you can downgrade from websocket and switch to http if you want.
*/

const hardcode_target = true
const hardcode_host = '127.0.0.1'
const hardcode_tcp_port = 5001
const hardcode_udp_port = 5002

import { dns, randomUUIDv7 } from "bun"

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

const server_ip = process.argv[2] ?? '0.0.0.0'
const server_port = Number(process.argv[3])

if (!Number.isInteger(server_port) || server_port < 1 || server_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> {
    if (hardcode_target) {
        ip = hardcode_host
        port = hardcode_tcp_port
    }

    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(hostname: string, port: number, clientIp: string): Promise<string | false> {
    if (hardcode_target) {
        hostname = hardcode_host
        port = hardcode_udp_port
    }

    const id = randomUUIDv7()
    const ip = (await dns.lookup(hostname, { family: 'IPv4' }))[0]?.address
    if (!ip) return false
    console.log(ip)

    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(() => {
                        console.log('heartbeat')
                        udp.send('N\n', conn.port, conn.ip)
                    }, Math.max(tm - 5, 1) * 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 = 15_000
const ratelimit_count = 32

setInterval(() => {
    let hasEvicted = false
    for (const key of Object.keys(sessions)) {
        const sess = sessions[key]!
        if (sess.lastSeen > Date.now() - eviction_timeout) continue
        // skip if websocket
        if (sess.ws) continue
        // evict
        hasEvicted = true
        if (sess.mode === 'tcp') sess.socket?.close()
        else sess.die()
        delete sessions[key]
    }
    if (hasEvicted) console.log('clients:' + Object.keys(sessions).length)
}, eviction_timeout)

Bun.serve({
    port: '10001',
    hostname: server_ip,
    routes: {
        '/': new Response('this is a broadlayer server.'),
        '/motd': new Response('default message-of-the-day!'),
        '/timeout': new Response(eviction_timeout.toString()),
        '/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 >= ratelimit_count) {
                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)
            console.log('clients:' + Object.keys(sessions).length)
            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 >= ratelimit_count) {
                return new Response('err_toomany', { status: 429 })
            }

            const res = await establishUDP(req.params.ip, port, clientIp)
            if (!res) return new Response('err_badhost')
            connectionsPerIp.set(clientIp, count + 1)
            console.log('clients:' + Object.keys(sessions).length)
            return new Response('ok_' + res)
        },
        '/alive/:id': async req => { // heartbeat
            const sess = sessions[req.params.id]
            if (!sess) return new Response('err_badid')
            sess.lastSeen = Date.now()
            return new Response('me too')
        },
        '/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]
            }

            console.log('clients:' + Object.keys(sessions).length)
            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)

                console.log('sendto')
                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 >= ratelimit_count) {
                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}`
                }
            })

            console.log('clients:' + Object.keys(sessions).length)
            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 >= ratelimit_count) {
                return new Response('err_toomany', { status: 429 })
            }
            connectionsPerIp.set(clientIp, count + 1)

            const res = await establishUDP(req.params.ip, port, clientIp)
            if (!res) return new Response('err_badhost')

            server.upgrade(req, {
                data: { sessionId: res },
                headers: {
                    'Set-Cookie': `SessionID=${res}`
                }
            })

            console.log('clients:' + Object.keys(sessions).length)
            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()
                console.log('connclose')
            }
        }
    }
})

Bun.serve({
    port: server_port,
    hostname: server_ip,
    tls: {
        cert: Bun.file('/etc/letsencrypt/live/downsta.ir/cert.pem'),
        key: Bun.file('/etc/letsencrypt/live/downsta.ir/privkey.pem')
    },
    routes: {
        '/': new Response('this is a broadlayer server.'),
        '/motd': new Response('default message-of-the-day!'),
        '/timeout': new Response(eviction_timeout.toString()),
        '/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 >= ratelimit_count) {
                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)
            console.log('clients:' + Object.keys(sessions).length)
            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 >= ratelimit_count) {
                return new Response('err_toomany', { status: 429 })
            }

            const res = await establishUDP(req.params.ip, port, clientIp)
            if (!res) return new Response('err_badhost')
            connectionsPerIp.set(clientIp, count + 1)
            console.log('clients:' + Object.keys(sessions).length)
            return new Response('ok_' + res)
        },
        '/alive/:id': async req => { // heartbeat
            const sess = sessions[req.params.id]
            if (!sess) return new Response('err_badid')
            sess.lastSeen = Date.now()
            return new Response('me too')
        },
        '/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]
            }

            console.log('clients:' + Object.keys(sessions).length)
            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)

                console.log('sendto')
                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 >= ratelimit_count) {
                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}`
                }
            })

            console.log('clients:' + Object.keys(sessions).length)
            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 >= ratelimit_count) {
                return new Response('err_toomany', { status: 429 })
            }
            connectionsPerIp.set(clientIp, count + 1)

            const res = await establishUDP(req.params.ip, port, clientIp)
            if (!res) return new Response('err_badhost')

            server.upgrade(req, {
                data: { sessionId: res },
                headers: {
                    'Set-Cookie': `SessionID=${res}`
                }
            })

            console.log('clients:' + Object.keys(sessions).length)
            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()
                console.log('connclose')
            }
        }
    }
})