sharded token bucket for admission control

tokenbucket.js

// SPDX-License-Identifier: CC0-1.0
// token bucket with shuffle-sharding for ip admissions
class ShardedTokenBucket {
    #depth = 5 // buckets within buckets
    #replenish = null // replenish interval
    #allow = true // const
    #deny = false // const
    #marked = new Set()
    #onesec = 1000 // rate at which to fill bucket capacity
    #fps = 0 // fill rate per ratems
    #max4 = 0xFF + 0x1
    #max6 = 0xFFFF + 0x1

    // buckets: (size ^ depth)
    // tokens: buckets * cap
    // tokens per ipsubnet: size * cap
    constructor(size, rate, cap) {
        // size of each shard, determines total tokens
        this.size = Math.abs(size) || 1
        // tokens per bucket
        this.capacity = cap
        // replenish this many tokens every sec
        this.#fps = Math.abs(rate) || 1
        // clean slate
        this.cleanStart()
        // resume resumes bg replenishing of marked bucket entries
        this.resume()
    }
    
    cleanStart() {
        // inits an empty bucket
        this.bucket = this.emptyBucket()
        // marked tracks buckets to be replenished
        this.#marked.clear()
    }

    resume() {
        let that = this
        clearInterval(that.#replenish)
        that.#replenish = setInterval(function() {
            // fill capacity only for marked buckets
            const tofill = that.#marked.values()
            for (let m of tofill) {
                let filled = true
                for (let e in m) {
                    if (m[e] === that.capacity) continue
                    logd("rep idx", e, "in", m)
                    m[e] = Math.min(that.capacity, m[e] + that.#fps)
                    filled &&= (m[e] === that.capacity)
                }
                // unmark buckets at full capacity
                if (filled) that.#marked.delete(m)
            }
        }, that.#onesec)
    }

    pause() {
        clearInterval(this.#replenish)
    }

    local4(o1, o2) {
        if (o1 == "0") return true // non-routable
        if (o1 == "10") return true
        if (o1 == "224") return true // multicast
        if (o1 == "127") return true // lo
        if (o1 == "100") return true // cgnat
        if (o1 == "192" && o2 == "168") return true
        if (o1 == "169" && o2 == "254") return true // link-local
        if (o1 == "172" && o2 == "16") return true
        return false
    }

    local6(o1) {
        if (o1 == "0") return true
        if (o1 == "fe80") return true // link-local
        if (o1 == "ff00") return true // mutlicast
        return false
    }

    emptyBucket() {
     	const dimens = new Array(this.#depth)
    	dimens.fill(this.size)
        // Array(size) [
        //     Array(size) [
        //         Array(size) [
        //             Array(size) [
        //                 ...depth no of times
        //             ]
        //         ]
        //     ]
        // ]
    	return this.bucketOf(dimens, this.capacity)
    }

    // https://stackoverflow.com/a/966938/
    bucketOf(dimens, v) {
        const len = dimens[0] || 0
        const tb = new Array(len)
        if (dimens.length <= 1) {
    	    tb.fill(v) // fill default value, v
            return tb
        }

        const d = dimens.slice(1)
        for (let i = 0; i < len; i++) tb[i] = this.bucketOf(d, v)

        return tb
    }

    random(min, max) { // (min-open, max-closed]
        return Math.floor(Math.random() * (max - min)) + min
    }

    admit(i0, i1, i2, i3, i4, cost) {
        const w = this.bucket[i0][i1][i2][i3]

        if (w[i4] <= 0) {
            return this.#deny
        }

        w[i4] -= cost
        this.#marked.add(w)
        logd("marked", i0, i1, i2, i3, i4, " / weight", w[i4])

        return (w[i4] >= 0) ? this.#allow : this.#deny
    }

    admit4(ip4, cost) {
        // ex: 1.22.133.244
        const octets = ip4.split(".")
        const ipsubnet = Math.floor(this.#max4 / this.size)
        cost = cost || 1

        if (this.local4(octets[0], octets[1])) return this.#deny

        const i0 = Math.floor(parseInt(octets[0] || 0) / ipsubnet)
        const i1 = Math.floor(parseInt(octets[1] || 0) / ipsubnet)
        const i2 = Math.floor(parseInt(octets[2] || 0) / ipsubnet)
        const i3 = Math.floor(parseInt(octets[3] || 0) / ipsubnet)
        const ir = this.random(0, this.size) // shuffle-of-size
        const ix = (ir % 2 === 0) ? this.size - (i0 + 1) : i0 // power-of-2

        return this.admit(ix, i1, i2, i3, ir, cost)
    }

    admit6(ip6, cost) {
        // ex: 2606:2800:220:1:248:1893:25c8:1946
        // ex: 2609:80:3::2:a10
        const hex = ip6.split(":")
        const ipsubnet = Math.floor(this.#max6 / this.size)
        cost = cost || 1

        if (this.local6(hex[0])) return this.#deny

        const i0 = Math.floor(parseInt(hex[0] || 0, 16) / ipsubnet)
        const i1 = Math.floor(parseInt(hex[1] || 0, 16) / ipsubnet)
        const i2 = Math.floor(parseInt(hex[hex.length-2] || 0, 16) / ipsubnet)
        const i3 = Math.floor(parseInt(hex[hex.length-1] || 0, 16) / ipsubnet)
        const ir = this.random(0, this.size) // shuffle-of-size
        const ix = (ir % 2 === 0) ? this.size - (i0 + 1) : i0 // power-of-2

        return this.admit(ix, i1, i2, i3, ir, cost)
    }

}

function logd() {
    if (debug) console.debug(...arguments)
}

let debug = false
const tb = new ShardedTokenBucket(/*size*/ 4, /*fill-rate*/ 1, /*cap*/ 5)

function naivetest() {
    let ok = 0
    let nok = 0
    let p = []
    for (let i = 0; i < 1000; i++) {
        if (tb.admit4("1.2.3.4", 1)) {
            ok++
            p.push(i)
        } else {
            nok++
        }
    }
    logd("p/allow/deny", p, ok, nok)
}

function naivetest6() {
    let ok = 0
    let nok = 0
    let p = []
    let ips = []
    for (let i = 0; i < 1000; i++) {
        const ip6 = tb.random(1, 0x10000).toString(16) + ":" +
              tb.random(1, 0x10000).toString(16) + ":" +
              tb.random(0, 0x10000).toString(16) + ":" +
              tb.random(0, 0x10000).toString(16) + ":" +
              tb.random(0, 0x10000).toString(16) + ":" +
              tb.random(0, 0x10000).toString(16) + ":" +
              tb.random(1, 0x10000).toString(16) + ":" +
              tb.random(1, 0x10000).toString(16)
        ips.push(ip6)
    }
    for (let i = 0; i < 10000; i++) {
        if (tb.admit6(ips[i%ips.length], 1)) {
            ok++
            p.push(i)
        } else {
            nok++
        }
    }
    logd("p6/allow6/deny6", p, ok, nok)
}