network chunk v7

2024-11-13 10:11:11 +02:00
parent 60e71c3e2a
commit 97cc54066d
3 changed files with 126 additions and 182 deletions
@@ -4,15 +4,14 @@ import { MessageHandler } from '../message_handler';
 import { SocketCommunicatorBase } from './socket_communicator_base';
 import { OperationHandler } from '../operations_base/operation_handler';

-const MAX_CHUNK_SIZE = 2048;  // Define chunk size
-
 export class TcpServerCommunicator extends SocketCommunicatorBase {
    private readonly socket: Socket;
    private privateKey: string | null;
    private publicKey: string | null;
    private aesKey: Buffer | null;
    private aesIv: Buffer | null;
-    private chunkBuffers: { [messageId: string]: string[] };  // Buffer for reassembling chunks
+    private chunkBuffers: { [messageId: string]: string[] };
+    private networkSpeed: number | null = null;  // Estimated network speed in bytes/ms

    constructor(socket: Socket, ip: string, port: number, operationHandler: OperationHandler) {
        super(ip, port, operationHandler);
@@ -21,11 +20,10 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
        this.publicKey = null;  // Client public key will be set later
        this.aesKey = null;
        this.aesIv = null;
-        this.chunkBuffers = {};  // Buffer for reassembling incoming messages
-        this.generateKeyPair();  // Generate RSA key pair for encryption
+        this.chunkBuffers = {};
+        this.generateKeyPair();
    }

-    // Generate RSA key pair (public and private keys)
    generateKeyPair(): void {
        const { privateKey, publicKey } = generateKeyPairSync('rsa', {
            modulusLength: 2048,
@@ -36,52 +34,42 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
        this.publicKey = publicKey;
    }

-    // Send the server's public key to the client
    async sendPublicKey(): Promise<void> {
        if (!this.publicKey) {
            throw new Error('Public key is not available. Please generate RSA key pair.');
        }
-
        await this.sendChunkedMessage('SET_PUBLIC_KEY', { publicKey: this.publicKey });
    }

-    // Generate AES key and IV, then send them to the client
    async sendAesKey(): Promise<void> {
-        this.aesKey = randomBytes(32);  // 256-bit AES key
-        this.aesIv = randomBytes(16);   // AES IV
-
+        this.aesKey = randomBytes(32);
+        this.aesIv = randomBytes(16);
        const aesKeyBase64 = this.aesKey.toString('base64');
        const aesIvBase64 = this.aesIv.toString('base64');
-
        await this.sendChunkedMessage('SET_AES_KEY', { aesKey: aesKeyBase64, aesIv: aesIvBase64 });
    }

-    // Encrypt a message with the server's private key (RSA encryption)
    private encryptWithRsa(message: string): string {
        if (!this.privateKey) throw new Error('Server private key not set.');
-
        return privateEncrypt(
            {
                key: this.privateKey,
-                padding: constants.RSA_PKCS1_PADDING,  // PKCS1 padding
+                padding: constants.RSA_PKCS1_PADDING,
            },
            Buffer.from(message)
        ).toString('base64');
    }

-    // Decrypt AES-encrypted messages
    private decryptWithAes(encryptedMessage: string): string {
        if (!this.aesKey || !this.aesIv) {
            throw new Error('AES key not set.');
        }
-
        const decipher = createDecipheriv('aes-256-cbc', this.aesKey, this.aesIv);
        let decrypted = decipher.update(Buffer.from(encryptedMessage, 'base64'));
        decrypted = Buffer.concat([decrypted, decipher.final()]);
        return decrypted.toString('utf-8');
    }

-    // Encrypt a message with AES
    private encryptWithAes(message: string): string {
        if (!this.aesKey || !this.aesIv) {
            throw new Error('AES key or IV is not set.');
@@ -92,55 +80,58 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
        return encrypted.toString('base64');
    }

-    // Handle incoming chunks of data
    async handleIncomingChunk(data: Buffer): Promise<void> {
        const incomingMessage = data.toString().trim();
-
-        // Extract header and chunk content from incoming data
        const [headerJson, chunkContent] = incomingMessage.split('|');
        const header = JSON.parse(headerJson);
-
-        console.log(`\n\n${incomingMessage}\n\n`);
-
-        // Initialize chunk array if this is the first chunk for this messageId
        if (!this.chunkBuffers[header.messageId]) {
-            this.chunkBuffers[header.messageId] = new Array(header.totalChunks);
+            this.chunkBuffers[header.messageId] = new Array(header.totalChunks).fill(undefined);
        }
-
-        // Place the chunk in the correct position in the chunk buffer
-        this.chunkBuffers[header.messageId][header.sequenceNumber] = chunkContent;
-        if(header.sequenceNumber === header.totalChunks) {
-            await this.processCompleteMessage();
+        this.chunkBuffers[header.messageId][header.sequenceNumber - 1] = chunkContent;
+        if (this.chunkBuffers[header.messageId].every((chunk) => chunk !== undefined)) {
+            await this.processCompleteMessage(header.messageId);
        }
    }

-    // Process the complete message when EOM is received
-    private async processCompleteMessage(): Promise<void> {
-        for (const [messageId, chunks] of Object.entries(this.chunkBuffers)) {
-            if (chunks.every((chunk) => chunk !== undefined)) {
-                // Join all chunks to form the full message
-                const fullMessage = chunks.join('');
-
-                // Handle the completed and possibly decrypted message
-                await this.handleIncomingMessage(fullMessage);
-
-                // Clean up buffer after processing
-                delete this.chunkBuffers[messageId];
-            }
+    private async processCompleteMessage(messageId: string): Promise<void> {
+        const chunks = this.chunkBuffers[messageId];
+        if (chunks && chunks.every((chunk) => chunk !== undefined)) {
+            const fullMessage = chunks.join('');
+            await this.handleIncomingMessage(fullMessage);
+            delete this.chunkBuffers[messageId];
        }
    }

-    // Send chunked message
+    private async scanNetworkSpeed(): Promise<number> {
+        const testMessage = 'PING_TEST'.repeat(100); // A test message of known size
+        const startTime = Date.now();
+        await this.writeToSocket(testMessage);
+        await new Promise((resolve) => this.socket.once('data', resolve));
+        const endTime = Date.now();
+        const duration = endTime - startTime;  // Time in ms
+        this.networkSpeed = testMessage.length / duration;  // Bytes/ms
+        return this.networkSpeed;
+    }
+
+    private calculateOptimalChunkSize(messageLength: number): number {
+        let chunkSize = this.networkSpeed ? Math.min(Math.floor(this.networkSpeed * 100), 1024) : 1024;
+        while (messageLength % chunkSize !== 0 && chunkSize > 0) {
+            chunkSize -= 4;
+        }
+        return chunkSize || 1024;
+    }
+
    async sendChunkedMessage(
        operationCode: string,
        metaInfo?: { [key: string]: any },
        fileContent?: Buffer
    ): Promise<void> {
+        if (!this.networkSpeed) {
+            await this.scanNetworkSpeed();
+        }
        const message = MessageHandler.formatMessage(operationCode, metaInfo, fileContent);
        let outgoingMessage: string;
-
-        // Encrypt or format message based on the operation code
-        switch(operationCode) {
+        switch (operationCode) {
            case 'SET_PUBLIC_KEY':
                outgoingMessage = Buffer.from(message, 'utf-8').toString('base64');
                break;
@@ -150,20 +141,9 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
            default:
                outgoingMessage = this.encryptWithAes(message);
        }
-
-        // Determine optimal chunk size (multiple of 4 and ≤ 1024 to align with base64 encoding)
-        let optimalChunkSize = MAX_CHUNK_SIZE;
-        while (outgoingMessage.length % optimalChunkSize !== 0 && optimalChunkSize > 0) {
-            optimalChunkSize -= 4;
-        }
-
-        if (optimalChunkSize === 0) optimalChunkSize = MAX_CHUNK_SIZE;  // Fallback in case of odd alignment
-
-        // Calculate total chunks and generate unique message ID
+        const optimalChunkSize = this.calculateOptimalChunkSize(outgoingMessage.length);
        const totalChunks = Math.ceil(outgoingMessage.length / optimalChunkSize);
        const messageId = Date.now().toString();
-
-        // Send each chunk as a string with delay to manage network flow
        for (let i = 0; i < totalChunks; i++) {
            const chunk = outgoingMessage.slice(i * optimalChunkSize, (i + 1) * optimalChunkSize);
            const chunkHeader = JSON.stringify({
@@ -171,27 +151,20 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
                sequenceNumber: i + 1,
                totalChunks,
            });
-
            const chunkWithHeader = `${chunkHeader}|${chunk}`;
-
-            // Write the chunk string directly to the socket
            await this.writeToSocket(chunkWithHeader);
-            await new Promise((resolve) => setTimeout(resolve, 300));  // Simulate network delay
+            await new Promise((resolve) => setTimeout(resolve, 300));
        }
    }

-    // Handle incoming message (decrypt with AES if available)
    async handleIncomingMessage(incomingMessage: string): Promise<void> {
        let messageToProcess = incomingMessage;
-
        if (this.aesKey && this.aesIv) {
            messageToProcess = this.decryptWithAes(incomingMessage);
        }
-
        this.handlerResult = await this.operationHandler.handleOperation(messageToProcess);
    }

-    // Write message to socket
    private writeToSocket(message: string): Promise<void> {
        return new Promise((resolve, reject) => {
            this.socket.write(message, (err: any) => {