openchain-sdk-yxl/lib/websocket/client.js

openchain sdk

'use strict';

let WebSocket;
if (typeof window !== 'undefined' && window.WebSocket) {
    WebSocket = window.WebSocket;
} else {
    WebSocket = require('ws');
}
const protobuf = require('protobufjs');
const path = require('path');
const EventEmitter = require('events');
const { Buffer } = require('buffer');

class WebSocketClient extends EventEmitter {
    constructor(options = {}) {
        super();
        this.url = options.url || 'ws://localhost:17002';
        this.autoReconnect = options.autoReconnect !== false;
        this.reconnectInterval = options.reconnectInterval || 5000;
        this.maxReconnectAttempts = options.maxReconnectAttempts || 10;
        this.subscriptions = new Set();
        this.reconnectAttempts = 0;
        this.connected = false;
        this.ws = null;
        this.root = null;

        // 初始化protobuf
        this.initProtobuf();
    }

    async initProtobuf() {
        try {
            const protoPath = path.join(__dirname, '../crypto/protobuf/bundle.json');
            this.root = await protobuf.load(protoPath);
        } catch (error) {
            this.emit('error', new Error('Failed to initialize protobuf: ' + error.message));
        }
    }

    connect() {
        if (this.ws) {
            return;
        }

        try {
            this.ws = new WebSocket(this.url);
            this.bindEvents();
        } catch (error) {
            this.emit('error', new Error('Failed to create WebSocket connection: ' + error.message));
        }
    }

    bindEvents() {
        if (!this.ws) return;

        const handleOpen = () => {
            this.connected = true;
            this.reconnectAttempts = 0;
            this.emit('connected');
            this.sendHello();
        };

        const handleMessage = (event) => {
            try {
                this.handleMessage(event.data || event);
            } catch (error) {
                this.emit('error', new Error('Failed to handle message: ' + error.message));
            }
        };

        const handleClose = () => {
            this.connected = false;
            this.emit('disconnected');
            this.handleReconnect();
        };

        const handleError = (error) => {
            this.emit('error', error);
        };

        if (typeof window !== 'undefined' && window.WebSocket) {
            this.ws.addEventListener('open', handleOpen);
            this.ws.addEventListener('message', handleMessage);
            this.ws.addEventListener('close', handleClose);
            this.ws.addEventListener('error', handleError);
        } else {
            this.ws.on('open', handleOpen);
            this.ws.on('message', handleMessage);
            this.ws.on('close', handleClose);
            this.ws.on('error', handleError);
        }
    }

    handleReconnect() {
        if (!this.autoReconnect || this.reconnectAttempts >= this.maxReconnectAttempts) {
            return;
        }

        setTimeout(() => {
            this.reconnectAttempts++;
            this.connect();
        }, this.reconnectInterval);
    }

    async sendHello() {
        const WsMessage = this.root.lookupType('protocol.WsMessage');
        const message = WsMessage.create({
            type: 1, // CHAIN_HELLO
            data: Buffer.from('Hello')
        });

        this.send(WsMessage.encode(message).finish());
    }

    async subscribeTx(addresses) {
        if (!Array.isArray(addresses) || addresses.length === 0) {
            throw new Error('Addresses must be a non-empty array');
        }

        if (addresses.length > 100) {
            throw new Error('Cannot subscribe to more than 100 addresses');
        }

        const WsMessage = this.root.lookupType('protocol.WsMessage');
        const ChainSubscribeTx = this.root.lookupType('protocol.ChainSubscribeTx');

        const subscribeMessage = ChainSubscribeTx.create({
            address: addresses
        });

        const message = WsMessage.create({
            type: 4, // CHAIN_SUBSCRIBE_TX
            data: ChainSubscribeTx.encode(subscribeMessage).finish()
        });

        addresses.forEach(addr => this.subscriptions.add(addr));
        this.send(WsMessage.encode(message).finish());
    }

    async handleMessage(data) {
        const WsMessage = this.root.lookupType('protocol.WsMessage');
        const message = WsMessage.decode(new Uint8Array(data));

        switch (message.type) {
            case 1: // CHAIN_HELLO response
                this.emit('hello', message);
                break;
            case 4: // CHAIN_SUBSCRIBE_TX response
                this.emit('subscription', message);
                break;
            case 5: // CHAIN_TX_ENV_STORE
                this.emit('transaction', message);
                break;
            default:
                this.emit('message', message);
        }
    }

    send(data) {
        if (!this.connected) {
            throw new Error('WebSocket is not connected');
        }

        try {
            this.ws.send(data);
        } catch (error) {
            this.emit('error', new Error('Failed to send message: ' + error.message));
        }
    }

    close() {
        if (this.ws) {
            this.autoReconnect = false;
            this.ws.close();
            this.ws = null;
        }
    }
}

module.exports = WebSocketClient;