File size: 9,842 Bytes
6d69a4c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
66919e5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6d69a4c
66919e5
 
 
 
 
 
 
 
 
 
6d69a4c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e8313c1
 
6d69a4c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
import initWasm, {
    decrypt_serialized_u64_radix_flat_wasm
} from './concrete-ml-extensions-wasm/concrete_ml_extensions_wasm.js';

const SERVER = 'https://backend-server.com';

let clientKey, serverKey;
let encTokens;
let encServerResult;
let keygenWorker;
let encryptWorker;
let sessionUid;

// Memory-efficient base64 encoding for large Uint8Array
function uint8ToBase64(uint8) {
    return new Promise((resolve, reject) => {
        const blob = new Blob([uint8]);
        const reader = new FileReader();
        reader.onload = function () {
            const base64 = reader.result.split(',')[1];
            resolve(base64);
        };
        reader.onerror = reject;
        reader.readAsDataURL(blob);
    });
}

const $ = id => document.getElementById(id);
const enable = (id, ok=true) => $(id).disabled = !ok;
const show = (id, visible=true) => $(id).hidden = !visible;

// Hide all spinners immediately
show('keygenSpin', false);
show('spin', false);
show('encIcon', false);
show('tokenizerSpin', false);

// Initialize WASM
(async () => {
    try {
        console.log('[Main] Initializing WASM module...');
        await initWasm();
        console.log('[Main] WASM module initialized successfully');

        // Initialize the keygen worker
        keygenWorker = new Worker(new URL('./keygen-worker.js', import.meta.url), { type: 'module' });
        keygenWorker.onmessage = async function(e) {
            if (e.data.type === 'success') {
                const res = e.data.result;
                console.log('[Main] Key generation successful');
                console.log(`[Main] Client key size: ${res.clientKey.length} bytes`);
                console.log(`[Main] Server key size: ${res.serverKey.length} bytes`);
                clientKey = res.clientKey; serverKey = res.serverKey;
                
                try {
                    // Initialize encryption worker
                    encryptWorker = new Worker(new URL('./encrypt-worker.js', import.meta.url), { type: 'module' });
                    encryptWorker.onmessage = function(e) {
                        if (e.data.type === 'ready') {
                            console.log('[Main] Encryption worker ready');
                        } else if (e.data.type === 'success') {
                            encTokens = e.data.result;
                            console.log(`[Main] Encryption completed: ${encTokens.length} bytes`);
                            show('encryptSpin', false);
                            show('encIcon', true);
                            enable('btnEncrypt', true);
                            enable('btnSend');
                            $('encStatus').textContent = 'Your text is encrypted 🔒';
                        } else if (e.data.type === 'error') {
                            console.error('[Main] Encryption error:', e.data.error);
                            show('encryptSpin', false);
                            enable('btnEncrypt', true);
                            $('encStatus').textContent = `Encryption failed: ${e.data.error}`;
                            alert(`Encryption failed: ${e.data.error}`);
                        }
                    };
                    
                    // Initialize the worker with the client key
                    encryptWorker.postMessage({ type: 'init', clientKey });
                    
                    console.log('[Main] Encoding server key to base64...');
                    const serverKeyB64 = await uint8ToBase64(serverKey);
                    console.log(`[Main] Server key base64 size: ${serverKeyB64.length} bytes`);
                    const handshakeResponse = await fetch(`${SERVER}/handshake`, {
                        method:'POST', 
                        headers:{'Content-Type':'application/json'},
                        body: JSON.stringify({
                            server_key_b64: serverKeyB64
                        })
                    });
                    
                    if (!handshakeResponse.ok) {
                        const errorText = await handshakeResponse.text();
                        throw new Error(`Server handshake failed: ${handshakeResponse.status} ${errorText}`);
                    }
                    
                    const { uid } = await handshakeResponse.json();
                    sessionUid = uid;
                    console.log('[Main] Server handshake successful');
                    $('keygenStatus').textContent = 'keys generated ✓';
                    enable('btnEncrypt');
                } catch (error) {
                    console.error('[Main] Server handshake error:', error);
                    $('keygenStatus').textContent = `Server handshake failed: ${error.message}`;
                    enable('btnEncrypt', false);
                }
            } else {
                console.error('[Main] Key generation error:', e.data.error);
                $('keygenStatus').textContent = `Error generating keys: ${e.data.error}`;
            }
            show('keygenSpin', false);
        };
    } catch (e) {
        console.error('[Main] Failed to initialize WASM module:', e);
        $('keygenStatus').textContent = `Initialization Error: ${e.message}`;
        throw e;
    }
})();

$('btnKeygen').onclick = async () => {
    if ($('keygenSpin').hidden === false) {
        console.log('[Main] Keygen already in progress, ignoring click');
        return;
    }
    show('keygenSpin', true);
    $('keygenStatus').textContent = 'generating…';
    try {
        keygenWorker.postMessage({});
    } catch (e) {
        console.error('[Main] Key generation error:', e);
        $('keygenStatus').textContent = `Error generating keys: ${e.message}`;
        show('keygenSpin', false);
    }
};

$('btnEncrypt').onclick = async () => {
    const text = $('tokenInput').value.trim();
    if (!text) {
        console.error('[Main] No text provided for tokenization/encryption');
        alert('Please enter text to encrypt.');
        return;
    }
    if (!encryptWorker) {
        console.error('[Main] Encryption worker not initialized');
        alert('Encryption worker is not ready. Please generate keys first.');
        return;
    }

    show('encryptSpin', true);
    show('encIcon', false);
    enable('btnEncrypt', false);

    try {
        console.log('[Main] Tokenizing text:', text);
        const tokenIds = llama3Tokenizer.encode(text);
        console.log('[Main] Token IDs:', tokenIds);

        encryptWorker.postMessage({ type: 'encrypt', tokenIds });
    } catch (error) {
        console.error('[Main] Tokenization or encryption initiation error:', error);
        show('encryptSpin', false);
        enable('btnEncrypt', true);
        alert(`Error during tokenization/encryption: ${error.message}`);
    }
};

$('btnSend').onclick = async () => {
    if ($('spin').hidden === false) {
        console.log('[Main] Send already in progress, ignoring click');
        return;
    }
    show('encIcon', false);
    show('spin', true);
    $('srvStatus').textContent = 'Preparing request…';
    $('srvComputing').hidden = true;
    const startTime = performance.now();
    try {
        const ciphertext_b64 = await uint8ToBase64(encTokens);

        $('srvStatus').textContent = 'Sending…';
        console.log('[Main] Sending request to server...');
        const response = await fetch(`${SERVER}/detect`, {
            method:'POST', 
            headers:{'Content-Type':'application/json'},
            body: JSON.stringify({
                uid: sessionUid,
                ciphertext_b64
            })
        }).catch(error => {
            throw new Error(`Server unreachable: ${error.message}`);
        });
        
        console.log(`[Main] Server response status: ${response.status}`);
        if (!response.ok) {
            const errorText = await response.text();
            throw new Error(`Server error: ${response.status} ${errorText}`);
        }

        $('srvStatus').textContent = '✓ sent';
        $('srvComputing').hidden = false;
        
        const {result_b64} = await response.json();
        const endTime = performance.now();
        const duration = ((endTime - startTime) / 1000).toFixed(2);
        console.log(`[Main] Server request completed in ${duration}s`);
        console.log('[Main] Processing server response...');
        encServerResult = Uint8Array.from(atob(result_b64), c=>c.charCodeAt(0));
        console.log(`[Main] Received encrypted result: ${encServerResult.length} bytes`);
        $('encResult').value = `(${encServerResult.length} B)`;
        $('srvComputing').hidden = true;
        $('srvStatus').textContent = `✓ received (${duration}s)`;
        enable('btnDecrypt');
    } catch (e) {
        const endTime = performance.now();
        const duration = ((endTime - startTime) / 1000).toFixed(2);
        console.error(`[Main] Server request failed after ${duration}s:`, e);
        $('srvComputing').hidden = true;
        $('srvStatus').textContent = `Server error: ${e.message} (${duration}s)`;
        show('spin', false);
    }
};

$('btnDecrypt').onclick = () => {
    try {
        console.log('[Main] Starting decryption...');
        const dec = decrypt_serialized_u64_radix_flat_wasm(encServerResult, clientKey);
        const [flag, score_scaled, total_g] = Array.from(dec);
        const score = (Number(score_scaled) / 1e6).toFixed(6);
        console.log('[Main] Decryption successful');
        console.log(`[Main] Result - flag: ${flag}, score: ${score}, total_g: ${total_g}`);
        $('decResult').textContent = `flag=${flag}, score=${score}, total_g=${total_g}`;
    } catch (e) {
        console.error('[Main] Decryption error:', e);
        $('decResult').textContent = `Error decrypting result: ${e.message}`;
    }
};