network chunk v14

CEO/src/network/operations_custom/general_operations.ts

@@ -7,7 +7,6 @@ export class GeneralOperations implements OperationPlugin {
     public static readonly operationCodes = {
         OK: 'OK',
         ERR: 'ERR',
-        END: 'END',
         HEARTBEAT: 'HEARTBEAT',
         ALIVE: 'ALIVE',
         SET_PUBLIC_KEY: 'SET_PUBLIC_KEY',

CEO/src/network/operations_custom/user_to_user_operations.ts

@@ -116,6 +116,8 @@ export class UserToUserOperations implements OperationPlugin {
             const appInfo = await jsonManager.readValue('shareDirectory');
             const shareDirectory = appInfo?.path || '';
 
+            console.log(`\n\nShare directory: ${shareDirectory}\n\n`);
+
             if (!shareDirectory) {
                 return { operationCode: operationCodes.ERR, metaInfo: { message: 'Share directory missing.' }};
             }

CEO/src/network/socket_communicator/socket_communicator_base.ts

@@ -1,21 +1,197 @@
 import { ParsedMessage } from '../message_handler';
 import { OperationHandler } from '../operations_base/operation_handler';
+import ping from "ping";
+import {
+    constants,
+    createCipheriv,
+    createDecipheriv,
+    generateKeyPairSync,
+    privateEncrypt,
+    publicDecrypt,
+    randomBytes
+} from "crypto";
 
 export abstract class SocketCommunicatorBase {
     protected readonly ip: string;
     protected readonly port: number;
     protected readonly operationHandler: OperationHandler;
     protected handlerResult: ParsedMessage | null;
+    protected networkSpeed: number | null = null;
+
+    protected chunkBuffers: { [messageId: string]: string[] };
+    protected readonly EOP = '<EOP>';
+
+    protected privateKey: string | null;
+    protected publicKey: string | null;
+    protected aesKey: Buffer | null;
+    protected aesIv: Buffer | null;
+
+    private incompleteChunkBuffer: string = '';
 
     protected constructor(ip: string, port: number, operationHandler: OperationHandler) {
         this.ip = ip;
         this.port = port;
         this.operationHandler = operationHandler
         this.handlerResult = null;
+        this.chunkBuffers = {};
+
+        this.privateKey = null;
+        this.publicKey = null;
+        this.aesKey = null;
+        this.aesIv = null;
     }
 
     // Getter for the handler result
     getHandlerResult(): ParsedMessage | null {
-        return this.handlerResult;
+        const result = this.handlerResult;
+        this.handlerResult = null;
+        return result;
+    }
+
+    protected generateKeyPair(): void {
+        const { privateKey, publicKey } = generateKeyPairSync('rsa', {
+            modulusLength: 2048,
+            publicKeyEncoding: { type: 'spki', format: 'pem' },
+            privateKeyEncoding: { type: 'pkcs8', format: 'pem' },
+        });
+        this.privateKey = privateKey;
+        this.publicKey = publicKey;
+    }
+
+    protected generateAesKey(): void {
+        this.aesKey = randomBytes(32);
+        this.aesIv = randomBytes(16);
+    }
+
+    protected encryptWithAes(message: string): string {
+        if (!this.aesKey || !this.aesIv) {
+            throw new Error('AES key or IV is not set.');
+        }
+        const cipher = createCipheriv('aes-256-cbc', this.aesKey, this.aesIv);
+        let encrypted = cipher.update(message, 'utf-8');
+        encrypted = Buffer.concat([encrypted, cipher.final()]);
+        return encrypted.toString('base64');
+    }
+
+    protected decryptWithAes(encryptedMessage: string): string {
+        if (!this.aesKey || !this.aesIv) {
+            throw new Error('AES key or IV is 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');
+    }
+
+    protected decryptWithRsa(message: string): string {
+        if (!this.publicKey) {
+            throw new Error('Server public key not set.');
+        }
+        try {
+            const encryptedMessage = Buffer.from(message, 'base64');
+            const decrypted = publicDecrypt(
+                {
+                    key: this.publicKey,
+                    padding: constants.RSA_PKCS1_PADDING,
+                },
+                encryptedMessage
+            );
+            return decrypted.toString('utf-8');
+        } catch (error) {
+            throw new Error('Failed to decrypt RSA message.');
         }
     }
+
+    protected 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,
+            },
+            Buffer.from(message)
+        ).toString('base64');
+    }
+
+    protected async scanNetworkLatency(): Promise<number> {
+        const targetIp = this.ip; // Use the IP from the superclass
+
+        try {
+            const response = await ping.promise.probe(targetIp);
+
+            if (!response.alive || response.time === "unknown") {
+                console.warn(`Ping failed to reach ${targetIp}. Using default network speed.`);
+                return 200; // Default latency in ms if ping fails
+            }
+
+            return response.time; // Latency in ms from ping response
+        } catch (error: any) {
+            console.error(`Ping error: ${error.message}. Using default network speed.`);
+            return 200; // Default latency in ms if an error occurs
+        }
+    }
+
+    // Calculate optimal chunk size based on network latency, with fallback if necessary
+    protected async calculateOptimalChunkSize(messageLength: number): Promise<number> {
+        const latency = await this.scanNetworkLatency();
+        this.networkSpeed = latency > 0 ? 1000 / latency : 1; // Speed in bytes/ms based on latency
+
+        // Calculate initial chunk size based on latency (bounded between 512 and 1024 bytes)
+        let chunkSize = Math.min(Math.max(512, Math.floor(5000 / this.networkSpeed)), 1024);
+
+        // Adjust chunk size for base64 alignment (multiple of 4)
+        while (messageLength % chunkSize !== 0 && chunkSize > 0) {
+            chunkSize -= 4;
+        }
+
+        return chunkSize || 1024; // Fallback to 1024 if alignment adjustment results in 0
+    }
+
+    async handleIncomingChunk(data: Buffer): Promise<void> {
+        // Append incoming data to the incomplete buffer
+        this.incompleteChunkBuffer += data.toString();
+
+        // Split the buffer by <EOP> to separate complete and incomplete messages
+        const messages = this.incompleteChunkBuffer.split(this.EOP);
+
+        // Save the last item back to the buffer if it's incomplete (no <EOP> at the end)
+        this.incompleteChunkBuffer = messages.pop() || "";
+
+        // Process each complete message in the split results
+        for (const incomingMessage of messages) {
+            try {
+                const [headerJson, chunkContent] = incomingMessage.split('|');
+                const header = JSON.parse(headerJson);
+
+                // Initialize an array for chunks if it's the first chunk for this messageId
+                if (!this.chunkBuffers[header.messageId]) {
+                    this.chunkBuffers[header.messageId] = [];
+                }
+
+                // Store the chunk in the correct position based on sequenceNumber (1-based indexing)
+                this.chunkBuffers[header.messageId][header.sequenceNumber - 1] = chunkContent;
+
+                console.log(`Received chunk: ${incomingMessage}`);
+                console.log(`Chunks received so far: ${this.chunkBuffers[header.messageId].filter(chunk => chunk !== undefined).length}`);
+
+                // Check if all chunks have been received
+                if (this.chunkBuffers[header.messageId].filter(chunk => chunk !== undefined).length === header.totalChunks) {
+                    // Join all chunks to form the full message
+                    const fullMessage = this.chunkBuffers[header.messageId].join('');
+
+                    // Process the complete message
+                    await this.handleIncomingMessage(fullMessage);
+
+                    // Clear the chunk buffer for this messageId
+                    delete this.chunkBuffers[header.messageId];
+                }
+            } catch (error: any) {
+                console.error(`Error handling chunk: ${error.message}`);
+            }
+        }
+    }
+
+    abstract handleIncomingMessage(incomingMessage: string): Promise<void>;
+
+    abstract sendMessage(message: string): Promise<void>;
+}

CEO/src/network/socket_communicator/tcp_client_communicator.ts

@@ -1,158 +1,46 @@
 import { Socket } from 'net';
-import { createCipheriv, createDecipheriv, constants, publicDecrypt } from 'crypto';
-import {MessageHandler, ParsedMessage} from '../message_handler';
 import { SocketCommunicatorBase } from './socket_communicator_base';
 import { OperationHandler } from '../operations_base/operation_handler';
 import { operationCodes } from '../operation_codes';
-
-const END_OF_MESSAGE = '<EOM>';  // Unique marker for the end of message
-const CHUNK_SIZE = 1024;  // Define chunk size
+import {MessageHandler} from "../message_handler";
 
 export class TcpClientCommunicator extends SocketCommunicatorBase {
     private readonly socket: Socket;
-    private aesKey: Buffer | null;
-    private aesIv: Buffer | null;
-    private messageBuffer: string;
-    private serverPublicKey: string | null;
     private isAesKeySetFlag: boolean;
-    private chunkBuffers: { [messageId: string]: string[] };  // Buffer for reassembling chunks
 
     constructor(socket: Socket, ip: string, port: number, operationHandler: OperationHandler) {
-        super(ip, port, operationHandler);  // Call parent constructor
+        super(ip, port, operationHandler);
         this.socket = socket;
         this.aesKey = null;
         this.aesIv = null;
-        this.serverPublicKey = null;
-        this.messageBuffer = '';  // Buffer for message reassembly
         this.isAesKeySetFlag = false;
-        this.chunkBuffers = {};  // Initialize chunk buffer for reassembling messages
+        this.chunkBuffers = {};
+    }
+
+    isAesKeySet(): boolean {
+        return this.isAesKeySetFlag;
     }
 
     setServerPublicKey(publicKey: string): void {
-        this.serverPublicKey = publicKey;
+        console.log('\n\nSetting server public key\n\n');
+        this.publicKey = publicKey;
     }
 
     setAesKey(aesKey: string, aesIv: string): void {
+        console.log('\n\nSetting AES key\n\n');
         this.aesKey = Buffer.from(aesKey, 'base64');
         this.aesIv = Buffer.from(aesIv, 'base64');
     }
 
-    private encryptWithAes(message: string): string {
-        if (!this.aesKey || !this.aesIv) {
-            throw new Error('AES key or IV is not set.');
-        }
-        const cipher = createCipheriv('aes-256-cbc', this.aesKey, this.aesIv);
-        let encrypted = cipher.update(message, 'utf-8');
-        encrypted = Buffer.concat([encrypted, cipher.final()]);
-        return encrypted.toString('base64');
-    }
-
-    private decryptWithAes(encryptedMessage: string): string {
-        if (!this.aesKey || !this.aesIv) {
-            throw new Error('AES key or IV is 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');
-    }
-
-    private decryptWithRsa(message: string): string {
-        if (!this.serverPublicKey) {
-            throw new Error('Server public key not set.');
-        }
-        try {
-            const encryptedMessage = Buffer.from(message.toString(), 'base64');
-            const decrypted = publicDecrypt(
-                {
-                    key: this.serverPublicKey,
-                    padding: constants.RSA_PKCS1_PADDING,
-                },
-                encryptedMessage
-            );
-            return decrypted.toString('utf-8');
-        } catch (error) {
-            throw new Error('Failed to decrypt RSA message.');
-        }
-    }
-
-    async sendChunkedMessage(operationCode: string, metaInfo?: { [key: string]: any }, fileContent?: Buffer): Promise<void> {
-        const message = MessageHandler.formatMessage(operationCode, metaInfo, fileContent);
-        const outgoingMessage = this.encryptWithAes(message);
-
-        const totalChunks = Math.ceil(outgoingMessage.length / CHUNK_SIZE);
-        const messageId = Date.now().toString();
-
-        for (let i = 0; i < totalChunks; i++) {
-            const chunk = outgoingMessage.slice(i * CHUNK_SIZE, (i + 1) * CHUNK_SIZE);
-            const chunkHeader = JSON.stringify({
-                messageId,
-                sequenceNumber: i,
-                totalChunks,
-            });
-
-            const chunkWithHeader = `${chunkHeader}|${chunk}`;
-
-            await this.writeToSocket(chunkWithHeader);
-
-            if (i === totalChunks - 1) {
-                await this.writeToSocket(END_OF_MESSAGE);
-            }
-        }
-    }
-
-    private writeToSocket(message: string): Promise<void> {
-        return new Promise((resolve, reject) => {
-            this.socket.write(message, (err: any) => {
-                if (err) {
-                    return reject(err);
-                }
-                resolve();
-            });
-        });
-    }
-
-    async handleIncomingChunk(data: Buffer): Promise<void> {
-        const incomingMessage = data.toString();
-        this.messageBuffer += incomingMessage;
-
-        if (this.messageBuffer.includes(END_OF_MESSAGE)) {
-            const messages = this.messageBuffer.split(END_OF_MESSAGE);
-
-            for (let i = 0; i < messages.length - 1; i++) {
-                const completeMessage = messages[i];
-                if (completeMessage) {
-                    await this.processCompleteMessage(completeMessage);
-                }
-            }
-
-            this.messageBuffer = messages[messages.length - 1];
-        }
-    }
-
-    private async processCompleteMessage(completeMessage: string): Promise<void> {
-        const [headerJson, chunkContent] = completeMessage.split('|');
-        const header = JSON.parse(headerJson);
-
-        if (!this.chunkBuffers[header.messageId]) {
-            this.chunkBuffers[header.messageId] = new Array(header.totalChunks);
-        }
-
-        this.chunkBuffers[header.messageId][header.sequenceNumber] = chunkContent;
-
-        if (this.chunkBuffers[header.messageId].every((chunk) => chunk !== undefined)) {
-            const fullMessage = this.chunkBuffers[header.messageId].join('');
-            await this.handleIncomingMessage(fullMessage);
-            delete this.chunkBuffers[header.messageId];
-        }
-    }
-
     async handleIncomingMessage(incomingMessage: string): Promise<void> {
-        let messageToProcess = incomingMessage;
+        let messageToProcess;
+
         if (this.aesKey && this.aesIv) {
             messageToProcess = this.decryptWithAes(incomingMessage);
-        } else if (this.serverPublicKey) {
+        } else if (this.publicKey) {
             messageToProcess = this.decryptWithRsa(incomingMessage);
+        } else {
+            messageToProcess = Buffer.from(incomingMessage, 'base64').toString('utf-8');
         }
 
         const result = await this.operationHandler.handleOperation(messageToProcess);
@@ -171,13 +59,31 @@ export class TcpClientCommunicator extends SocketCommunicatorBase {
         this.handlerResult = result;
     }
 
-    isAesKeySet(): boolean {
-        return this.isAesKeySetFlag;
+    async sendMessage(operationCode: string, metaInfo?: { [key: string]: any }, fileContent?: Buffer): Promise<void> {
+        if (!this.networkSpeed) {
+            this.networkSpeed = await this.scanNetworkLatency();
         }
 
-    getHandlerResult(): ParsedMessage | null {
-        const message = this.handlerResult;
-        this.handlerResult = null;
-        return message;
+        const message = MessageHandler.formatMessage(operationCode, metaInfo, fileContent);
+
+        const outgoingMessage = this.encryptWithAes(message);
+
+        // Calculate optimal chunk size based on network latency
+        const optimalChunkSize = await this.calculateOptimalChunkSize(outgoingMessage.length);
+        const totalChunks = Math.ceil(outgoingMessage.length / optimalChunkSize);
+        const messageId = Date.now().toString();
+
+        // Send each chunk with a delay between them
+        for (let i = 0; i < totalChunks; i++) {
+            const chunk = outgoingMessage.slice(i * optimalChunkSize, (i + 1) * optimalChunkSize);
+            const chunkHeader = JSON.stringify({
+                messageId,
+                sequenceNumber: i + 1,
+                totalChunks,
+            });
+            const chunkWithHeader = `${chunkHeader}|${chunk}${this.EOP}`;
+            if(!this.socket.write(chunkWithHeader)) this.socket.end();
         }
     }
+
+}

CEO/src/network/socket_communicator/tcp_server_communicator.ts

@@ -1,194 +1,76 @@
 import { Socket } from 'net';
-import { privateEncrypt, generateKeyPairSync, createCipheriv, createDecipheriv, randomBytes, constants } from 'crypto';
 import { MessageHandler } from '../message_handler';
 import { SocketCommunicatorBase } from './socket_communicator_base';
 import { OperationHandler } from '../operations_base/operation_handler';
-
-const END_OF_MESSAGE = '<EOM>';  // Define a unique marker for end of message
-const CHUNK_SIZE = 1024;  // Define chunk size
+import {operationCodes} from "../operation_codes";
 
 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 messageBuffer: string;
-    private chunkBuffers: { [messageId: string]: string[] };  // Buffer for reassembling chunks
 
     constructor(socket: Socket, ip: string, port: number, operationHandler: OperationHandler) {
         super(ip, port, operationHandler);
         this.socket = socket;
-        this.privateKey = null;
-        this.publicKey = null;  // Client public key will be set later
-        this.aesKey = null;
-        this.aesIv = null;
-        this.messageBuffer = '';  // Initialize the message buffer
-        this.chunkBuffers = {};  // Buffer for reassembling incoming messages
-        this.generateKeyPair();  // Generate RSA key pair for encryption
     }
 
-    // Generate RSA key pair (public and private keys)
-    generateKeyPair(): void {
-        const { privateKey, publicKey } = generateKeyPairSync('rsa', {
-            modulusLength: 2048,
-            publicKeyEncoding: { type: 'spki', format: 'pem' },
-            privateKeyEncoding: { type: 'pkcs8', format: 'pem' },
-        });
-        this.privateKey = privateKey;
-        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.');
+        this.generateKeyPair();
+        if (!this.publicKey || !this.privateKey) {
+            throw new Error('RSA key pair is not available. Please generate RSA key pair.');
         }
 
-        await this.sendChunkedMessage('SET_PUBLIC_KEY', { publicKey: this.publicKey });
+        await this.sendMessage(operationCodes.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.generateAesKey();
+        if (!this.aesKey || !this.aesIv) {
+            throw new Error('AES key or IV is not available. Please generate AES key.');
+        }
 
         const aesKeyBase64 = this.aesKey.toString('base64');
         const aesIvBase64 = this.aesIv.toString('base64');
-
-        await this.sendChunkedMessage('SET_AES_KEY', { aesKey: aesKeyBase64, aesIv: aesIvBase64 });
+        await this.sendMessage(operationCodes.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
-            },
-            Buffer.from(message)
-        ).toString('base64');
+    async handleIncomingMessage(incomingMessage: string): Promise<void> {
+        const messageToProcess = this.decryptWithAes(incomingMessage);
+        this.handlerResult = await this.operationHandler.handleOperation(messageToProcess);
     }
 
-    // Decrypt AES-encrypted messages
-    private decryptWithAes(encryptedMessage: string): string {
-        if (!this.aesKey || !this.aesIv) {
-            throw new Error('AES key not set.');
+    async sendMessage(operationCode: string, metaInfo?: { [key: string]: any }, fileContent?: Buffer): Promise<void> {
+        if (!this.networkSpeed) {
+            this.networkSpeed = await this.scanNetworkLatency();
         }
 
-        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.');
-        }
-        const cipher = createCipheriv('aes-256-cbc', this.aesKey, this.aesIv);
-        let encrypted = cipher.update(message, 'utf-8');
-        encrypted = Buffer.concat([encrypted, cipher.final()]);
-        return encrypted.toString('base64');
-    }
-
-    // Handle incoming chunks of data
-    async handleIncomingChunk(data: Buffer): Promise<void> {
-        const incomingMessage = data.toString();
-        this.messageBuffer += incomingMessage;
-
-        if (this.messageBuffer.includes(END_OF_MESSAGE)) {
-            const messages = this.messageBuffer.split(END_OF_MESSAGE);
-
-            for (let i = 0; i < messages.length - 1; i++) {
-                const completeMessage = messages[i];
-                if (completeMessage) {
-                    await this.processCompleteMessage(completeMessage);
-                }
-            }
-
-            this.messageBuffer = messages[messages.length - 1];
-        }
-    }
-
-    private async processCompleteMessage(completeMessage: string): Promise<void> {
-        const [headerJson, chunkContent] = completeMessage.split('|');
-        const header = JSON.parse(headerJson);
-
-        if (!this.chunkBuffers[header.messageId]) {
-            this.chunkBuffers[header.messageId] = new Array(header.totalChunks);
-        }
-
-        this.chunkBuffers[header.messageId][header.sequenceNumber] = chunkContent;
-
-        if (this.chunkBuffers[header.messageId].every((chunk) => chunk !== undefined)) {
-            const fullMessage = this.chunkBuffers[header.messageId].join('');
-            await this.handleIncomingMessage(fullMessage);
-            delete this.chunkBuffers[header.messageId];
-        }
-    }
-
-    // Send chunked message
-    async sendChunkedMessage(operationCode: string, metaInfo?: { [key: string]: any }, fileContent?: Buffer): Promise<void> {
         const message = MessageHandler.formatMessage(operationCode, metaInfo, fileContent);
-        let outgoingMessage: string;
 
+        let outgoingMessage: string;
         switch (operationCode) {
-            case 'SET_PUBLIC_KEY':
-                outgoingMessage = message;
+            case operationCodes.SET_PUBLIC_KEY:
+                outgoingMessage = Buffer.from(message, 'utf-8').toString('base64');
                 break;
-            case 'SET_AES_KEY':
+            case operationCodes.SET_AES_KEY:
                 outgoingMessage = this.encryptWithRsa(message);
                 break;
             default:
                 outgoingMessage = this.encryptWithAes(message);
         }
 
-        const totalChunks = Math.ceil(outgoingMessage.length / CHUNK_SIZE);
+        // Calculate optimal chunk size based on network latency
+        const optimalChunkSize = await this.calculateOptimalChunkSize(outgoingMessage.length);
+        const totalChunks = Math.ceil(outgoingMessage.length / optimalChunkSize);
         const messageId = Date.now().toString();
 
+        // Send each chunk with a delay between them
         for (let i = 0; i < totalChunks; i++) {
-            const chunk = outgoingMessage.slice(i * CHUNK_SIZE, (i + 1) * CHUNK_SIZE);
+            const chunk = outgoingMessage.slice(i * optimalChunkSize, (i + 1) * optimalChunkSize);
             const chunkHeader = JSON.stringify({
                 messageId,
-                sequenceNumber: i,
+                sequenceNumber: i + 1,
                 totalChunks,
             });
-
-            const chunkWithHeader = `${chunkHeader}|${chunk}`;
-
-            await this.writeToSocket(chunkWithHeader);
-
-            if (i === totalChunks - 1) {
-                await this.writeToSocket(END_OF_MESSAGE);
+            const chunkWithHeader = `${chunkHeader}|${chunk}${this.EOP}`;
+            if(!this.socket.write(chunkWithHeader)) this.socket.end();
         }
     }
 }
-
-    // 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) => {
-                if (err) {
-                    console.error('Error sending message over TCP:', err);
-                    return reject(err);
-                }
-                resolve();
-            });
-        });
-    }
-}

CEO/src/network/socket_communicator/udp_socket_communicator.ts

@@ -11,20 +11,9 @@ export class UdpSocketCommunicator extends SocketCommunicatorBase {
         this.socket = socket;
     }
 
-    // Handle incoming message (no decryption needed for UDP)
-    async handleIncomingMessage(incomingMessage: string): Promise<void> {
-        this.handlerResult = await this.operationHandler.handleOperation(incomingMessage);
-    }
-
     // Send plain message over UDP (metaInfo as JSON and fileContent as Buffer)
     async sendMessage(operationCode: string, metaInfo?: { [key: string]: any }, fileContent?: Buffer): Promise<void> {
         const message = MessageHandler.formatMessage(operationCode, metaInfo, fileContent);
-
-        await this.sendUdpMessage(message);
-    }
-
-    // Helper method to wrap socket.send in a Promise for async/await support
-    private sendUdpMessage(message: string): Promise<void> {
         return new Promise((resolve, reject) => {
             this.socket.send(message, this.port, this.ip, (err: any) => {
                 if (err) {
@@ -36,4 +25,9 @@ export class UdpSocketCommunicator extends SocketCommunicatorBase {
             });
         });
     }
+
+    // Handle incoming message (no decryption needed for UDP)
+    async handleIncomingMessage(incomingMessage: string): Promise<void> {
+        this.handlerResult = await this.operationHandler.handleOperation(incomingMessage);
+    }
 }

CEO/src/network/tcp/tcp_client.ts

@@ -80,7 +80,7 @@ export class TcpClient {
         }
 
         this.log(`Sending message with operationCode: ${operationCode}`);
-        await this.communicator.sendChunkedMessage(operationCode, metaInfo, fileContent);
+        await this.communicator.sendMessage(operationCode, metaInfo, fileContent);
         return true;
     }
 

CEO/src/network/tcp/tcp_server.ts

@@ -43,7 +43,6 @@ export class TcpServer {
 
             const tcpCommunicator = new TcpServerCommunicator(socket, ip, port, this.operationHandler);
             this.connectionManager.addConnection(ip, port, tcpCommunicator);
-            tcpCommunicator.generateKeyPair();
 
             tcpCommunicator.sendPublicKey()
                 .then(() => tcpCommunicator.sendAesKey())
@@ -106,7 +105,7 @@ export class TcpServer {
         const handlerResult = communicator.getHandlerResult();
         if (handlerResult) {
             try {
-                await communicator.sendChunkedMessage(
+                await communicator.sendMessage(
                     handlerResult.operationCode,
                     handlerResult.metaInfo,
                     handlerResult.fileContent

UC/src/network/socket_communicator/socket_communicator_base.ts

@@ -26,6 +26,8 @@ export abstract class SocketCommunicatorBase {
     protected aesKey: Buffer | null;
     protected aesIv: Buffer | null;
 
+    private incompleteChunkBuffer: string = '';
+
     protected constructor(ip: string, port: number, operationHandler: OperationHandler) {
         this.ip = ip;
         this.port = port;
@@ -146,35 +148,46 @@ export abstract class SocketCommunicatorBase {
     }
 
     async handleIncomingChunk(data: Buffer): Promise<void> {
-        const incomingData = data.toString().trim();
+        // Append incoming data to the incomplete buffer
+        this.incompleteChunkBuffer += data.toString();
 
-        // Split the incoming data by <EOP> to handle multiple chunks concatenated by TCP
-        const messages = incomingData.split(this.EOP).filter(Boolean); // Filter out any empty strings from split
+        // Split the buffer by <EOP> to separate complete and incomplete messages
+        const messages = this.incompleteChunkBuffer.split(this.EOP);
 
+        // Save the last item back to the buffer if it's incomplete (no <EOP> at the end)
+        this.incompleteChunkBuffer = messages.pop() || "";
+
+        // Process each complete message in the split results
         for (const incomingMessage of messages) {
+            try {
                 const [headerJson, chunkContent] = incomingMessage.split('|');
                 const header = JSON.parse(headerJson);
 
-            // Initialize an array for chunks if it's the first chunk received for this messageId
+                // Initialize an array for chunks if it's the first chunk for this messageId
                 if (!this.chunkBuffers[header.messageId]) {
                     this.chunkBuffers[header.messageId] = [];
                 }
 
-            // Directly set the chunk at the correct index, adjusting for 1-based indexing
+                // Store the chunk in the correct position based on sequenceNumber (1-based indexing)
                 this.chunkBuffers[header.messageId][header.sequenceNumber - 1] = chunkContent;
 
                 console.log(`Received chunk: ${incomingMessage}`);
                 console.log(`Chunks received so far: ${this.chunkBuffers[header.messageId].filter(chunk => chunk !== undefined).length}`);
 
-            // Check if all chunks have been received by confirming the length
+                // Check if all chunks have been received
                 if (this.chunkBuffers[header.messageId].filter(chunk => chunk !== undefined).length === header.totalChunks) {
-                const chunks = this.chunkBuffers[header.messageId];
-                const fullMessage = chunks.join('');
+                    // Join all chunks to form the full message
+                    const fullMessage = this.chunkBuffers[header.messageId].join('');
 
+                    // Process the complete message
                     await this.handleIncomingMessage(fullMessage);
 
+                    // Clear the chunk buffer for this messageId
                     delete this.chunkBuffers[header.messageId];
                 }
+            } catch (error: any) {
+                console.error(`Error handling chunk: ${error.message}`);
+            }
         }
     }
 

User/src/network/socket_communicator/socket_communicator_base.ts

@@ -26,6 +26,8 @@ export abstract class SocketCommunicatorBase {
     protected aesKey: Buffer | null;
     protected aesIv: Buffer | null;
 
+    private incompleteChunkBuffer: string = '';
+
     protected constructor(ip: string, port: number, operationHandler: OperationHandler) {
         this.ip = ip;
         this.port = port;
@@ -146,35 +148,46 @@ export abstract class SocketCommunicatorBase {
     }
 
     async handleIncomingChunk(data: Buffer): Promise<void> {
-        const incomingData = data.toString().trim();
+        // Append incoming data to the incomplete buffer
+        this.incompleteChunkBuffer += data.toString();
 
-        // Split the incoming data by <EOP> to handle multiple chunks concatenated by TCP
-        const messages = incomingData.split(this.EOP).filter(Boolean); // Filter out any empty strings from split
+        // Split the buffer by <EOP> to separate complete and incomplete messages
+        const messages = this.incompleteChunkBuffer.split(this.EOP);
 
+        // Save the last item back to the buffer if it's incomplete (no <EOP> at the end)
+        this.incompleteChunkBuffer = messages.pop() || "";
+
+        // Process each complete message in the split results
         for (const incomingMessage of messages) {
+            try {
                 const [headerJson, chunkContent] = incomingMessage.split('|');
                 const header = JSON.parse(headerJson);
 
-            // Initialize an array for chunks if it's the first chunk received for this messageId
+                // Initialize an array for chunks if it's the first chunk for this messageId
                 if (!this.chunkBuffers[header.messageId]) {
                     this.chunkBuffers[header.messageId] = [];
                 }
 
-            // Directly set the chunk at the correct index, adjusting for 1-based indexing
+                // Store the chunk in the correct position based on sequenceNumber (1-based indexing)
                 this.chunkBuffers[header.messageId][header.sequenceNumber - 1] = chunkContent;
 
                 console.log(`Received chunk: ${incomingMessage}`);
                 console.log(`Chunks received so far: ${this.chunkBuffers[header.messageId].filter(chunk => chunk !== undefined).length}`);
 
-            // Check if all chunks have been received by confirming the length
+                // Check if all chunks have been received
                 if (this.chunkBuffers[header.messageId].filter(chunk => chunk !== undefined).length === header.totalChunks) {
-                const chunks = this.chunkBuffers[header.messageId];
-                const fullMessage = chunks.join('');
+                    // Join all chunks to form the full message
+                    const fullMessage = this.chunkBuffers[header.messageId].join('');
 
+                    // Process the complete message
                     await this.handleIncomingMessage(fullMessage);
 
+                    // Clear the chunk buffer for this messageId
                     delete this.chunkBuffers[header.messageId];
                 }
+            } catch (error: any) {
+                console.error(`Error handling chunk: ${error.message}`);
+            }
         }
     }