Manipulate CIDRs with Node.js (Trie implementation)

ipset.js

import { parseIp, stringifyIp } from "ip-bigint";

/**
 * @typedef { { version: 4 | 6; address: bigint; mask: number; } } CIDR
 */

/**
 * @param {string} str
 * @returns {CIDR}
 */
function parseCIDR(str) {
    let mask = null;
    const idxSlash = str.indexOf("/");
    if (idxSlash !== -1) {
        const maskStr = str.substring(idxSlash + 1);
        if (!/^\d+$/.test(maskStr)) {
            return new Error("Invalid CIDR");
        }
        mask = parseInt(maskStr);
        str = str.substring(0, idxSlash);
    }

    const { number, version } = parseIp(str);
    if (mask === null) {
        mask = version === 4 ? 32 : 128;
    }
    if ((version === 4 && mask <= 32) || (version === 6 && mask <= 128)) {
        return {
            version: version,
            address: number,
            mask: mask
        };
    }
    return new Error("Invalid CIDR");
}

/**
 * @param {CIDR} cidr
 * @returns {string}
 */
function stringifyCIDR(cidr) {
    return stringifyIp({
        number: cidr.address,
        version: cidr.version
    }) + "/" + cidr.mask;
}

export default class IPSet {

    /**
     * @param {4|6} version
     */
    constructor(version) {
        if (version !== 4 && version !== 6) {
            throw new Error("IP version must be 4 or 6");
        }
        this.version = version;
        this.highestBit = version == 4 ? 31 : 127;
        this.trie = { flag: false, c0: null, c1: null };
    }

    /**
     * @param {string} cidr
     */
    add(cidr) {
        const parsed = parseCIDR(cidr);
        if (parsed.version !== this.version) {
            throw new Error("IP version mismatch");
        }
        this.trieSetFlag(parsed, true);
    }

    /**
     * @param {string} cidr
     */
    remove(cidr) {
        const parsed = parseCIDR(cidr);
        if (parsed.version !== this.version) {
            throw new Error("IP version mismatch");
        }
        this.trieSetFlag(parsed, false);
    }

    /**
     * @param {string} cidr
     * @returns {boolean}
     */
    has(cidr) {
        return this.trieQuery(parseCIDR(cidr));
    }

    /**
     * @returns {string[]}
     */
    toCIDRs() {
        const result = [];
        for (const cidr of this.trieToCIDR()) {
            result.push(stringifyCIDR(cidr));
        }
        return result;
    }

    *trieToCIDR(node = this.trie, depth = 0, prefix = 0n) {
        if (node.flag === null) {
            yield* this.trieToCIDR(node.c0, depth + 1, prefix);
            yield* this.trieToCIDR(node.c1, depth + 1, prefix | (1n << BigInt(this.highestBit - depth)));

        } else if (node.flag === true) {
            yield {
                version: this.version,
                address: prefix,
                mask: depth
            };
        }
    }

    trieSetFlag(cidr, newFlag, node = this.trie, depth = 0) {
        if (depth === cidr.mask) {
            node.flag = newFlag;
            node.c0 = null;
            node.c1 = null;
            return;
        }

        if (node.flag === newFlag) {
            return;
        }

        if (node.flag !== null) {
            node.c0 = { flag: node.flag, c0: null, c1: null };
            node.c1 = { flag: node.flag, c0: null, c1: null };
            node.flag = null;
        }

        if (((cidr.address >> BigInt(this.highestBit - depth)) & 1n) === 0n) {
            this.trieSetFlag(cidr, newFlag, node.c0, depth + 1);
        } else {
            this.trieSetFlag(cidr, newFlag, node.c1, depth + 1);
        }

        if (node.c0.flag !== null && node.c0.flag === node.c1.flag) {
            node.flag = node.c0.flag;
            node.c0 = null;
            node.c1 = null;
        }
    }

    trieQuery(cidr, node = this.trie, depth = 0) {
        if (depth === cidr.mask) {
            return node.flag === true;
        }

        if (node.flag !== null) {
            return node.flag;
        }

        if (((cidr.address >> BigInt(this.highestBit - depth)) & 1n) === 0n) {
            return this.trieQuery(cidr, node.c0, depth + 1);
        } else {
            return this.trieQuery(cidr, node.c1, depth + 1);
        }
    }
};