fix: use notarized transcript instead of pre-fetched response for red…

…action (#113)

* wip

* fix: use transcript for redaction

* refactor: add map secret to ranges utility

* fix: body selector

* fix

src/entries/Background/plugins/utils.ts

@@ -26,4 +26,18 @@ export async function installPlugin(
     filePath,
   });
   return hash;
+}
+
+export function mapSecretsToRange(secrets: string[], text: string) {
+  return secrets
+    .map((secret: string) => {
+      const index = text.indexOf(secret);
+      return index > -1
+        ? {
+          start: index,
+          end: index + secret.length,
+        }
+        : null;
+    })
+    .filter((data: any) => !!data) as { start: number; end: number }[]
 }

src/entries/Background/rpc.ts

@@ -34,6 +34,7 @@ import {
   hexToArrayBuffer,
   makePlugin,
   PluginConfig,
+  safeParseJSON,
 } from '../../utils/misc';
 import {
   getLoggingFilter,
@@ -46,6 +47,8 @@ import { deferredPromise } from '../../utils/promise';
 import { minimatch } from 'minimatch';
 import { OffscreenActionTypes } from '../Offscreen/types';
 import { SidePanelActionTypes } from '../SidePanel/types';
+import { subtractRanges } from '../Offscreen/utils';
+import { mapSecretsToRange } from './plugins/utils';
 
 const charwise = require('charwise');
 
@@ -382,8 +385,9 @@ async function runPluginProver(request: BackgroundAction, now = Date.now()) {
     method,
     headers,
     body,
-    secretHeaders,
-    secretResps,
+    secretHeaders = [],
+    getSecretResponse,
+    getSecretResponseFn,
     notaryUrl: _notaryUrl,
     websocketProxyUrl: _websocketProxyUrl,
     maxSentData: _maxSentData,
@@ -394,6 +398,8 @@ async function runPluginProver(request: BackgroundAction, now = Date.now()) {
   const maxSentData = _maxSentData || (await getMaxSent());
   const maxRecvData = _maxRecvData || (await getMaxRecv());
 
+  let secretResps: string[] = [];
+
   const { id } = await addNotaryRequest(now, {
     url,
     method,
@@ -417,8 +423,77 @@ async function runPluginProver(request: BackgroundAction, now = Date.now()) {
     action: addRequestHistory(await getNotaryRequest(id)),
   });
 
-  await browser.runtime.sendMessage({
-    type: BackgroundActiontype.process_prove_request,
+  const onProverResponse = async (request: any) => {
+    const { data, type } = request;
+
+    if (type !== OffscreenActionTypes.create_prover_response) {
+      return;
+    }
+
+    if (data.error) {
+      console.error(data.error);
+      return;
+    }
+
+    if (data.id !== id) {
+      return;
+    }
+
+    if (getSecretResponse) {
+      const body = data.transcript.recv.split('\r\n').reduce(
+        (
+          state: { headerEnd: boolean; isBody: boolean; body: string[] },
+          line: string,
+        ) => {
+          if (state.headerEnd) {
+            state.body.push(line);
+          } else if (!line) {
+            state.headerEnd = true;
+          }
+
+          return state;
+        },
+        { headerEnd: false, body: [] },
+      ).body;
+
+      if (body.length == 1) {
+        secretResps = await getSecretResponseFn(body[0]);
+      } else {
+        secretResps = await getSecretResponse(
+          body.filter((txt: string) => {
+            const json = safeParseJSON(txt);
+            return typeof json === 'object';
+          })[0],
+        );
+      }
+    }
+
+    const commit = {
+      sent: subtractRanges(
+        data.transcript.ranges.sent.all,
+        mapSecretsToRange(secretHeaders, data.transcript.sent),
+      ),
+      recv: subtractRanges(
+        data.transcript.ranges.recv.all,
+        mapSecretsToRange(secretResps, data.transcript.recv),
+      ),
+    };
+
+    browser.runtime.sendMessage({
+      type: OffscreenActionTypes.create_presentation_request,
+      data: {
+        id,
+        commit,
+      },
+    });
+
+    browser.runtime.onMessage.removeListener(onProverResponse);
+  };
+
+  browser.runtime.onMessage.addListener(onProverResponse);
+
+  browser.runtime.sendMessage({
+    type: OffscreenActionTypes.create_prover_request,
     data: {
       id,
       url,
@@ -429,8 +504,6 @@ async function runPluginProver(request: BackgroundAction, now = Date.now()) {
       websocketProxyUrl,
       maxRecvData,
       maxSentData,
-      secretHeaders,
-      secretResps,
     },
   });
 }

src/entries/Offscreen/Offscreen.tsx

@@ -14,12 +14,16 @@ import { BackgroundActiontype } from '../Background/rpc';
 import browser from 'webextension-polyfill';
 import { PresentationJSON } from '../../utils/types';
 import { PresentationJSON as PresentationJSONa7 } from 'tlsn-js/build/types';
-import { Method } from 'tlsn-js/wasm/pkg';
+import { Commit, Method } from 'tlsn-wasm';
+import { subtractRanges } from './utils';
+import { mapSecretsToRange } from '../Background/plugins/utils';
 
 const { init, Prover, Presentation }: any = Comlink.wrap(
   new Worker(new URL('./worker.ts', import.meta.url)),
 );
 
+const provers: { [id: string]: TProver } = {};
+
 const Offscreen = () => {
   useEffect(() => {
     (async () => {
@@ -74,6 +78,76 @@ const Offscreen = () => {
 
             break;
           }
+          case OffscreenActionTypes.create_prover_request: {
+            const { id } = request.data;
+
+            (async () => {
+              try {
+                const prover = await createProver(request.data);
+
+                provers[id] = prover;
+
+                browser.runtime.sendMessage({
+                  type: OffscreenActionTypes.create_prover_response,
+                  data: {
+                    id,
+                    transcript: await prover.transcript(),
+                  },
+                });
+              } catch (error) {
+                console.error(error);
+                browser.runtime.sendMessage({
+                  type: OffscreenActionTypes.create_prover_response,
+                  data: {
+                    id,
+                    error,
+                  },
+                });
+              }
+            })();
+            break;
+          }
+          case OffscreenActionTypes.create_presentation_request: {
+            const { id, commit } = request.data;
+            (async () => {
+              const prover = provers[id];
+
+              try {
+                if (!prover) throw new Error(`Cannot find prover ${id}.`);
+
+                const notarizationOutputs = await prover.notarize(commit);
+
+                const presentation = (await new Presentation({
+                  attestationHex: notarizationOutputs.attestation,
+                  secretsHex: notarizationOutputs.secrets,
+                  notaryUrl: notarizationOutputs.notaryUrl,
+                  websocketProxyUrl: notarizationOutputs.websocketProxyUrl,
+                  reveal: commit,
+                })) as TPresentation;
+                const presentationJSON = await presentation.json();
+
+                browser.runtime.sendMessage({
+                  type: BackgroundActiontype.finish_prove_request,
+                  data: {
+                    id,
+                    proof: presentationJSON,
+                  },
+                });
+
+                delete provers[id];
+              } catch (error) {
+                console.error(error);
+                browser.runtime.sendMessage({
+                  type: BackgroundActiontype.finish_prove_request,
+                  data: {
+                    id,
+                    error,
+                  },
+                });
+              }
+            })();
+            break;
+          }
           case BackgroundActiontype.process_prove_request: {
             const { id } = request.data;
 
@@ -141,46 +215,6 @@ const Offscreen = () => {
 
 export default Offscreen;
 
-function subtractRanges(
-  ranges: { start: number; end: number },
-  negatives: { start: number; end: number }[],
-): { start: number; end: number }[] {
-  const returnVal: { start: number; end: number }[] = [ranges];
-
-  negatives
-    .sort((a, b) => (a.start < b.start ? -1 : 1))
-    .forEach(({ start, end }) => {
-      const last = returnVal.pop()!;
-
-      if (start < last.start || end > last.end) {
-        console.error('invalid ranges');
-        return;
-      }
-
-      if (start === last.start && end === last.end) {
-        return;
-      }
-
-      if (start === last.start && end < last.end) {
-        returnVal.push({ start: end, end: last.end });
-        return;
-      }
-
-      if (start > last.start && end < last.end) {
-        returnVal.push({ start: last.start, end: start });
-        returnVal.push({ start: end, end: last.end });
-        return;
-      }
-
-      if (start > last.start && end === last.end) {
-        returnVal.push({ start: last.start, end: start });
-        return;
-      }
-    });
-
-  return returnVal;
-}
-
 async function createProof(options: {
   url: string;
   notaryUrl: string;
@@ -233,31 +267,11 @@ async function createProof(options: {
   const commit = {
     sent: subtractRanges(
       transcript.ranges.sent.all,
-      secretHeaders
-        .map((secret: string) => {
-          const index = transcript.sent.indexOf(secret);
-          return index > -1
-            ? {
-                start: index,
-                end: index + secret.length,
-              }
-            : null;
-        })
-        .filter((data: any) => !!data) as { start: number; end: number }[],
+      mapSecretsToRange(secretHeaders, transcript.sent),
     ),
     recv: subtractRanges(
       transcript.ranges.recv.all,
-      secretResps
-        .map((secret: string) => {
-          const index = transcript.recv.indexOf(secret);
-          return index > -1
-            ? {
-                start: index,
-                end: index + secret.length,
-              }
-            : null;
-        })
-        .filter((data: any) => !!data) as { start: number; end: number }[],
+      mapSecretsToRange(secretResps, transcript.recv),
     ),
   };
 
@@ -270,8 +284,54 @@ async function createProof(options: {
     websocketProxyUrl: notarizationOutputs.websocketProxyUrl,
     reveal: commit,
   })) as TPresentation;
-  const presentationJSON = await presentation.json();
-  return presentationJSON;
+
+  return presentation.json();
+}
+
+async function createProver(options: {
+  url: string;
+  notaryUrl: string;
+  websocketProxyUrl: string;
+  method?: Method;
+  headers?: {
+    [name: string]: string;
+  };
+  body?: any;
+  maxSentData?: number;
+  maxRecvData?: number;
+  id: string;
+}): Promise<TProver> {
+  const {
+    url,
+    method = 'GET',
+    headers = {},
+    body,
+    maxSentData,
+    maxRecvData,
+    notaryUrl,
+    websocketProxyUrl,
+    id,
+  } = options;
+
+  const hostname = urlify(url)?.hostname || '';
+  const notary = NotaryServer.from(notaryUrl);
+  const prover: TProver = await new Prover({
+    id,
+    serverDns: hostname,
+    maxSentData,
+    maxRecvData,
+  });
+
+  await prover.setup(await notary.sessionUrl(maxSentData, maxRecvData));
+
+  await prover.sendRequest(websocketProxyUrl + `?token=${hostname}`, {
+    url,
+    method,
+    headers,
+    body,
+  });
+
+  return prover;
 }
 
 async function verifyProof(

src/entries/Offscreen/types.ts

@@ -1,4 +1,8 @@
 export enum OffscreenActionTypes {
   notarization_request = 'offscreen/notarization_request',
   notarization_response = 'offscreen/notarization_response',
+  create_prover_request = 'offscreen/create_prover_request',
+  create_prover_response = 'offscreen/create_prover_response',
+  create_presentation_request = 'offscreen/create_presentation_request',
+  create_presentation_response = 'offscreen/create_presentation_response',
 }