jfrery-zama commited on
Commit
af99c56
·
unverified ·
1 Parent(s): 75bf83f

compatible with server

Browse files
Files changed (1) hide show
  1. wasm-demo.js +134 -81
wasm-demo.js CHANGED
@@ -2,7 +2,7 @@ import initWasm, {
2
  decrypt_serialized_u64_radix_flat_wasm
3
  } from './concrete-ml-extensions-wasm/concrete_ml_extensions_wasm.js';
4
 
5
- const SERVER = 'https://backend-server.com';
6
 
7
  let clientKey, serverKey;
8
  let encTokens;
@@ -10,6 +10,7 @@ let encServerResult;
10
  let keygenWorker;
11
  let encryptWorker;
12
  let sessionUid;
 
13
 
14
  // Memory-efficient base64 encoding for large Uint8Array
15
  function uint8ToBase64(uint8) {
@@ -78,37 +79,43 @@ show('tokenizerSpin', false);
78
  // Initialize the worker with the client key
79
  encryptWorker.postMessage({ type: 'init', clientKey });
80
 
81
- console.log('[Main] Encoding server key to base64...');
82
- const serverKeyB64 = await uint8ToBase64(serverKey);
83
- console.log(`[Main] Server key base64 size: ${serverKeyB64.length} bytes`);
84
- const handshakeResponse = await fetch(`${SERVER}/handshake`, {
85
- method:'POST',
86
- headers:{'Content-Type':'application/json'},
87
- body: JSON.stringify({
88
- server_key_b64: serverKeyB64
89
- })
 
 
 
 
90
  });
91
-
92
- if (!handshakeResponse.ok) {
93
- const errorText = await handshakeResponse.text();
94
- throw new Error(`Server handshake failed: ${handshakeResponse.status} ${errorText}`);
95
  }
96
-
97
- const { uid } = await handshakeResponse.json();
98
  sessionUid = uid;
99
- console.log('[Main] Server handshake successful');
100
- $('keygenStatus').textContent = 'keys generated ✓';
101
  enable('btnEncrypt');
102
  } catch (error) {
103
- console.error('[Main] Server handshake error:', error);
104
- $('keygenStatus').textContent = `Server handshake failed: ${error.message}`;
105
  enable('btnEncrypt', false);
 
 
106
  }
107
  } else {
108
  console.error('[Main] Key generation error:', e.data.error);
109
  $('keygenStatus').textContent = `Error generating keys: ${e.data.error}`;
 
110
  }
111
- show('keygenSpin', false);
112
  };
113
  } catch (e) {
114
  console.error('[Main] Failed to initialize WASM module:', e);
@@ -164,79 +171,125 @@ $('btnEncrypt').onclick = async () => {
164
  }
165
  };
166
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
167
  $('btnSend').onclick = async () => {
168
  if ($('spin').hidden === false) {
169
- console.log('[Main] Send already in progress, ignoring click');
 
 
 
 
170
  return;
171
  }
 
172
  show('encIcon', false);
173
  show('spin', true);
174
- $('srvStatus').textContent = 'Preparing request…';
175
  $('srvComputing').hidden = true;
176
- $('srvProgress').hidden = true;
177
- const startTime = performance.now();
178
  try {
179
- const ciphertext_b64 = await uint8ToBase64(encTokens);
180
-
181
- $('srvStatus').textContent = 'Sending…';
182
- console.log('[Main] Sending request to server...');
183
- const response = await fetch(`${SERVER}/detect`, {
184
- method:'POST',
185
- headers:{'Content-Type':'application/json'},
186
- body: JSON.stringify({
187
- uid: sessionUid,
188
- ciphertext_b64
189
- })
190
- }).catch(error => {
191
- throw new Error(`Server unreachable: ${error.message}`);
192
  });
193
-
194
- console.log(`[Main] Server response status: ${response.status}`);
195
- if (!response.ok) {
196
- const errorText = await response.text();
197
- throw new Error(`Server error: ${response.status} ${errorText}`);
198
  }
199
 
200
- $('srvStatus').textContent = '✓ sent';
 
 
 
201
  $('srvComputing').hidden = false;
202
- $('srvProgress').hidden = false;
203
-
204
- const tokenCount = llama3Tokenizer.encode($('tokenInput').value.trim()).length;
205
- const estimatedTime = tokenCount * 15;
206
- const endTime = startTime + (estimatedTime * 1000);
207
-
208
- const progressInterval = setInterval(() => {
209
- const now = performance.now();
210
- const elapsed = (now - startTime) / 1000;
211
- const remaining = Math.max(0, estimatedTime - elapsed);
212
- $('srvProgress').textContent = `Estimated time remaining: ${Math.ceil(remaining)}s`;
213
-
214
- if (remaining <= 0) {
215
- clearInterval(progressInterval);
216
  }
217
- }, 1000);
218
-
219
- const {result_b64} = await response.json();
220
- clearInterval(progressInterval);
221
- const actualEndTime = performance.now();
222
- const duration = ((actualEndTime - startTime) / 1000).toFixed(2);
223
- console.log(`[Main] Server request completed in ${duration}s`);
224
- console.log('[Main] Processing server response...');
225
- encServerResult = Uint8Array.from(atob(result_b64), c=>c.charCodeAt(0));
226
- console.log(`[Main] Received encrypted result: ${encServerResult.length} bytes`);
227
- $('encResult').value = `(${encServerResult.length} B)`;
228
- $('srvComputing').hidden = true;
229
- $('srvProgress').hidden = true;
230
- $('srvStatus').textContent = `✓ received (${duration}s)`;
231
- enable('btnDecrypt');
232
  } catch (e) {
233
- const endTime = performance.now();
234
- const duration = ((endTime - startTime) / 1000).toFixed(2);
235
- console.error(`[Main] Server request failed after ${duration}s:`, e);
236
- $('srvComputing').hidden = true;
237
- $('srvProgress').hidden = true;
238
- $('srvStatus').textContent = `Server error: ${e.message} (${duration}s)`;
239
  show('spin', false);
 
240
  }
241
  };
242
 
@@ -248,9 +301,9 @@ $('btnDecrypt').onclick = () => {
248
  const score = (Number(score_scaled) / 1e6).toFixed(6);
249
  console.log('[Main] Decryption successful');
250
  console.log(`[Main] Result - flag: ${flag}, score: ${score}, total_g: ${total_g}`);
251
- $('decResult').textContent = `flag=${flag}, score=${score}, total_g=${total_g}`;
252
  } catch (e) {
253
  console.error('[Main] Decryption error:', e);
254
- $('decResult').textContent = `Error decrypting result: ${e.message}`;
255
  }
256
  };
 
2
  decrypt_serialized_u64_radix_flat_wasm
3
  } from './concrete-ml-extensions-wasm/concrete_ml_extensions_wasm.js';
4
 
5
+ const SERVER = 'https://api.zama.ai';
6
 
7
  let clientKey, serverKey;
8
  let encTokens;
 
10
  let keygenWorker;
11
  let encryptWorker;
12
  let sessionUid;
13
+ let taskId;
14
 
15
  // Memory-efficient base64 encoding for large Uint8Array
16
  function uint8ToBase64(uint8) {
 
79
  // Initialize the worker with the client key
80
  encryptWorker.postMessage({ type: 'init', clientKey });
81
 
82
+ console.log('[Main] Sending server key to server...');
83
+ $('keygenStatus').textContent = 'Keys generated, sending server key...';
84
+ show('keygenSpin', true);
85
+
86
+ const formData = new FormData();
87
+ const serverKeyBlob = new Blob([serverKey], { type: 'application/octet-stream' });
88
+ const serverKeyFile = new File([serverKeyBlob], "server.key");
89
+ formData.append('key', serverKeyFile);
90
+ formData.append('task_name', 'synthid');
91
+
92
+ const addKeyResponse = await fetch(`${SERVER}/add_key`, {
93
+ method: 'POST',
94
+ body: formData
95
  });
96
+
97
+ if (!addKeyResponse.ok) {
98
+ const errorText = await addKeyResponse.text();
99
+ throw new Error(`Server /add_key failed: ${addKeyResponse.status} ${errorText}`);
100
  }
101
+
102
+ const { uid } = await addKeyResponse.json();
103
  sessionUid = uid;
104
+ console.log('[Main] Server key sent and UID received:', sessionUid);
105
+ $('keygenStatus').textContent = 'Keys generated & UID received ✓';
106
  enable('btnEncrypt');
107
  } catch (error) {
108
+ console.error('[Main] Server key submission error:', error);
109
+ $('keygenStatus').textContent = `Server key submission failed: ${error.message}`;
110
  enable('btnEncrypt', false);
111
+ } finally {
112
+ show('keygenSpin', false);
113
  }
114
  } else {
115
  console.error('[Main] Key generation error:', e.data.error);
116
  $('keygenStatus').textContent = `Error generating keys: ${e.data.error}`;
117
+ show('keygenSpin', false);
118
  }
 
119
  };
120
  } catch (e) {
121
  console.error('[Main] Failed to initialize WASM module:', e);
 
171
  }
172
  };
173
 
174
+ async function pollTaskStatus(currentTaskId, currentUid) {
175
+ try {
176
+ const statusResponse = await fetch(`${SERVER}/get_task_status?task_id=${currentTaskId}&uid=${currentUid}`);
177
+ if (!statusResponse.ok) {
178
+ const errorText = await statusResponse.text();
179
+ console.error(`[Poll] Error fetching status: ${statusResponse.status} ${errorText}`);
180
+ $('srvStatus').textContent = `Status check error: ${statusResponse.status}`;
181
+ show('spin', false);
182
+ return null;
183
+ }
184
+
185
+ const statusData = await statusResponse.json();
186
+ console.log('[Poll] Task status:', statusData);
187
+ $('srvStatus').textContent = `Status: ${statusData.status} - ${statusData.details}`;
188
+
189
+ if (statusData.status === 'success' || statusData.status === 'completed') {
190
+ return statusData;
191
+ } else if (['failure', 'revoked', 'unknown', 'error'].includes(statusData.status.toLowerCase())) {
192
+ console.error('[Poll] Task failed or unrecoverable:', statusData);
193
+ $('srvStatus').textContent = `Task failed: ${statusData.status}`;
194
+ show('spin', false);
195
+ return null;
196
+ } else {
197
+ setTimeout(() => pollTaskStatus(currentTaskId, currentUid).then(finalStatus => {
198
+ if (finalStatus && (finalStatus.status === 'success' || finalStatus.status === 'completed')) {
199
+ getTaskResult(currentTaskId, currentUid, 'synthid');
200
+ }
201
+ }), 5000);
202
+ return null;
203
+ }
204
+ } catch (e) {
205
+ console.error('[Poll] Polling exception:', e);
206
+ $('srvStatus').textContent = `Polling error: ${e.message}`;
207
+ show('spin', false);
208
+ return null;
209
+ }
210
+ }
211
+
212
+ async function getTaskResult(currentTaskId, currentUid, taskName) {
213
+ $('srvStatus').textContent = 'Fetching result...';
214
+ try {
215
+ const resultResponse = await fetch(`${SERVER}/get_task_result?task_name=${taskName}&task_id=${currentTaskId}&uid=${currentUid}`);
216
+
217
+ if (!resultResponse.ok) {
218
+ const errorText = await resultResponse.text();
219
+ throw new Error(`Server /get_task_result error: ${resultResponse.status} ${errorText}`);
220
+ }
221
+
222
+ const resultArrayBuffer = await resultResponse.arrayBuffer();
223
+ encServerResult = new Uint8Array(resultArrayBuffer);
224
+
225
+ console.log(`[Main] Received encrypted result: ${encServerResult.length} bytes`);
226
+ $('encResult').value = `(${encServerResult.length} B)`;
227
+ $('srvStatus').textContent = `✓ result received (${((performance.now() - window.taskStartTime) / 1000).toFixed(2)}s total)`;
228
+ enable('btnDecrypt');
229
+
230
+ } catch (e) {
231
+ const duration = window.taskStartTime ? ((performance.now() - window.taskStartTime) / 1000).toFixed(2) : 'N/A';
232
+ console.error(`[Main] /get_task_result failed after ${duration}s:`, e);
233
+ $('srvStatus').textContent = `Result fetch error: ${e.message} (${duration}s)`;
234
+ } finally {
235
+ show('spin', false);
236
+ $('srvComputing').hidden = true;
237
+ }
238
+ }
239
+
240
  $('btnSend').onclick = async () => {
241
  if ($('spin').hidden === false) {
242
+ console.log('[Main] Task submission/polling already in progress, ignoring click');
243
+ return;
244
+ }
245
+ if (!sessionUid || !encTokens) {
246
+ alert('Please generate keys and encrypt text first.');
247
  return;
248
  }
249
+
250
  show('encIcon', false);
251
  show('spin', true);
252
+ $('srvStatus').textContent = 'Submitting task…';
253
  $('srvComputing').hidden = true;
254
+ window.taskStartTime = performance.now();
255
+
256
  try {
257
+ const formData = new FormData();
258
+ formData.append('uid', sessionUid);
259
+ formData.append('task_name', 'synthid');
260
+
261
+ const encryptedInputBlob = new Blob([encTokens], { type: 'application/octet-stream' });
262
+ const encryptedInputFile = new File([encryptedInputBlob], "input.fheencrypted");
263
+ formData.append('encrypted_input', encryptedInputFile);
264
+
265
+ const startTaskResponse = await fetch(`${SERVER}/start_task`, {
266
+ method: 'POST',
267
+ body: formData
 
 
268
  });
269
+
270
+ if (!startTaskResponse.ok) {
271
+ const errorText = await startTaskResponse.text();
272
+ throw new Error(`Server /start_task error: ${startTaskResponse.status} ${errorText}`);
 
273
  }
274
 
275
+ const { task_id: newTaskId } = await startTaskResponse.json();
276
+ taskId = newTaskId;
277
+ console.log('[Main] Task submitted to server. Task ID:', taskId);
278
+ $('srvStatus').textContent = `Task submitted (ID: ${taskId.substring(0,8)}...). Polling status...`;
279
  $('srvComputing').hidden = false;
280
+
281
+ pollTaskStatus(taskId, sessionUid).then(finalStatus => {
282
+ if (finalStatus && (finalStatus.status === 'success' || finalStatus.status === 'completed')) {
283
+ getTaskResult(taskId, sessionUid, 'synthid');
 
 
 
 
 
 
 
 
 
 
284
  }
285
+ });
286
+
 
 
 
 
 
 
 
 
 
 
 
 
 
287
  } catch (e) {
288
+ const duration = ((performance.now() - window.taskStartTime) / 1000).toFixed(2);
289
+ console.error(`[Main] Task submission failed after ${duration}s:`, e);
290
+ $('srvStatus').textContent = `Task submission error: ${e.message} (${duration}s)`;
 
 
 
291
  show('spin', false);
292
+ $('srvComputing').hidden = true;
293
  }
294
  };
295
 
 
301
  const score = (Number(score_scaled) / 1e6).toFixed(6);
302
  console.log('[Main] Decryption successful');
303
  console.log(`[Main] Result - flag: ${flag}, score: ${score}, total_g: ${total_g}`);
304
+ $('decResult').textContent = `Flag: ${flag}, Score: ${score}, Total G: ${total_g}`;
305
  } catch (e) {
306
  console.error('[Main] Decryption error:', e);
307
+ $('decResult').textContent = `Decryption failed: ${e.message}`;
308
  }
309
  };