farmcontrol-server/src/network/websocketScanner.js

import { EventEmitter } from 'events';
import os from 'os';
import { Worker } from 'worker_threads';
import { WebSocketScannerWorker } from './websocketScannerWorker.js';

const isBun = typeof globalThis.Bun !== 'undefined';
const workerUrl = new URL('./websocketScannerWorker.js', import.meta.url);

export class WebSocketScanner extends EventEmitter {
    constructor(options = {}) {
        super();
        this.scanning = false;
        this.stopped = false;
        this.workers = [];
        this.inProcessWorker = null;
        this.useWorkers = !isBun;
        this.maxThreads = this.useWorkers
            ? (options.maxThreads || os.cpus().length)
            : Math.max(1, Math.min(options.maxThreads || 16, 24));
        this.totalIPs = 0;
        this.scannedIPs = 0;
    }

    /**
     * Scans the local network for services on a specified port.
     * When loop is true, repeats until stopScan() is called.
     */
    async scanNetwork(port, protocol, objectType = 'printer', { loop = false } = {}) {
        this.stopped = false;
        this.scanning = true;
        let foundServices = [];

        try {
            do {
                await this.cleanupWorkers();
                if (this.stopped) {
                    break;
                }

                foundServices = await this.scanOnce(port, protocol, objectType);

                if (this.stopped) {
                    break;
                }

                this.emit('scanPassComplete', foundServices);
            } while (loop && !this.stopped);
        } finally {
            await this.cleanupWorkers();
            this.scanning = false;
        }

        return foundServices;
    }

    async scanOnce(port, protocol, objectType) {
        this.scannedIPs = 0;
        const foundServices = [];

        const { startIP, endIP } = this.getLocalNetworkRange();
        const start = this.ipToNumber(startIP);
        const end = this.ipToNumber(endIP);

        this.totalIPs = end - start + 1;

        if (this.useWorkers) {
            await this.scanWithWorkers(start, end, port, protocol, objectType, foundServices);
        } else {
            await this.scanInProcess(startIP, endIP, port, protocol, objectType, foundServices);
        }

        await this.cleanupWorkers();
        return foundServices;
    }

    handleWorkerMessage(message, foundServices) {
        if (this.stopped) {
            return;
        }

        switch (message.type) {
            case 'serviceFound':
                foundServices.push({ ip: message.ip, hostname: message.hostname });
                this.emit('serviceFound', { ip: message.ip, hostname: message.hostname });
                break;
            case 'scanProgress':
                this.scannedIPs += message.increment;
                this.emit('scanProgress', {
                    currentIP: message.currentIP,
                    progress: (this.scannedIPs / this.totalIPs) * 100
                });
                break;
            default:
                break;
        }
    }

    async scanInProcess(startIP, endIP, port, protocol, objectType, foundServices) {
        const worker = new WebSocketScannerWorker((message) => {
            this.handleWorkerMessage(message, foundServices);
        });
        this.inProcessWorker = worker;

        await worker.scanRange(startIP, endIP, port, protocol, objectType, {
            concurrency: this.maxThreads
        });
    }

    async scanWithWorkers(start, end, port, protocol, objectType, foundServices) {
        const ipsPerWorker = Math.ceil(this.totalIPs / this.maxThreads);
        const workerPromises = [];

        for (let i = 0; i < this.maxThreads; i++) {
            if (this.stopped) {
                break;
            }

            const workerStart = start + (i * ipsPerWorker);
            const workerEnd = Math.min(workerStart + ipsPerWorker - 1, end);

            if (workerStart > end) {
                break;
            }

            const worker = new Worker(workerUrl);
            this.workers.push(worker);

            const workerPromise = new Promise((resolve) => {
                let settled = false;
                const finish = (results) => {
                    if (settled) {
                        return;
                    }
                    settled = true;
                    resolve(results || []);
                };

                worker.on('message', (message) => {
                    if (this.stopped) {
                        finish([]);
                        return;
                    }

                    if (message.type === 'scanComplete') {
                        finish(message.results);
                        return;
                    }

                    this.handleWorkerMessage(message, foundServices);
                });
                worker.on('error', () => finish([]));
                worker.on('exit', () => finish([]));
            });

            worker.postMessage({
                type: 'scan',
                startIP: this.numberToIP(workerStart),
                endIP: this.numberToIP(workerEnd),
                port,
                protocol,
                objectType
            });

            workerPromises.push(workerPromise);
        }

        await Promise.all(workerPromises);
    }

    async cleanupWorkers() {
        if (this.inProcessWorker) {
            this.inProcessWorker.stop();
            this.inProcessWorker = null;
        }

        const workers = this.workers.splice(0, this.workers.length);
        await Promise.all(workers.map((worker) => this.terminateWorker(worker)));
    }

    terminateWorker(worker) {
        return Promise.race([
            Promise.resolve().then(() => worker.terminate()),
            new Promise((resolve) => setTimeout(resolve, 500))
        ]).catch(() => {});
    }

    getLocalNetworkRange() {
        const interfaces = os.networkInterfaces();
        let localIP = null;
        let subnetMask = null;

        for (const name of Object.keys(interfaces)) {
            for (const iface of interfaces[name]) {
                const isIPv4 = iface.family === 'IPv4' || iface.family === 4;
                if (isIPv4 && !iface.internal) {
                    localIP = iface.address;
                    subnetMask = iface.netmask;
                    break;
                }
            }
            if (localIP) {
                break;
            }
        }

        if (!localIP) {
            throw new Error('Could not determine local network IP address');
        }

        const ipNum = this.ipToNumber(localIP);
        const maskNum = this.ipToNumber(subnetMask);
        const networkNum = ipNum & maskNum;
        const broadcastNum = networkNum | (~maskNum >>> 0);

        return {
            startIP: this.numberToIP(networkNum + 1),
            endIP: this.numberToIP(broadcastNum - 1)
        };
    }

    ipToNumber(ip) {
        return ip.split('.')
            .reduce((acc, octet) => (acc << 8) + parseInt(octet), 0) >>> 0;
    }

    numberToIP(num) {
        return [
            (num >>> 24) & 255,
            (num >>> 16) & 255,
            (num >>> 8) & 255,
            num & 255
        ].join('.');
    }

    async stopScan() {
        this.stopped = true;
        this.scanning = false;
        await this.cleanupWorkers();
    }
}