chunks v1

This commit is contained in:
andrei-mihnea-cerbu
2024-11-12 12:32:29 +02:00
parent d84cb3b18d
commit 8c88d73665
8 changed files with 240 additions and 126 deletions
BIN
View File
Binary file not shown.
@@ -1,11 +1,11 @@
import { Socket } from 'net';
import { privateEncrypt, privateDecrypt, generateKeyPairSync, createCipheriv, createDecipheriv, randomBytes } from 'crypto';
import { MessageHandler, ParsedMessage } from '../message_handler';
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';
import {constants} from "node:crypto";
const END_OF_MESSAGE = '<EOM>'; // Define a unique marker for end of message
const CHUNK_SIZE = 1024; // Define chunk size
export class TcpServerCommunicator extends SocketCommunicatorBase {
private readonly socket: Socket;
@@ -14,6 +14,7 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
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);
@@ -23,6 +24,7 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
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
}
@@ -35,7 +37,6 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
});
this.privateKey = privateKey;
this.publicKey = publicKey;
console.log('RSA key pair generated.');
}
// Send the server's public key to the client
@@ -44,9 +45,7 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
throw new Error('Public key is not available. Please generate RSA key pair.');
}
const message = MessageHandler.formatMessage('SET_PUBLIC_KEY', { publicKey: this.publicKey });
await this.writeToSocket(message + END_OF_MESSAGE);
console.log('Public key sent to client.');
await this.sendChunkedMessage('SET_PUBLIC_KEY', { publicKey: this.publicKey });
}
// Generate AES key and IV, then send them to the client
@@ -57,16 +56,11 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
const aesKeyBase64 = this.aesKey.toString('base64');
const aesIvBase64 = this.aesIv.toString('base64');
const formattedMessage = MessageHandler.formatMessage('SET_AES_KEY', { aesKey: aesKeyBase64, aesIv: aesIvBase64 });
const encryptedMessage = this.encryptWithRsa(Buffer.from(formattedMessage)).toString('base64');
await this.writeToSocket(encryptedMessage + END_OF_MESSAGE);
console.log('AES key and IV sent to client.');
await this.sendChunkedMessage('SET_AES_KEY', { aesKey: aesKeyBase64, aesIv: aesIvBase64 });
}
// Encrypt a message with the server's private key (RSA encryption)
private encryptWithRsa(message: Buffer): Buffer {
const bufferMessage = Buffer.isBuffer(message) ? message : Buffer.from(message);
private encryptWithRsa(message: string): string {
if (!this.privateKey) throw new Error('Server private key not set.');
return privateEncrypt(
@@ -74,8 +68,8 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
key: this.privateKey,
padding: constants.RSA_PKCS1_PADDING, // PKCS1 padding
},
bufferMessage
);
Buffer.from(message)
).toString('base64');
}
// Decrypt AES-encrypted messages
@@ -106,17 +100,34 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
const incomingMessage = data.toString();
this.messageBuffer += incomingMessage;
// Check if the message ends with END_OF_MESSAGE
if (this.messageBuffer.endsWith(END_OF_MESSAGE)) {
// Remove the END_OF_MESSAGE marker and process the message
const completeMessage = this.messageBuffer.slice(0, -END_OF_MESSAGE.length);
if (this.messageBuffer.includes(END_OF_MESSAGE)) {
const messages = this.messageBuffer.split(END_OF_MESSAGE);
console.log(`\n\nComplete Message:\n${completeMessage}\n\n`);
for (let i = 0; i < messages.length - 1; i++) {
const completeMessage = messages[i];
if (completeMessage) {
this.processCompleteMessage(completeMessage);
}
}
this.handleIncomingMessage(completeMessage);
this.messageBuffer = messages[messages.length - 1];
}
}
// Clear the message buffer after processing
this.messageBuffer = '';
private processCompleteMessage(completeMessage: string): 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('');
this.handleIncomingMessage(fullMessage);
delete this.chunkBuffers[header.messageId];
}
}
@@ -125,14 +136,29 @@ export class TcpServerCommunicator extends SocketCommunicatorBase {
const message = MessageHandler.formatMessage(operationCode, metaInfo, fileContent);
let outgoingMessage: string;
if (this.aesKey && this.aesIv) {
outgoingMessage = this.encryptWithAes(message);
} else {
outgoingMessage = message;
}
if (this.aesKey && this.aesIv) outgoingMessage = this.encryptWithAes(message);
else if(this.privateKey) outgoingMessage = this.encryptWithRsa(message);
else outgoingMessage = message;
outgoingMessage += END_OF_MESSAGE; // Append end-of-message marker
await this.writeToSocket(outgoingMessage);
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);
}
}
}
// Handle incoming message (decrypt with AES if available)