chunks v1
This commit is contained in:
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)
|
||||
|
||||
Reference in New Issue
Block a user