Fix tee attestation (#102)

app/src/screens/Onboarding/ConfirmBelongingScreen.tsx

@@ -14,8 +14,7 @@ import { notificationSuccess } from '../../utils/haptic';
 import { styles } from '../ProveFlow/ProofRequestStatusScreen';
 
 const ConfirmBelongingScreen: React.FC = () => {
-  const onOkPress = useHapticNavigation('LoadingScreen');
-
+  const onOkPress = useHapticNavigation('LoadingScreen', 'default');
   useEffect(() => {
     notificationSuccess();
   }, []);

app/src/screens/Onboarding/LoadingScreen.tsx

@@ -1,4 +1,4 @@
-import React, { useEffect, useState } from 'react';
+import React, { useEffect, useRef, useState } from 'react';
 import { StyleSheet } from 'react-native';
 
 import LottieView from 'lottie-react-native';
@@ -9,16 +9,25 @@ import { initPassportDataParsing } from '../../../../common/src/utils/passports/
 // Import animations
 import failAnimation from '../../assets/animations/loading/fail.json';
 import miscAnimation from '../../assets/animations/loading/misc.json';
+import successAnimation from '../../assets/animations/loading/success.json';
 import useHapticNavigation from '../../hooks/useHapticNavigation';
-import { usePassport } from '../../stores/passportDataProvider';
 import { ProofStatusEnum, useProofInfo } from '../../stores/proofProvider';
+import useUserStore from '../../stores/userStore';
 import { registerPassport } from '../../utils/proving/payload';
 
 const LoadingScreen: React.FC = () => {
-  const { setData } = usePassport();
-  const goToSuccessScreen = useHapticNavigation('AccountVerifiedSuccess', {
-    action: 'default',
-  });
+  const goToSuccessScreen = useHapticNavigation('AccountVerifiedSuccess');
+  const goToErrorScreen = useHapticNavigation('ConfirmBelongingScreen');
+  const goToSuccessScreenWithDelay = () => {
+    setTimeout(() => {
+      goToSuccessScreen();
+    }, 3000);
+  };
+  const goToErrorScreenWithDelay = () => {
+    setTimeout(() => {
+      goToErrorScreen();
+    }, 3000);
+  };
   const [animationSource, setAnimationSource] = useState<any>(miscAnimation);
   const { status, setStatus } = useProofInfo();
 
@@ -30,36 +39,45 @@ const LoadingScreen: React.FC = () => {
   useEffect(() => {
     // Change animation based on the global proof status.
     if (status === ProofStatusEnum.SUCCESS) {
-      goToSuccessScreen();
+      setAnimationSource(successAnimation);
+      goToSuccessScreenWithDelay();
     } else if (
       status === ProofStatusEnum.FAILURE ||
       status === ProofStatusEnum.ERROR
     ) {
       setAnimationSource(failAnimation);
+      goToErrorScreenWithDelay();
     }
   }, [status]);
 
+  // Use a ref to make sure processPayload is only executed once during the component's lifecycle.
+  const processPayloadCalled = useRef(false);
+
   useEffect(() => {
-    const processPayload = async () => {
-      try {
-        // Generate passport data and update the store.
-        const passportData = genMockPassportData(
-          'sha1',
-          'sha256',
-          'rsa_sha256_65537_2048',
-          'FRA',
-          '000101',
-          '300101',
-        );
-        const passportDataInit = initPassportDataParsing(passportData);
-        setData(passportDataInit);
-        await registerPassport(passportDataInit);
-      } catch (error) {
-        console.error('Error processing payload:', error);
-        setStatus(ProofStatusEnum.ERROR);
-      }
-    };
-    processPayload();
+    if (!processPayloadCalled.current) {
+      processPayloadCalled.current = true;
+      const processPayload = async () => {
+        try {
+          // Generate passport data and update the store.
+          const passportData = genMockPassportData(
+            'sha1',
+            'sha256',
+            'rsa_sha256_65537_2048',
+            'FRA',
+            '000101',
+            '300101',
+          );
+          const passportDataInit = initPassportDataParsing(passportData);
+          await useUserStore.getState().registerPassportData(passportDataInit);
+          // This will trigger sendPayload(), which updates global status via your tee.ts code.
+          await registerPassport(passportDataInit);
+        } catch (error) {
+          console.error('Error processing payload:', error);
+          setStatus(ProofStatusEnum.ERROR);
+        }
+      };
+      processPayload();
+    }
   }, []);
 
   return (

app/src/stores/appProvider.tsx

@@ -2,7 +2,7 @@ import React, { createContext, useContext, useEffect, useRef } from 'react';
 
 import io, { Socket } from 'socket.io-client';
 
-import { WS_DB_RELAYER_NEW } from '../../../common/src/constants/constants';
+import { WS_DB_RELAYER } from '../../../common/src/constants/constants';
 import { SelfApp } from '../../../common/src/utils/appType';
 
 interface IAppContext {
@@ -45,9 +45,9 @@ export const AppProvider: React.FC<{ children: React.ReactNode }> = ({
       }
 
       // Ensure the URL uses the proper WebSocket scheme.
-      const connectionUrl = WS_DB_RELAYER_NEW.startsWith('https')
-        ? WS_DB_RELAYER_NEW.replace(/^https/, 'wss')
-        : WS_DB_RELAYER_NEW;
+      const connectionUrl = WS_DB_RELAYER.startsWith('https')
+        ? WS_DB_RELAYER.replace(/^https/, 'wss')
+        : WS_DB_RELAYER;
       const socketUrl = `${connectionUrl}/websocket`;
 
       // Create a new socket connection using the updated URL.

app/src/stores/proofProvider.tsx

@@ -9,7 +9,7 @@ import React, {
 
 import io, { Socket } from 'socket.io-client';
 
-import { WS_DB_RELAYER_NEW } from '../../../common/src/constants/constants';
+import { WS_DB_RELAYER } from '../../../common/src/constants/constants';
 import { SelfApp } from '../../../common/src/utils/appType';
 import { setupUniversalLinkListener } from '../utils/qrCodeNew';
 
@@ -76,6 +76,9 @@ export function ProofProvider({ children }: PropsWithChildren) {
       appName: '',
       logoBase64: '',
       scope: '',
+      endpointType: 'https',
+      endpoint: '',
+      header: '',
       sessionId: '',
       userId: '',
       userIdType: 'uuid',
@@ -141,10 +144,10 @@ function useWebsocket(
     if (!selectedApp.sessionId) {
       return;
     }
-    console.log('creating ws', WS_DB_RELAYER_NEW, selectedApp.sessionId);
+    console.log('creating ws', WS_DB_RELAYER, selectedApp.sessionId);
 
     try {
-      newSocket = io(WS_DB_RELAYER_NEW + '/websocket', {
+      newSocket = io(WS_DB_RELAYER + '/websocket', {
         path: '/',
         transports: ['websocket'],
         query: { sessionId: selectedApp.sessionId, clientType: 'mobile' },

app/src/utils/proving/attest.ts

@@ -5,6 +5,7 @@ import * as asn1 from 'asn1.js';
 import * as asn1js from 'asn1js';
 import { Buffer } from 'buffer';
 import elliptic from 'elliptic';
+import { sha384 } from 'js-sha512';
 import { Certificate } from 'pkijs';
 
 import { IMAGE_HASH } from '../../../../common/src/constants/constants';
@@ -70,21 +71,18 @@ export const numberInRange = (
  * @param certChainStr An array of certificates in PEM format, ordered from leaf to root.
  * @return True if the certificate chain is valid, false otherwise.
  */
-export const verifyCertChain = (
+export const verifyCertChain = async (
   rootPem: string,
   certChainStr: string[],
-): boolean => {
+): Promise<boolean> => {
   try {
     // Parse all certificates
-    const rootCert = new X509Certificate(rootPem);
     const certChain = certChainStr.map(cert => new X509Certificate(cert));
 
     // Verify the chain from leaf to root
-    for (let i = 0; i < certChain.length; i++) {
+    // certChain[0] is the root, we use the hardcoded rootPem
+    for (let i = 1; i < certChain.length; i++) {
       const currentCert = certChain[i];
-      const issuerCert =
-        i === certChain.length - 1 ? rootCert : certChain[i + 1];
-
       // Verify certificate validity period
       const now = new Date();
       if (now < currentCert.notBefore || now > currentCert.notAfter) {
@@ -94,7 +92,10 @@ export const verifyCertChain = (
 
       // Verify signature
       try {
-        const isValid = currentCert.verify(issuerCert);
+        const isValid = verifyCertificateSignature(
+          certChainStr[i],
+          i === 1 ? rootPem : certChainStr[i - 1],
+        );
         if (!isValid) {
           console.error(`Certificate at index ${i} has invalid signature`);
           return false;
@@ -199,7 +200,7 @@ export const verifyAttestation = async (attestation: Array<number>) => {
   }
   console.log('TEE image hash verified');
 
-  if (!verifyCertChain(AWS_ROOT_PEM, [...certChain, cert])) {
+  if (!(await verifyCertChain(AWS_ROOT_PEM, [...certChain, cert]))) {
     throw new Error('Invalid certificate chain');
   }
 
@@ -340,9 +341,51 @@ export function getCertificateFromPem(pemContent: string): Certificate {
   }
 
   const asn1Data = asn1js.fromBER(arrayBuffer);
-  if (asn1.offset === -1) {
+  if (asn1Data.offset === -1) {
     throw new Error(`ASN.1 parsing error: ${asn1Data.result.error}`);
   }
 
   return new Certificate({ schema: asn1Data.result });
 }
+
+function verifyCertificateSignature(child: string, parent: string): boolean {
+  const certBuffer_csca = Buffer.from(
+    parent.replace(/(-----(BEGIN|END) CERTIFICATE-----|\n)/g, ''),
+    'base64',
+  );
+  const asn1Data_csca = asn1js.fromBER(certBuffer_csca);
+  const cert_csca = new Certificate({ schema: asn1Data_csca.result });
+  const publicKeyInfo_csca = cert_csca.subjectPublicKeyInfo;
+  const publicKeyBuffer_csca =
+    publicKeyInfo_csca.subjectPublicKey.valueBlock.valueHexView;
+  const curve = 'p384';
+  const ec_csca = new elliptic.ec(curve);
+  const key_csca = ec_csca.keyFromPublic(publicKeyBuffer_csca);
+
+  const tbsHash = getTBSHash(child);
+
+  const certBuffer_dsc = Buffer.from(
+    child.replace(/(-----(BEGIN|END) CERTIFICATE-----|\n)/g, ''),
+    'base64',
+  );
+  const asn1Data_dsc = asn1js.fromBER(certBuffer_dsc);
+  const cert_dsc = new Certificate({ schema: asn1Data_dsc.result });
+  const signatureValue = cert_dsc.signatureValue.valueBlock.valueHexView;
+  const signature_crypto = Buffer.from(signatureValue).toString('hex');
+  return key_csca.verify(tbsHash, signature_crypto);
+}
+
+export function getTBSHash(pem: string): string {
+  const certBuffer = Buffer.from(
+    pem.replace(/(-----(BEGIN|END) CERTIFICATE-----|\n)/g, ''),
+    'base64',
+  );
+  const asn1Data_cert = asn1js.fromBER(certBuffer);
+  const cert = new Certificate({ schema: asn1Data_cert.result });
+  const tbsAsn1 = cert.encodeTBS();
+  const tbsDer = tbsAsn1.toBER(false);
+  const tbsBytes = Buffer.from(tbsDer);
+  const tbsBytesArray = Array.from(tbsBytes);
+  const msgHash = sha384(tbsBytesArray);
+  return msgHash as string;
+}

app/src/utils/proving/payload.ts

@@ -27,8 +27,11 @@ import { sendPayload } from './tee';
 
 const mock_secret = '0'; //TODO: retrieve the secret from keychain
 
-function generateTeeInputsRegister(secret: string, passportData: PassportData) {
-  const inputs = generateCircuitInputsRegister(secret, passportData);
+async function generateTeeInputsRegister(
+  secret: string,
+  passportData: PassportData,
+) {
+  const inputs = await generateCircuitInputsRegister(secret, passportData);
   const circuitName = getCircuitNameFromPassportData(passportData, 'register');
   if (circuitName == null) {
     throw new Error('Circuit name is null');
@@ -83,7 +86,15 @@ export async function sendRegisterPayload(passportData: PassportData) {
     mock_secret,
     passportData,
   );
-  await sendPayload(inputs, circuitName, WS_RPC_URL_REGISTER);
+  console.log('WS_RPC_URL_REGISTER', WS_RPC_URL_REGISTER);
+  await sendPayload(
+    inputs,
+    'register',
+    circuitName,
+    'https',
+    'https://self.xyz',
+    WS_RPC_URL_REGISTER,
+  );
 }
 
 function generateTeeInputsDsc(passportData: PassportData) {
@@ -106,7 +117,14 @@ export async function sendDscPayload(passportData: PassportData): Promise<any> {
   }
   const { inputs, circuitName } = generateTeeInputsDsc(passportData);
   console.log('circuitName', circuitName);
-  const result = await sendPayload(inputs, circuitName, WS_RPC_URL_DSC);
+  const result = await sendPayload(
+    inputs,
+    'dsc',
+    circuitName,
+    'https',
+    'https://self.xyz',
+    WS_RPC_URL_DSC,
+  );
   return result;
 }
 
@@ -167,7 +185,14 @@ export async function sendVcAndDisclosePayload(
     return;
   }
   const { inputs, circuitName } = generateTeeInputsVCAndDisclose(passportData);
-  await sendPayload(inputs, circuitName, WS_RPC_URL_VC_AND_DISCLOSE);
+  await sendPayload(
+    inputs,
+    'vc_and_disclose',
+    circuitName,
+    'https',
+    'https://self.xyz',
+    WS_RPC_URL_VC_AND_DISCLOSE,
+  );
 }
 
 /*** Logic Flow ****/

app/src/utils/proving/tee.ts

@@ -3,12 +3,16 @@ import forge from 'node-forge';
 import io, { Socket } from 'socket.io-client';
 import { v4 } from 'uuid';
 
-import { WS_DB_RELAYER_OLD } from '../../../../common/src/constants/constants';
+import {
+  CIRCUIT_TYPES,
+  WS_DB_RELAYER,
+} from '../../../../common/src/constants/constants';
+import { EndpointType } from '../../../../common/src/utils/appType';
 import {
   ProofStatusEnum,
   updateGlobalProofStatus,
 } from '../../stores/proofProvider';
-import { verifyAttestation } from './attest';
+import { getPublicKey, verifyAttestation } from './attest';
 
 const { ec: EC } = elliptic;
 
@@ -50,13 +54,17 @@ const pubkey =
  */
 export async function sendPayload(
   inputs: any,
+  circuit: (typeof CIRCUIT_TYPES)[number],
   circuitName: string,
+  endpointType: EndpointType,
+  endpoint: string,
   wsRpcUrl: string,
   timeoutMs = 1200000,
 ): Promise<void> {
   return new Promise(resolve => {
     let finalized = false;
     function finalize(status: ProofStatusEnum) {
+      console.log('Finalizing with status:', status);
       if (!finalized) {
         finalized = true;
         updateGlobalProofStatus(status);
@@ -92,20 +100,59 @@ export async function sendPayload(
         const result = JSON.parse(event.data);
         // If attestation is present, process it.
         if (result.result?.attestation !== undefined) {
-          // const serverPubkey = getPublicKey(result.result.attestation);
+          const serverPubkey = getPublicKey(result.result.attestation);
           const verified = await verifyAttestation(result.result.attestation);
-          console.log('AWS Root Certificate verified:', verified);
-          if (verified) {
-            finalize(ProofStatusEnum.SUCCESS);
-          } else {
+          if (!verified) {
             finalize(ProofStatusEnum.FAILURE);
+            throw new Error('Attestation verification failed');
           }
+          const key2 = ec.keyFromPublic(serverPubkey as string, 'hex');
+          const sharedKey = key1.derive(key2.getPublic());
+          const forgeKey = forge.util.createBuffer(
+            Buffer.from(
+              sharedKey.toString('hex').padStart(64, '0'),
+              'hex',
+            ).toString('binary'),
+          );
+          const payload = getPayload(
+            inputs,
+            circuit,
+            circuitName,
+            endpointType,
+            endpoint,
+          );
+          const encryptionData = encryptAES256GCM(
+            JSON.stringify(payload),
+            forgeKey,
+          );
+          const submitBody = {
+            jsonrpc: '2.0',
+            method: 'openpassport_submit_request',
+            id: 1,
+            params: {
+              uuid: result.result.uuid,
+              ...encryptionData,
+            },
+          };
+          console.log('Sending submit body');
+          const truncatedBody = {
+            ...submitBody,
+            params: {
+              uuid: submitBody.params.uuid,
+              nonce: submitBody.params.nonce.slice(0, 3) + '...',
+              cipher_text: submitBody.params.cipher_text.slice(0, 3) + '...',
+              auth_tag: submitBody.params.auth_tag.slice(0, 3) + '...',
+            },
+          };
+          console.log('Truncated submit body:', truncatedBody);
+          ws.send(JSON.stringify(submitBody));
         } else {
           // Otherwise, assume it's a UUID. Set up SocketIO to get further progress.
           const receivedUuid = result.result;
           console.log('Received UUID:', receivedUuid);
+          console.log(result);
           if (!socket) {
-            socket = io(WS_DB_RELAYER_OLD, {
+            socket = io(WS_DB_RELAYER, {
               path: '/',
               transports: ['websocket'],
             });
@@ -113,17 +160,18 @@ export async function sendPayload(
               console.log('SocketIO: Connection opened');
               socket?.emit('subscribe', receivedUuid);
             });
-            socket.on('message', message => {
+            socket.on('status', message => {
               const data =
                 typeof message === 'string' ? JSON.parse(message) : message;
               console.log('SocketIO message:', data);
               // When the proof has generated, disconnect and close the WebSocket.
-              if (data.new_status === 2) {
+              if (data.status === 2) {
                 console.log('Proof generation completed');
                 socket?.disconnect();
                 if (ws.readyState === WebSocket.OPEN) {
                   ws.close();
                 }
+                finalize(ProofStatusEnum.SUCCESS);
               }
             });
             socket.on('disconnect', reason => {
@@ -168,3 +216,55 @@ export async function sendPayload(
 }
 
 export { encryptAES256GCM };
+
+/***
+ *  types
+ * ***/
+export type TEEPayloadDisclose = {
+  type: 'disclose';
+  endpointType: string;
+  endpoint: string;
+  circuit: {
+    name: string;
+    inputs: string;
+  };
+};
+
+export type TEEPayload = {
+  type: 'register' | 'dsc';
+  onchain: true;
+  circuit: {
+    name: string;
+    inputs: string;
+  };
+};
+export function getPayload(
+  inputs: any,
+  circuit: string,
+  circuitName: string,
+  endpointType: string,
+  endpoint: string,
+) {
+  if (circuit === 'vc_and_disclose') {
+    const payload: TEEPayloadDisclose = {
+      type: 'disclose',
+      endpointType: endpointType,
+      endpoint: endpoint,
+      circuit: {
+        name: circuitName,
+        inputs: JSON.stringify(inputs),
+      },
+    };
+    return payload;
+  } else if (circuit === 'register' || circuit === 'dsc') {
+    const payload: TEEPayload = {
+      type: circuit as 'register' | 'dsc',
+      onchain: true,
+      circuit: {
+        name: circuitName,
+        inputs: JSON.stringify(inputs),
+      },
+    };
+    return payload;
+  }
+}