Encrypt and Transfer Files over websocket

I’m developing a service for personal use, where i can exchange public key between 2 clients, and transfer encrypted files. I am new to websocket, so i don’t know what i am doing wrong here so the file isn’t encrypted when received on the other end.
So what i am doing here is:
1- generating private/public rsa key
2- sending it to be stored in websocket server
3- fetch public key by the other party
4- use the public key to encrypt the aes
5- encrypt the file with aes
6- send the file
7- Try to decrypt the aes with the private key => FAILING

Client Side

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Secure File Transfer</title>
</head>
<body>
<h2>Secure File Transfer</h2>
<div>
    <label for="clientId">Your Client ID:</label>
    <input type="text" id="clientId">
    <button onclick="setClientId()">Set Client ID</button>
    <p>Client ID: <span id="displayedClientId"></span></p>
</div>
<div>
    <label for="fileInput">Select File:</label>
    <input type="file" id="fileInput">
    <button onclick="handleFileUpload()">Upload File</button>
</div>
<div>
    <label for="targetClientId">Target Client ID:</label>
    <input type="text" id="targetClientId">
</div>
<div id="status"></div>
<div id="progress"></div>

<script>
    let socket;
    let keyPair;
    let targetPublicKey;
    let clientId;
    let receivedChunks = {};
    let totalFileSize = 0;
    let aesKey; // Global variable for AES key

    async function initializeWebSocket() {
        socket = new WebSocket('ws://localhost:1111'); // WebSocket server URL

        socket.onopen = async function() {
            console.log('Connected to WebSocket server.');

            // Generate RSA key pair
            try {
                keyPair = await generateRSAKeyPair();
                setStatus('RSA key pair generated.');

                if (clientId) {
                    await sendPublicKey();
                }
            } catch (error) {
                console.error('Error generating RSA key pair:', error);
                setStatus('Error generating RSA key pair.');
            }
        };

        socket.onmessage = async function(event) {
            const message = JSON.parse(event.data);
            console.log('Received message:', message);

            if (message.type === 'publicKey') {
                if (message.clientId !== clientId) {
                    // Received the target client's public key
                    targetPublicKey = await importPublicKey(message.publicKey);
                    setStatus('Received public key from target client.');
                }
            } else if (message.type === 'file') {
                await handleIncomingFileChunk(message);
            } else if (message.type === 'error') {
                setStatus('Error: ' + message.message);
            }
        };

        socket.onclose = function() {
            console.log('Disconnected from WebSocket server.');
        };

        socket.onerror = function(error) {
            console.error('WebSocket error:', error);
            setStatus('WebSocket error: ' + error.message);
        };
    }

    function setClientId() {
        clientId = document.getElementById('clientId').value;
        document.getElementById('displayedClientId').textContent = clientId;

        if (!keyPair) {
            setStatus('Generating RSA key pair...');
            generateRSAKeyPair().then(generatedKeyPair => {
                keyPair = generatedKeyPair;
                sendPublicKey();
            }).catch(error => {
                console.error('Error generating RSA key pair:', error);
                setStatus('Error generating RSA key pair.');
            });
        } else {
            sendPublicKey();
        }
    }

    async function sendPublicKey() {
        if (!keyPair || !keyPair.publicKey) {
            setStatus('RSA key pair not generated yet.');
            return;
        }

        const publicKeyString = await exportPublicKey(keyPair.publicKey);
        socket.send(JSON.stringify({ type: 'publicKey', clientId: clientId, publicKey: publicKeyString }));
        setStatus('Public key sent to server.');
    }

    async function handleFileUpload() {
        const fileInput = document.getElementById('fileInput');
        const files = fileInput.files;
        const targetClientId = document.getElementById('targetClientId').value;

        if (!files || files.length === 0) {
            setStatus('No file selected.');
            return;
        }

        if (!targetClientId) {
            setStatus('Target client ID is required.');
            return;
        }

        if (!targetPublicKey) {
            setStatus('Requesting public key from target client...');
            socket.send(JSON.stringify({ type: 'requestPublicKey', targetClientId: targetClientId }));
            return;
        }

        const file = files[0];
        totalFileSize = file.size;
        const chunkSize = 245;
        let offset = 0;

        try {
            // Generate AES key for encrypting file chunks
            aesKey = await generateAESKey();

            // Encrypt AES key with RSA
            const encryptedAesKey = await encryptAESKeyWithRSA(aesKey, targetPublicKey);

            setStatus('Uploading file...');
            setProgress(0);

            while (offset < file.size) {
                const chunk = await readChunk(file, offset, chunkSize);
                const encryptedChunk = await encryptChunkWithAES(chunk);

                const message = {
                    type: 'file',
                    targetClientId: targetClientId,
                    filename: file.name,
                    fileType: file.type,
                    offset: offset,
                    chunkSize: encryptedChunk.length,
                    encryptedData: encryptedChunk,
                    totalFileSize: file.size,
                    aesKey: encryptedAesKey // Include encrypted AES key
                };

                socket.send(JSON.stringify(message));

                offset += chunkSize;
                const percent = Math.min((offset / file.size) * 100, 100);
                setProgress(percent);
            }

            setStatus('File uploaded successfully.');
        } catch (error) {
            console.error('Error uploading file:', error);
            setStatus('Error uploading file: ' + error.message);
        }
    }

    async function readChunk(file, offset, length) {
        return new Promise((resolve, reject) => {
            const reader = new FileReader();
            reader.onload = function(event) {
                const chunk = new Uint8Array(event.target.result);
                resolve(chunk);
            };
            reader.onerror = function(event) {
                reject(event.target.error);
            };

            const blob = file.slice(offset, offset + length);
            reader.readAsArrayBuffer(blob);
        });
    }

    async function handleIncomingFileChunk(message) {
        const { filename, fileType, offset, encryptedData, totalFileSize, aesKey } = message;
        console.log('Received file chunk:', message);

        try {

            // Decrypt AES key with RSA private key
            const decryptedAesKey = await decryptAESKeyWithRSA(aesKey);

            // Decrypt chunk data using decrypted AES key
            const decryptedBytes = await decryptChunkWithAES(encryptedData, decryptedAesKey);

            // Store decrypted chunk
            receivedChunks[offset] = new Uint8Array(decryptedBytes);

            // Check if all chunks are received
            if (offset + 245 >= totalFileSize) {
                const combinedChunks = combineChunks(receivedChunks, totalFileSize);
                handleCompleteFile({
                    filename: filename,
                    fileType: fileType,
                    combinedChunks: combinedChunks
                });
            }

            setStatus(`Received chunk for file: ${filename}`);
        } catch (error) {
            console.error('Error handling incoming file chunk:', error);
            setStatus('Error handling incoming file chunk: ' + error.message);
        }
    }


    async function generateRSAKeyPair() {
        const keyPair = await window.crypto.subtle.generateKey(
            {
                name: 'RSA-OAEP',
                modulusLength: 2048,
                publicExponent: new Uint8Array([0x01, 0x00, 0x01]), // 65537
                hash: { name: 'SHA-256' }
            },
            true,
            ['encrypt', 'decrypt']
        );
        return keyPair;
    }

    async function exportPublicKey(publicKey) {
        const exportedKey = await window.crypto.subtle.exportKey(
            'spki',
            publicKey
        );
        const exportedKeyArray = new Uint8Array(exportedKey);
        return btoa(String.fromCharCode.apply(null, exportedKeyArray));
    }

    async function importPublicKey(publicKeyString) {
        const publicKeyBuffer = Uint8Array.from(atob(publicKeyString), c => c.charCodeAt(0));
        const importedKey = await window.crypto.subtle.importKey(
            'spki',
            publicKeyBuffer.buffer,
            {
                name: 'RSA-OAEP',
                hash: { name: 'SHA-256' }
            },
            true,
            ['encrypt']
        );
        return importedKey;
    }

    async function encryptAESKeyWithRSA(aesKey, publicKey) {
        const exportedKey = await window.crypto.subtle.exportKey(
            'raw',
            aesKey
        );

        const encryptedKey = await window.crypto.subtle.encrypt(
            {
                name: 'RSA-OAEP'
            },
            publicKey,
            exportedKey
        );

        return new Uint8Array(encryptedKey);
    }

    async function decryptAESKeyWithRSA(encryptedAesKey) {
        // Ensure keyPair is defined and has privateKey
        if (!keyPair || !keyPair.privateKey) {
            throw new Error('RSA private key not available.');
        }

        try {
            const encryptedAesKeyArray = new Uint8Array(encryptedAesKey);

            const decryptedAesKey = await window.crypto.subtle.decrypt(
                {
                    name: 'RSA-OAEP'
                },
                keyPair.privateKey,
                encryptedAesKeyArray
            );

            return decryptedAesKey;
        } catch (error) {
            console.error('Error decrypting AES key with RSA:', error);
            throw new Error('Failed to decrypt AES key with RSA.');
        }
    }


    async function generateAESKey() {
        const key = await window.crypto.subtle.generateKey(
            {
                name: 'AES-GCM',
                length: 256
            },
            true,
            ['encrypt', 'decrypt']
        );
        return key;
    }

    async function encryptChunkWithAES(chunk) {
        const iv = window.crypto.getRandomValues(new Uint8Array(12));

        const encrypted = await window.crypto.subtle.encrypt(
            {
                name: 'AES-GCM',
                iv: iv
            },
            aesKey,
            chunk
        );

        // Concatenate IV and encrypted data
        const encryptedArray = new Uint8Array(iv.byteLength + encrypted.byteLength);
        encryptedArray.set(new Uint8Array(iv), 0);
        encryptedArray.set(new Uint8Array(encrypted), iv.byteLength);

        return encryptedArray;
    }

    async function decryptChunkWithAES(encryptedChunk, aesKey) {
        const iv = encryptedChunk.slice(0, 12);
        const encryptedData = encryptedChunk.slice(12);

        const decrypted = await window.crypto.subtle.decrypt(
            {
                name: 'AES-GCM',
                iv: iv
            },
            aesKey,
            encryptedData
        );

        return new Uint8Array(decrypted);
    }

    function combineChunks(chunks, totalSize) {
        const combined = new Uint8Array(totalSize);
        let offset = 0;

        for (let chunkOffset in chunks) {
            const chunk = chunks[chunkOffset];
            combined.set(chunk, parseInt(chunkOffset));
            offset += chunk.length;
        }

        return combined;
    }

    function handleCompleteFile({ filename, fileType, combinedChunks }) {
        const blob = new Blob([combinedChunks], { type: fileType });
        downloadFile(blob, filename);
        setStatus(`File ${filename} downloaded successfully.`);
    }

    function downloadFile(blob, fileName) {
        const url = URL.createObjectURL(blob);
        const a = document.createElement('a');
        a.href = url;
        a.download = fileName || 'file';
        document.body.appendChild(a);
        a.click();
        document.body.removeChild(a);
        URL.revokeObjectURL(url);
    }

    function setStatus(message) {
        document.getElementById('status').textContent = message;
    }

    function setProgress(percent) {
        document.getElementById('progress').textContent = `Progress: ${percent}%`;
    }

    window.onload = initializeWebSocket;
</script>
</body>
</html>

Server Side

const WebSocket = require('ws');

// Initialize WebSocket server
const wss = new WebSocket.Server({ port: 1111 });

// Client storage
const clients = {};

// Handle incoming WebSocket connections
wss.on('connection', function connection(ws) {
  console.log('Client connected.');

  // Handle incoming messages from clients
  ws.on('message', function incoming(message) {
    try {
      const data = JSON.parse(message);
      console.log('Received message:', data);

      if (data.type === 'publicKey') {
        // Store client's public key
        clients[data.clientId] = {
          publicKey: data.publicKey,
          ws: ws,
        };
        console.log(`Stored public key for client ${data.clientId}`);
      } else if (data.type === 'requestPublicKey') {
        // Send requested client's public key to requester
        const targetClientId = data.targetClientId;
        const targetClient = clients[targetClientId];

        if (targetClient && targetClient.publicKey) {
          ws.send(JSON.stringify({ type: 'publicKey', clientId: targetClientId, publicKey: targetClient.publicKey }));
          console.log(`Sent public key of client ${targetClientId} to client ${data.clientId}`);
        } else {
          ws.send(JSON.stringify({ type: 'error', message: `Client ${targetClientId} not found or public key not available.` }));
        }
      } else if (data.type === 'file') {
        // Route file to target client
        const targetClientId = data.targetClientId;
        const targetClient = clients[targetClientId];

        if (targetClient && targetClient.ws) {
          targetClient.ws.send(JSON.stringify(data));
          console.log(`File routed to client ${targetClientId}`);
        } else {
          ws.send(JSON.stringify({ type: 'error', message: `Client ${targetClientId} not found or offline.` }));
        }
      } else {
        ws.send(JSON.stringify({ type: 'error', message: `Invalid message type or action.` }));
      }
    } catch (error) {
      console.error('Error parsing or processing message:', error);
    }
  });

  // Handle WebSocket close
  ws.on('close', function close() {
    console.log('Client disconnected.');
    // Clean up client's stored public key if exists
    Object.keys(clients).forEach(clientId => {
      if (clients[clientId].ws === ws) {
        delete clients[clientId];
        console.log(`Removed client ${clientId} from storage.`);
      }
    });
  });

  // Handle initial client connection setup
  ws.on('error', function error(err) {
    console.error('WebSocket error:', err);
  });
});

// Start WebSocket server
console.log('WebSocket server is running on port 1111');

I tried to base64 it instead of sending it as an buffer array

Filed under: Kiến thức lập trình - @ 05:29

Thẻ: javascriptnode.jswebsocket

Thiết kế website giá rẻ

Danh mục

Encrypt and Transfer Files over websocket