@cliz/inlets/lib/core/server/utils.js

Cloud Native Tunnel

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.createBandwidthLimiter = exports.createTrafficStatsContainer = exports.Notification = exports.createContainer = exports.createCallbackCtainer = exports.createHostContainer = exports.shortUUID = void 0;
const doreamon_1 = require("@zodash/doreamon");
const shorturl_1 = require("@zodash/shorturl");
const nobot_1 = require("@znode/nobot");
function shortUUID() {
    return (0, shorturl_1.default)(doreamon_1.default.uuid());
}
exports.shortUUID = shortUUID;
function createHostContainer() {
    const hosts = {};
    function get(id) {
        if (id) {
            const _id = id.toLowerCase();
            return hosts[_id];
        }
        return hosts;
    }
    const has = (id) => {
        return !!hosts[id];
    };
    const bindWs = (wsSocket, subDomain) => {
        const uid = (subDomain || shortUUID()).toLowerCase();
        hosts[uid] = {
            wsSocket,
        };
        return uid;
    };
    const unbindWs = (id) => {
        const _id = id.toLowerCase();
        delete hosts[_id];
    };
    const bindTcp = (id, tcp) => {
        const _id = id.toLowerCase();
        if (!hosts[_id])
            return;
        hosts[_id].tcpSocket = tcp;
    };
    return {
        get,
        has,
        bindWs,
        unbindWs,
        bindTcp,
    };
}
exports.createHostContainer = createHostContainer;
function createCallbackCtainer() {
    const callbacks = {};
    const get = (tcpId, requestId) => {
        if (callbacks[tcpId] && callbacks[tcpId][requestId]) {
            return callbacks[tcpId][requestId];
        }
        return null;
    };
    const set = (tcpId, requestId, callback) => {
        if (!callbacks[tcpId]) {
            callbacks[tcpId] = {};
        }
        callbacks[tcpId][requestId] = callback;
    };
    const remove = (tcpId) => {
        delete callbacks[tcpId];
    };
    return {
        get,
        set,
        remove,
    };
}
exports.createCallbackCtainer = createCallbackCtainer;
function createContainer() {
    const containers = {};
    function create(id, token, wsSocket, authentication) {
        const destroy = () => {
            const c = get(id);
            if (!c) {
                throw new Error(`invalid token for container`);
            }
            c.sourceServer && c.sourceServer.close();
            remove(id);
        };
        containers[id] = {
            wsSocket,
            token,
            type: authentication.type,
            port: authentication.port,
            sourcePort: authentication.tunnelPort,
            version: authentication.version,
            authType: authentication.authType,
            clientId: authentication.clientId,
            tunnelPort: authentication.tunnelPort,
            requests: {},
            signature: authentication.signature,
            clientTimestamp: authentication.timestamp,
            destroy,
        };
        wsSocket.on('disconnect', () => {
            destroy();
        });
    }
    function get(id) {
        return containers[id];
    }
    function set(id, key, value) {
        let one = get(id);
        if (!one) {
            throw new Error(`invalid id for container`);
        }
        one[key] = value;
    }
    function remove(id) {
        delete containers[id];
    }
    function registerRequest(containerId, requestId, sourceSocket) {
        const container = containers[containerId];
        container.requests[requestId] = sourceSocket;
    }
    function connectRequest(containerId, requestId, targetSocket) {
        const container = containers[containerId];
        const sourceSocket = container.requests[requestId];
        delete container.requests[requestId];
        if (requestId !== sourceSocket.requestId || requestId !== targetSocket.requestId) {
            throw new Error(`[connectRequest] request id 无法对应(current: ${requestId}, sourceSocket: ${sourceSocket.requestId}, targetSocket: ${targetSocket.requestId})`);
        }
        targetSocket.pipe(sourceSocket);
        sourceSocket.pipe(targetSocket);
    }
    return {
        create,
        get,
        set,
        remove,
        registerRequest,
        connectRequest,
    };
}
exports.createContainer = createContainer;
class Notification {
    constructor(config) {
        this.config = config;
    }
    async notify(title, message) {
        if (!this.config) {
            return;
        }
        const content = !Array.isArray(message)
            ? message
            : message.join('\n');
        return (0, nobot_1.sendMessage)(this.config.provider, this.config.url, {
            title,
            content,
        });
    }
}
exports.Notification = Notification;
function createTrafficStatsContainer() {
    const stats = {
        global: {
            uploadBytes: 0,
            downloadBytes: 0,
            connections: 0,
            requests: 0,
            startTime: Date.now(),
            lastUpdateTime: Date.now(),
        },
        byClientId: {},
    };
    function getOrCreateClientStats(clientId) {
        if (!clientId) {
            return stats.global;
        }
        if (!stats.byClientId[clientId]) {
            stats.byClientId[clientId] = {
                uploadBytes: 0,
                downloadBytes: 0,
                connections: 0,
                requests: 0,
                startTime: Date.now(),
                lastUpdateTime: Date.now(),
            };
        }
        return stats.byClientId[clientId];
    }
    function addUploadBytes(clientId, bytes) {
        const clientStats = getOrCreateClientStats(clientId);
        clientStats.uploadBytes += bytes;
        clientStats.lastUpdateTime = Date.now();
        if (clientId) {
            stats.global.uploadBytes += bytes;
            stats.global.lastUpdateTime = Date.now();
        }
    }
    function addDownloadBytes(clientId, bytes) {
        const clientStats = getOrCreateClientStats(clientId);
        clientStats.downloadBytes += bytes;
        clientStats.lastUpdateTime = Date.now();
        if (clientId) {
            stats.global.downloadBytes += bytes;
            stats.global.lastUpdateTime = Date.now();
        }
    }
    function addConnection(clientId) {
        const clientStats = getOrCreateClientStats(clientId);
        clientStats.connections += 1;
        clientStats.lastUpdateTime = Date.now();
        if (clientId) {
            stats.global.connections += 1;
            stats.global.lastUpdateTime = Date.now();
        }
    }
    function addRequest(clientId) {
        const clientStats = getOrCreateClientStats(clientId);
        clientStats.requests += 1;
        clientStats.lastUpdateTime = Date.now();
        if (clientId) {
            stats.global.requests += 1;
            stats.global.lastUpdateTime = Date.now();
        }
    }
    function getStats(clientId) {
        if (clientId) {
            return stats.byClientId[clientId] || null;
        }
        return stats;
    }
    function reset(clientId) {
        if (clientId) {
            delete stats.byClientId[clientId];
            return;
        }
        stats.global = {
            uploadBytes: 0,
            downloadBytes: 0,
            connections: 0,
            requests: 0,
            startTime: Date.now(),
            lastUpdateTime: Date.now(),
        };
        stats.byClientId = {};
    }
    function setStartTime(clientId) {
        const clientStats = getOrCreateClientStats(clientId);
        clientStats.startTime = Date.now();
        clientStats.endTime = undefined;
        clientStats.lastUpdateTime = Date.now();
    }
    function setEndTime(clientId) {
        const clientStats = getOrCreateClientStats(clientId);
        if (!clientStats.endTime) {
            clientStats.endTime = Date.now();
            clientStats.lastUpdateTime = Date.now();
        }
    }
    function formatStats(clientId) {
        const stats = getStats(clientId || undefined);
        if (!stats || (clientId && typeof stats === 'object' && 'global' in stats)) {
            return 'No stats available';
        }
        const clientStats = clientId ? stats : stats.global;
        if (!clientStats) {
            return 'No stats available';
        }
        const formatBytes = (bytes) => {
            if (bytes < 1024)
                return `${bytes} B`;
            if (bytes < 1024 * 1024)
                return `${(bytes / 1024).toFixed(2)} KB`;
            if (bytes < 1024 * 1024 * 1024)
                return `${(bytes / (1024 * 1024)).toFixed(2)} MB`;
            return `${(bytes / (1024 * 1024 * 1024)).toFixed(2)} GB`;
        };
        const totalBytes = clientStats.uploadBytes + clientStats.downloadBytes;
        const endTime = clientStats.endTime || Date.now();
        const duration = Math.floor((endTime - clientStats.startTime) / 1000);
        const hours = Math.floor(duration / 3600);
        const minutes = Math.floor((duration % 3600) / 60);
        const seconds = duration % 60;
        const durationStr = hours > 0
            ? `${hours}h ${minutes}m ${seconds}s`
            : minutes > 0
                ? `${minutes}m ${seconds}s`
                : `${seconds}s`;
        return [
            `上传: ${formatBytes(clientStats.uploadBytes)}`,
            `下载: ${formatBytes(clientStats.downloadBytes)}`,
            `总计: ${formatBytes(totalBytes)}`,
            `请求数: ${clientStats.requests}`,
            `连接数: ${clientStats.connections}`,
            `运行时长: ${durationStr}`,
        ].join(', ');
    }
    return {
        addUploadBytes,
        addDownloadBytes,
        addConnection,
        addRequest,
        getStats,
        reset,
        setStartTime,
        setEndTime,
        formatStats,
    };
}
exports.createTrafficStatsContainer = createTrafficStatsContainer;
class TokenBucket {
    constructor(capacity, refillRate) {
        this.capacity = capacity;
        this.refillRate = refillRate;
        this.tokens = capacity;
        this.lastRefill = Date.now();
    }
    canConsume(tokens) {
        this.refill();
        return this.tokens >= tokens;
    }
    consume(tokens) {
        this.refill();
        if (this.tokens >= tokens) {
            this.tokens -= tokens;
            return true;
        }
        return false;
    }
    getAvailableTokens() {
        this.refill();
        return this.tokens;
    }
    refill() {
        const now = Date.now();
        const elapsed = (now - this.lastRefill) / 1000;
        const tokensToAdd = elapsed * this.refillRate;
        this.tokens = Math.min(this.capacity, this.tokens + tokensToAdd);
        this.lastRefill = now;
    }
}
function createBandwidthLimiter(limits = { byClientId: {} }) {
    var _a, _b;
    const globalUploadBucket = ((_a = limits.global) === null || _a === void 0 ? void 0 : _a.uploadBytesPerSecond)
        ? new TokenBucket(limits.global.uploadBytesPerSecond, limits.global.uploadBytesPerSecond)
        : null;
    const globalDownloadBucket = ((_b = limits.global) === null || _b === void 0 ? void 0 : _b.downloadBytesPerSecond)
        ? new TokenBucket(limits.global.downloadBytesPerSecond, limits.global.downloadBytesPerSecond)
        : null;
    const clientUploadBuckets = {};
    const clientDownloadBuckets = {};
    function getClientUploadBucket(clientId) {
        var _a;
        if (!clientId) {
            return null;
        }
        const limit = (_a = limits.byClientId[clientId]) === null || _a === void 0 ? void 0 : _a.uploadBytesPerSecond;
        if (!limit) {
            return null;
        }
        if (!clientUploadBuckets[clientId]) {
            clientUploadBuckets[clientId] = new TokenBucket(limit, limit);
        }
        return clientUploadBuckets[clientId];
    }
    function getClientDownloadBucket(clientId) {
        var _a;
        if (!clientId) {
            return null;
        }
        const limit = (_a = limits.byClientId[clientId]) === null || _a === void 0 ? void 0 : _a.downloadBytesPerSecond;
        if (!limit) {
            return null;
        }
        if (!clientDownloadBuckets[clientId]) {
            clientDownloadBuckets[clientId] = new TokenBucket(limit, limit);
        }
        return clientDownloadBuckets[clientId];
    }
    function checkUpload(clientId, bytes) {
        if (!clientId) {
            if (globalUploadBucket) {
                return globalUploadBucket.consume(bytes);
            }
            return true;
        }
        const clientBucket = getClientUploadBucket(clientId);
        if (clientBucket && !clientBucket.canConsume(bytes)) {
            return false;
        }
        if (globalUploadBucket && !globalUploadBucket.canConsume(bytes)) {
            return false;
        }
        if (globalUploadBucket) {
            if (!globalUploadBucket.consume(bytes)) {
                return false;
            }
        }
        if (clientBucket) {
            if (!clientBucket.consume(bytes)) {
                return false;
            }
        }
        return true;
    }
    function checkDownload(clientId, bytes) {
        if (!clientId) {
            if (globalDownloadBucket) {
                return globalDownloadBucket.consume(bytes);
            }
            return true;
        }
        const clientBucket = getClientDownloadBucket(clientId);
        if (clientBucket && !clientBucket.canConsume(bytes)) {
            return false;
        }
        if (globalDownloadBucket && !globalDownloadBucket.canConsume(bytes)) {
            return false;
        }
        if (globalDownloadBucket) {
            if (!globalDownloadBucket.consume(bytes)) {
                return false;
            }
        }
        if (clientBucket) {
            if (!clientBucket.consume(bytes)) {
                return false;
            }
        }
        return true;
    }
    function setClientLimit(clientId, limit) {
        if (!limits.byClientId[clientId]) {
            limits.byClientId[clientId] = {};
        }
        limits.byClientId[clientId] = { ...limits.byClientId[clientId], ...limit };
        if (limit.uploadBytesPerSecond !== undefined) {
            clientUploadBuckets[clientId] = new TokenBucket(limit.uploadBytesPerSecond, limit.uploadBytesPerSecond);
        }
        if (limit.downloadBytesPerSecond !== undefined) {
            clientDownloadBuckets[clientId] = new TokenBucket(limit.downloadBytesPerSecond, limit.downloadBytesPerSecond);
        }
    }
    function removeClientLimit(clientId) {
        delete limits.byClientId[clientId];
        delete clientUploadBuckets[clientId];
        delete clientDownloadBuckets[clientId];
    }
    return {
        checkUpload,
        checkDownload,
        setClientLimit,
        removeClientLimit,
    };
}
exports.createBandwidthLimiter = createBandwidthLimiter;