network chunk v7

This commit is contained in:
andrei-mihnea-cerbu
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 { SocketCommunicatorBase } from './socket_communicator_base';
import { OperationHandler } from '../operations_base/operation_handler'; import { OperationHandler } from '../operations_base/operation_handler';
const MAX_CHUNK_SIZE = 2048; // Define chunk size
export class TcpServerCommunicator extends SocketCommunicatorBase { export class TcpServerCommunicator extends SocketCommunicatorBase {
private readonly socket: Socket; private readonly socket: Socket;
private privateKey: string | null; private privateKey: string | null;
private publicKey: string | null; private publicKey: string | null;
private aesKey: Buffer | null; private aesKey: Buffer | null;
private aesIv: 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) { constructor(socket: Socket, ip: string, port: number, operationHandler: OperationHandler) {
super(ip, port, 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.publicKey = null; // Client public key will be set later
this.aesKey = null; this.aesKey = null;
this.aesIv = null; this.aesIv = null;
this.chunkBuffers = {}; // Buffer for reassembling incoming messages this.chunkBuffers = {};
this.generateKeyPair(); // Generate RSA key pair for encryption this.generateKeyPair();
} }
// Generate RSA key pair (public and private keys)
generateKeyPair(): void { generateKeyPair(): void {
const { privateKey, publicKey } = generateKeyPairSync('rsa', { const { privateKey, publicKey } = generateKeyPairSync('rsa', {
modulusLength: 2048, modulusLength: 2048,
@@ -36,52 +34,42 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
this.publicKey = publicKey; this.publicKey = publicKey;
} }
// Send the server's public key to the client
async sendPublicKey(): Promise<void> { async sendPublicKey(): Promise<void> {
if (!this.publicKey) { if (!this.publicKey) {
throw new Error('Public key is not available. Please generate RSA key pair.'); throw new Error('Public key is not available. Please generate RSA key pair.');
} }
await this.sendChunkedMessage('SET_PUBLIC_KEY', { publicKey: this.publicKey }); await this.sendChunkedMessage('SET_PUBLIC_KEY', { publicKey: this.publicKey });
} }
// Generate AES key and IV, then send them to the client
async sendAesKey(): Promise<void> { async sendAesKey(): Promise<void> {
this.aesKey = randomBytes(32); // 256-bit AES key this.aesKey = randomBytes(32);
this.aesIv = randomBytes(16); // AES IV this.aesIv = randomBytes(16);
const aesKeyBase64 = this.aesKey.toString('base64'); const aesKeyBase64 = this.aesKey.toString('base64');
const aesIvBase64 = this.aesIv.toString('base64'); const aesIvBase64 = this.aesIv.toString('base64');
await this.sendChunkedMessage('SET_AES_KEY', { aesKey: aesKeyBase64, aesIv: aesIvBase64 }); 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 { private encryptWithRsa(message: string): string {
if (!this.privateKey) throw new Error('Server private key not set.'); if (!this.privateKey) throw new Error('Server private key not set.');
return privateEncrypt( return privateEncrypt(
{ {
key: this.privateKey, key: this.privateKey,
padding: constants.RSA_PKCS1_PADDING, // PKCS1 padding padding: constants.RSA_PKCS1_PADDING,
}, },
Buffer.from(message) Buffer.from(message)
).toString('base64'); ).toString('base64');
} }
// Decrypt AES-encrypted messages
private decryptWithAes(encryptedMessage: string): string { private decryptWithAes(encryptedMessage: string): string {
if (!this.aesKey || !this.aesIv) { if (!this.aesKey || !this.aesIv) {
throw new Error('AES key not set.'); throw new Error('AES key not set.');
} }
const decipher = createDecipheriv('aes-256-cbc', this.aesKey, this.aesIv); const decipher = createDecipheriv('aes-256-cbc', this.aesKey, this.aesIv);
let decrypted = decipher.update(Buffer.from(encryptedMessage, 'base64')); let decrypted = decipher.update(Buffer.from(encryptedMessage, 'base64'));
decrypted = Buffer.concat([decrypted, decipher.final()]); decrypted = Buffer.concat([decrypted, decipher.final()]);
return decrypted.toString('utf-8'); return decrypted.toString('utf-8');
} }
// Encrypt a message with AES
private encryptWithAes(message: string): string { private encryptWithAes(message: string): string {
if (!this.aesKey || !this.aesIv) { if (!this.aesKey || !this.aesIv) {
throw new Error('AES key or IV is not set.'); throw new Error('AES key or IV is not set.');
@@ -92,55 +80,58 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
return encrypted.toString('base64'); return encrypted.toString('base64');
} }
// Handle incoming chunks of data
async handleIncomingChunk(data: Buffer): Promise<void> { async handleIncomingChunk(data: Buffer): Promise<void> {
const incomingMessage = data.toString().trim(); const incomingMessage = data.toString().trim();
// Extract header and chunk content from incoming data
const [headerJson, chunkContent] = incomingMessage.split('|'); const [headerJson, chunkContent] = incomingMessage.split('|');
const header = JSON.parse(headerJson); 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]) { if (!this.chunkBuffers[header.messageId]) {
this.chunkBuffers[header.messageId] = new Array(header.totalChunks); this.chunkBuffers[header.messageId] = new Array(header.totalChunks).fill(undefined);
} }
this.chunkBuffers[header.messageId][header.sequenceNumber - 1] = chunkContent;
// Place the chunk in the correct position in the chunk buffer if (this.chunkBuffers[header.messageId].every((chunk) => chunk !== undefined)) {
this.chunkBuffers[header.messageId][header.sequenceNumber] = chunkContent; await this.processCompleteMessage(header.messageId);
if(header.sequenceNumber === header.totalChunks) {
await this.processCompleteMessage();
} }
} }
// Process the complete message when EOM is received private async processCompleteMessage(messageId: string): Promise<void> {
private async processCompleteMessage(): Promise<void> { const chunks = this.chunkBuffers[messageId];
for (const [messageId, chunks] of Object.entries(this.chunkBuffers)) { if (chunks && chunks.every((chunk) => chunk !== undefined)) {
if (chunks.every((chunk) => chunk !== undefined)) {
// Join all chunks to form the full message
const fullMessage = chunks.join(''); const fullMessage = chunks.join('');
// Handle the completed and possibly decrypted message
await this.handleIncomingMessage(fullMessage); await this.handleIncomingMessage(fullMessage);
// Clean up buffer after processing
delete this.chunkBuffers[messageId]; delete this.chunkBuffers[messageId];
} }
} }
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;
} }
// Send chunked message
async sendChunkedMessage( async sendChunkedMessage(
operationCode: string, operationCode: string,
metaInfo?: { [key: string]: any }, metaInfo?: { [key: string]: any },
fileContent?: Buffer fileContent?: Buffer
): Promise<void> { ): Promise<void> {
if (!this.networkSpeed) {
await this.scanNetworkSpeed();
}
const message = MessageHandler.formatMessage(operationCode, metaInfo, fileContent); const message = MessageHandler.formatMessage(operationCode, metaInfo, fileContent);
let outgoingMessage: string; let outgoingMessage: string;
switch (operationCode) {
// Encrypt or format message based on the operation code
switch(operationCode) {
case 'SET_PUBLIC_KEY': case 'SET_PUBLIC_KEY':
outgoingMessage = Buffer.from(message, 'utf-8').toString('base64'); outgoingMessage = Buffer.from(message, 'utf-8').toString('base64');
break; break;
@@ -150,20 +141,9 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
default: default:
outgoingMessage = this.encryptWithAes(message); outgoingMessage = this.encryptWithAes(message);
} }
const optimalChunkSize = this.calculateOptimalChunkSize(outgoingMessage.length);
// 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 totalChunks = Math.ceil(outgoingMessage.length / optimalChunkSize); const totalChunks = Math.ceil(outgoingMessage.length / optimalChunkSize);
const messageId = Date.now().toString(); const messageId = Date.now().toString();
// Send each chunk as a string with delay to manage network flow
for (let i = 0; i < totalChunks; i++) { for (let i = 0; i < totalChunks; i++) {
const chunk = outgoingMessage.slice(i * optimalChunkSize, (i + 1) * optimalChunkSize); const chunk = outgoingMessage.slice(i * optimalChunkSize, (i + 1) * optimalChunkSize);
const chunkHeader = JSON.stringify({ const chunkHeader = JSON.stringify({
@@ -171,27 +151,20 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
sequenceNumber: i + 1, sequenceNumber: i + 1,
totalChunks, totalChunks,
}); });
const chunkWithHeader = `${chunkHeader}|${chunk}`; const chunkWithHeader = `${chunkHeader}|${chunk}`;
// Write the chunk string directly to the socket
await this.writeToSocket(chunkWithHeader); 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> { async handleIncomingMessage(incomingMessage: string): Promise<void> {
let messageToProcess = incomingMessage; let messageToProcess = incomingMessage;
if (this.aesKey && this.aesIv) { if (this.aesKey && this.aesIv) {
messageToProcess = this.decryptWithAes(incomingMessage); messageToProcess = this.decryptWithAes(incomingMessage);
} }
this.handlerResult = await this.operationHandler.handleOperation(messageToProcess); this.handlerResult = await this.operationHandler.handleOperation(messageToProcess);
} }
// Write message to socket
private writeToSocket(message: string): Promise<void> { private writeToSocket(message: string): Promise<void> {
return new Promise((resolve, reject) => { return new Promise((resolve, reject) => {
this.socket.write(message, (err: any) => { this.socket.write(message, (err: any) => {
@@ -1,28 +1,27 @@
import { Socket } from 'net'; import { Socket } from 'net';
import { createCipheriv, createDecipheriv, constants, publicDecrypt } from 'crypto'; import { createCipheriv, createDecipheriv, constants, publicDecrypt } from 'crypto';
import {MessageHandler, ParsedMessage} from '../message_handler'; import { MessageHandler, ParsedMessage } from '../message_handler';
import { SocketCommunicatorBase } from './socket_communicator_base'; import { SocketCommunicatorBase } from './socket_communicator_base';
import { OperationHandler } from '../operations_base/operation_handler'; import { OperationHandler } from '../operations_base/operation_handler';
import { operationCodes } from '../operation_codes'; import { operationCodes } from '../operation_codes';
const MAX_CHUNK_SIZE = 2048;
export class TcpClientCommunicator extends SocketCommunicatorBase { export class TcpClientCommunicator extends SocketCommunicatorBase {
private readonly socket: Socket; private readonly socket: Socket;
private aesKey: Buffer | null; private aesKey: Buffer | null;
private aesIv: Buffer | null; private aesIv: Buffer | null;
private serverPublicKey: string | null; private serverPublicKey: string | null;
private isAesKeySetFlag: boolean; private isAesKeySetFlag: boolean;
private chunkBuffers: { [messageId: string]: string[] }; // Buffer for reassembling chunks private chunkBuffers: { [messageId: string]: string[] };
private networkSpeed: number | null = null;
constructor(socket: Socket, ip: string, port: number, operationHandler: OperationHandler) { constructor(socket: Socket, ip: string, port: number, operationHandler: OperationHandler) {
super(ip, port, operationHandler); // Call parent constructor super(ip, port, operationHandler);
this.socket = socket; this.socket = socket;
this.aesKey = null; this.aesKey = null;
this.aesIv = null; this.aesIv = null;
this.serverPublicKey = null; this.serverPublicKey = null;
this.isAesKeySetFlag = false; this.isAesKeySetFlag = false;
this.chunkBuffers = {}; // Initialize chunk buffer for reassembling messages this.chunkBuffers = {};
} }
setServerPublicKey(publicKey: string): void { setServerPublicKey(publicKey: string): void {
@@ -59,7 +58,7 @@ export class TcpClientCommunicator extends SocketCommunicatorBase {
throw new Error('Server public key not set.'); throw new Error('Server public key not set.');
} }
try { try {
const encryptedMessage = Buffer.from(message.toString(), 'base64'); const encryptedMessage = Buffer.from(message, 'base64');
const decrypted = publicDecrypt( const decrypted = publicDecrypt(
{ {
key: this.serverPublicKey, key: this.serverPublicKey,
@@ -73,23 +72,36 @@ export class TcpClientCommunicator extends SocketCommunicatorBase {
} }
} }
async sendChunkedMessage(operationCode: string, metaInfo?: { [key: string]: any }, fileContent?: Buffer): Promise<void> { private async scanNetworkSpeed(): Promise<number> {
const message = MessageHandler.formatMessage(operationCode, metaInfo, fileContent); const testMessage = 'PING_TEST'.repeat(100);
const outgoingMessage = this.encryptWithAes(message); const startTime = Date.now();
await this.writeToSocket(testMessage);
// Determine optimal chunk size (multiple of 4 and ≤ 1024 to align with base64 encoding) await new Promise((resolve) => this.socket.once('data', resolve));
let optimalChunkSize = MAX_CHUNK_SIZE; const endTime = Date.now();
while (outgoingMessage.length % optimalChunkSize !== 0 && optimalChunkSize > 0) { const duration = endTime - startTime;
optimalChunkSize -= 4; this.networkSpeed = testMessage.length / duration;
return this.networkSpeed;
} }
if (optimalChunkSize === 0) optimalChunkSize = MAX_CHUNK_SIZE; // Fallback in case of odd alignment 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;
}
// Calculate total chunks and generate unique message ID 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);
const outgoingMessage = this.encryptWithAes(message);
const optimalChunkSize = this.calculateOptimalChunkSize(outgoingMessage.length);
const totalChunks = Math.ceil(outgoingMessage.length / optimalChunkSize); const totalChunks = Math.ceil(outgoingMessage.length / optimalChunkSize);
const messageId = Date.now().toString(); const messageId = Date.now().toString();
// Send each chunk as a string with delay to manage network flow
for (let i = 0; i < totalChunks; i++) { for (let i = 0; i < totalChunks; i++) {
const chunk = outgoingMessage.slice(i * optimalChunkSize, (i + 1) * optimalChunkSize); const chunk = outgoingMessage.slice(i * optimalChunkSize, (i + 1) * optimalChunkSize);
const chunkHeader = JSON.stringify({ const chunkHeader = JSON.stringify({
@@ -100,9 +112,8 @@ export class TcpClientCommunicator extends SocketCommunicatorBase {
const chunkWithHeader = `${chunkHeader}|${chunk}`; const chunkWithHeader = `${chunkHeader}|${chunk}`;
// Write the chunk string directly to the socket
await this.writeToSocket(chunkWithHeader); await this.writeToSocket(chunkWithHeader);
await new Promise((resolve) => setTimeout(resolve, 300)); // Simulate network delay await new Promise((resolve) => setTimeout(resolve, 300));
} }
} }
@@ -119,40 +130,27 @@ export class TcpClientCommunicator extends SocketCommunicatorBase {
async handleIncomingChunk(data: Buffer): Promise<void> { async handleIncomingChunk(data: Buffer): Promise<void> {
const incomingMessage = data.toString().trim(); const incomingMessage = data.toString().trim();
console.log(`\n\n${incomingMessage}\n\n`);
// Extract header and chunk content from incoming data
const [headerJson, chunkContent] = incomingMessage.split('|'); const [headerJson, chunkContent] = incomingMessage.split('|');
const header = JSON.parse(headerJson); const header = JSON.parse(headerJson);
// Initialize chunk array if this is the first chunk for this messageId
if (!this.chunkBuffers[header.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 - 1] = chunkContent;
this.chunkBuffers[header.messageId][header.sequenceNumber] = chunkContent; if (this.chunkBuffers[header.messageId].every((chunk) => chunk !== undefined)) {
if(header.sequenceNumber === header.totalChunks) { await this.processCompleteMessage(header.messageId);
await this.processCompleteMessage();
} }
} }
// Process the complete message when EOM is received private async processCompleteMessage(messageId: string): Promise<void> {
private async processCompleteMessage(): Promise<void> { const chunks = this.chunkBuffers[messageId];
for (const [messageId, chunks] of Object.entries(this.chunkBuffers)) { if (chunks && chunks.every((chunk) => chunk !== undefined)) {
if (chunks.every((chunk) => chunk !== undefined)) {
// Join all chunks to form the full message
const fullMessage = chunks.join(''); const fullMessage = chunks.join('');
// Handle the completed and possibly decrypted message
await this.handleIncomingMessage(fullMessage); await this.handleIncomingMessage(fullMessage);
// Clean up buffer after processing
delete this.chunkBuffers[messageId]; delete this.chunkBuffers[messageId];
} }
} }
}
async handleIncomingMessage(incomingMessage: string): Promise<void> { async handleIncomingMessage(incomingMessage: string): Promise<void> {
let messageToProcess; let messageToProcess;
@@ -160,7 +158,7 @@ export class TcpClientCommunicator extends SocketCommunicatorBase {
messageToProcess = this.decryptWithAes(incomingMessage); messageToProcess = this.decryptWithAes(incomingMessage);
} else if (this.serverPublicKey) { } else if (this.serverPublicKey) {
messageToProcess = this.decryptWithRsa(incomingMessage); messageToProcess = this.decryptWithRsa(incomingMessage);
}else{ } else {
messageToProcess = Buffer.from(incomingMessage, 'base64').toString('utf-8'); messageToProcess = Buffer.from(incomingMessage, 'base64').toString('utf-8');
} }
@@ -4,15 +4,14 @@ import { MessageHandler } from '../message_handler';
import { SocketCommunicatorBase } from './socket_communicator_base'; import { SocketCommunicatorBase } from './socket_communicator_base';
import { OperationHandler } from '../operations_base/operation_handler'; import { OperationHandler } from '../operations_base/operation_handler';
const MAX_CHUNK_SIZE = 2048; // Define chunk size
export class TcpServerCommunicator extends SocketCommunicatorBase { export class TcpServerCommunicator extends SocketCommunicatorBase {
private readonly socket: Socket; private readonly socket: Socket;
private privateKey: string | null; private privateKey: string | null;
private publicKey: string | null; private publicKey: string | null;
private aesKey: Buffer | null; private aesKey: Buffer | null;
private aesIv: 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) { constructor(socket: Socket, ip: string, port: number, operationHandler: OperationHandler) {
super(ip, port, 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.publicKey = null; // Client public key will be set later
this.aesKey = null; this.aesKey = null;
this.aesIv = null; this.aesIv = null;
this.chunkBuffers = {}; // Buffer for reassembling incoming messages this.chunkBuffers = {};
this.generateKeyPair(); // Generate RSA key pair for encryption this.generateKeyPair();
} }
// Generate RSA key pair (public and private keys)
generateKeyPair(): void { generateKeyPair(): void {
const { privateKey, publicKey } = generateKeyPairSync('rsa', { const { privateKey, publicKey } = generateKeyPairSync('rsa', {
modulusLength: 2048, modulusLength: 2048,
@@ -36,52 +34,42 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
this.publicKey = publicKey; this.publicKey = publicKey;
} }
// Send the server's public key to the client
async sendPublicKey(): Promise<void> { async sendPublicKey(): Promise<void> {
if (!this.publicKey) { if (!this.publicKey) {
throw new Error('Public key is not available. Please generate RSA key pair.'); throw new Error('Public key is not available. Please generate RSA key pair.');
} }
await this.sendChunkedMessage('SET_PUBLIC_KEY', { publicKey: this.publicKey }); await this.sendChunkedMessage('SET_PUBLIC_KEY', { publicKey: this.publicKey });
} }
// Generate AES key and IV, then send them to the client
async sendAesKey(): Promise<void> { async sendAesKey(): Promise<void> {
this.aesKey = randomBytes(32); // 256-bit AES key this.aesKey = randomBytes(32);
this.aesIv = randomBytes(16); // AES IV this.aesIv = randomBytes(16);
const aesKeyBase64 = this.aesKey.toString('base64'); const aesKeyBase64 = this.aesKey.toString('base64');
const aesIvBase64 = this.aesIv.toString('base64'); const aesIvBase64 = this.aesIv.toString('base64');
await this.sendChunkedMessage('SET_AES_KEY', { aesKey: aesKeyBase64, aesIv: aesIvBase64 }); 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 { private encryptWithRsa(message: string): string {
if (!this.privateKey) throw new Error('Server private key not set.'); if (!this.privateKey) throw new Error('Server private key not set.');
return privateEncrypt( return privateEncrypt(
{ {
key: this.privateKey, key: this.privateKey,
padding: constants.RSA_PKCS1_PADDING, // PKCS1 padding padding: constants.RSA_PKCS1_PADDING,
}, },
Buffer.from(message) Buffer.from(message)
).toString('base64'); ).toString('base64');
} }
// Decrypt AES-encrypted messages
private decryptWithAes(encryptedMessage: string): string { private decryptWithAes(encryptedMessage: string): string {
if (!this.aesKey || !this.aesIv) { if (!this.aesKey || !this.aesIv) {
throw new Error('AES key not set.'); throw new Error('AES key not set.');
} }
const decipher = createDecipheriv('aes-256-cbc', this.aesKey, this.aesIv); const decipher = createDecipheriv('aes-256-cbc', this.aesKey, this.aesIv);
let decrypted = decipher.update(Buffer.from(encryptedMessage, 'base64')); let decrypted = decipher.update(Buffer.from(encryptedMessage, 'base64'));
decrypted = Buffer.concat([decrypted, decipher.final()]); decrypted = Buffer.concat([decrypted, decipher.final()]);
return decrypted.toString('utf-8'); return decrypted.toString('utf-8');
} }
// Encrypt a message with AES
private encryptWithAes(message: string): string { private encryptWithAes(message: string): string {
if (!this.aesKey || !this.aesIv) { if (!this.aesKey || !this.aesIv) {
throw new Error('AES key or IV is not set.'); throw new Error('AES key or IV is not set.');
@@ -92,55 +80,58 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
return encrypted.toString('base64'); return encrypted.toString('base64');
} }
// Handle incoming chunks of data
async handleIncomingChunk(data: Buffer): Promise<void> { async handleIncomingChunk(data: Buffer): Promise<void> {
const incomingMessage = data.toString().trim(); const incomingMessage = data.toString().trim();
// Extract header and chunk content from incoming data
const [headerJson, chunkContent] = incomingMessage.split('|'); const [headerJson, chunkContent] = incomingMessage.split('|');
const header = JSON.parse(headerJson); 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]) { if (!this.chunkBuffers[header.messageId]) {
this.chunkBuffers[header.messageId] = new Array(header.totalChunks); this.chunkBuffers[header.messageId] = new Array(header.totalChunks).fill(undefined);
} }
this.chunkBuffers[header.messageId][header.sequenceNumber - 1] = chunkContent;
// Place the chunk in the correct position in the chunk buffer if (this.chunkBuffers[header.messageId].every((chunk) => chunk !== undefined)) {
this.chunkBuffers[header.messageId][header.sequenceNumber] = chunkContent; await this.processCompleteMessage(header.messageId);
if(header.sequenceNumber === header.totalChunks) {
await this.processCompleteMessage();
} }
} }
// Process the complete message when EOM is received private async processCompleteMessage(messageId: string): Promise<void> {
private async processCompleteMessage(): Promise<void> { const chunks = this.chunkBuffers[messageId];
for (const [messageId, chunks] of Object.entries(this.chunkBuffers)) { if (chunks && chunks.every((chunk) => chunk !== undefined)) {
if (chunks.every((chunk) => chunk !== undefined)) {
// Join all chunks to form the full message
const fullMessage = chunks.join(''); const fullMessage = chunks.join('');
// Handle the completed and possibly decrypted message
await this.handleIncomingMessage(fullMessage); await this.handleIncomingMessage(fullMessage);
// Clean up buffer after processing
delete this.chunkBuffers[messageId]; delete this.chunkBuffers[messageId];
} }
} }
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;
} }
// Send chunked message
async sendChunkedMessage( async sendChunkedMessage(
operationCode: string, operationCode: string,
metaInfo?: { [key: string]: any }, metaInfo?: { [key: string]: any },
fileContent?: Buffer fileContent?: Buffer
): Promise<void> { ): Promise<void> {
if (!this.networkSpeed) {
await this.scanNetworkSpeed();
}
const message = MessageHandler.formatMessage(operationCode, metaInfo, fileContent); const message = MessageHandler.formatMessage(operationCode, metaInfo, fileContent);
let outgoingMessage: string; let outgoingMessage: string;
switch (operationCode) {
// Encrypt or format message based on the operation code
switch(operationCode) {
case 'SET_PUBLIC_KEY': case 'SET_PUBLIC_KEY':
outgoingMessage = Buffer.from(message, 'utf-8').toString('base64'); outgoingMessage = Buffer.from(message, 'utf-8').toString('base64');
break; break;
@@ -150,20 +141,9 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
default: default:
outgoingMessage = this.encryptWithAes(message); outgoingMessage = this.encryptWithAes(message);
} }
const optimalChunkSize = this.calculateOptimalChunkSize(outgoingMessage.length);
// 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 totalChunks = Math.ceil(outgoingMessage.length / optimalChunkSize); const totalChunks = Math.ceil(outgoingMessage.length / optimalChunkSize);
const messageId = Date.now().toString(); const messageId = Date.now().toString();
// Send each chunk as a string with delay to manage network flow
for (let i = 0; i < totalChunks; i++) { for (let i = 0; i < totalChunks; i++) {
const chunk = outgoingMessage.slice(i * optimalChunkSize, (i + 1) * optimalChunkSize); const chunk = outgoingMessage.slice(i * optimalChunkSize, (i + 1) * optimalChunkSize);
const chunkHeader = JSON.stringify({ const chunkHeader = JSON.stringify({
@@ -171,27 +151,20 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
sequenceNumber: i + 1, sequenceNumber: i + 1,
totalChunks, totalChunks,
}); });
const chunkWithHeader = `${chunkHeader}|${chunk}`; const chunkWithHeader = `${chunkHeader}|${chunk}`;
// Write the chunk string directly to the socket
await this.writeToSocket(chunkWithHeader); 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> { async handleIncomingMessage(incomingMessage: string): Promise<void> {
let messageToProcess = incomingMessage; let messageToProcess = incomingMessage;
if (this.aesKey && this.aesIv) { if (this.aesKey && this.aesIv) {
messageToProcess = this.decryptWithAes(incomingMessage); messageToProcess = this.decryptWithAes(incomingMessage);
} }
this.handlerResult = await this.operationHandler.handleOperation(messageToProcess); this.handlerResult = await this.operationHandler.handleOperation(messageToProcess);
} }
// Write message to socket
private writeToSocket(message: string): Promise<void> { private writeToSocket(message: string): Promise<void> {
return new Promise((resolve, reject) => { return new Promise((resolve, reject) => {
this.socket.write(message, (err: any) => { this.socket.write(message, (err: any) => {