From b4b883553a32b0e82a7338558e6a5aa83fd055ba Mon Sep 17 00:00:00 2001
From: Arman Yaraee <arman.yaraee@deversifi.com>
Date: Thu, 3 Aug 2023 17:05:47 +0100
Subject: [PATCH] [PI-139]: Add ledger support

---
 examples/33.generateL1RegistrationPayload.js |  3 +-
 src/lib/keystore/index.js                    |  5 ++-
 src/lib/ledger/index.js                      | 33 +++++++++++++++-----
 src/lib/stark/createRegistrationMessage.js   |  3 +-
 src/lib/stark/signRegistration.js            | 22 ++++---------
 5 files changed, 38 insertions(+), 28 deletions(-)

diff --git a/examples/33.generateL1RegistrationPayload.js b/examples/33.generateL1RegistrationPayload.js
index f8aee944..ce31995a 100755
--- a/examples/33.generateL1RegistrationPayload.js
+++ b/examples/33.generateL1RegistrationPayload.js
@@ -43,11 +43,10 @@ const rhinofiConfig = {
   const { starkKeyHex, ethAddress } = await rhinofi.getUserConfig()
 
   const l1RegistrationSignature = await rhinofi.stark.signRegistration(
-    starkPrivKey,
     ethAddress
   )
 
-  const callData = rhinofi.stark.l1RegistrationCallData(
+  const callData = await rhinofi.stark.l1RegistrationCallData(
     starkKeyHex,
     ethAddress,
     l1RegistrationSignature
diff --git a/src/lib/keystore/index.js b/src/lib/keystore/index.js
index 46bac70b..9b449b24 100644
--- a/src/lib/keystore/index.js
+++ b/src/lib/keystore/index.js
@@ -51,7 +51,10 @@ module.exports = sw => starkPrivateKey => {
   const sign = async tx => {
     const starkKeyPair = await getKeyPair()
 
-    const starkMessage = getMessage(sw)(tx)
+    const starkMessage = typeof tx === 'string'
+      ? tx
+      : getMessage(sw)(tx)
+
     const signature = FP.mapValues(
       x => '0x' + x,
       starkSign({ sw }, starkKeyPair, starkMessage)
diff --git a/src/lib/ledger/index.js b/src/lib/ledger/index.js
index 688c718b..4569942a 100644
--- a/src/lib/ledger/index.js
+++ b/src/lib/ledger/index.js
@@ -22,15 +22,36 @@ const getMessage = sw => tx => {
   return starkTransferTxToMessageHash(sw)(tx)
 }
 
+const withTransport = (dvf) => async (fn) => {
+  const Transport = selectTransport(dvf.isBrowser)
+  const transport = await Transport.create()
+  try {
+    return await fn(transport)
+  } finally {
+    await transport.close()
+  }
+}
+
 const getTxSignature = async (dvf, tx, path) => {
+  // Generic stark message Signing
+  if (typeof tx === 'string') {
+    return withTransport(dvf)(async (transport) => {
+      const eth = new Eth(transport)
+      const { address } = await eth.getAddress(path)
+      const starkPath = dvf.stark.ledger.getPath(address)
+      const paddedMessage = `0x${tx.padEnd(64, '0').substr(-64)}`
+      return eth.starkUnsafeSign(
+        starkPath,
+        paddedMessage
+      )
+    })
+  }
+
   if (tx.type != null) {
     if (!(transferTransactionTypes.includes(tx.type))) {
       throw new DVFError(`Unsupported stark transaction type: ${tx.type}`, { tx })
     }
-    let transport
-    try {
-      const Transport = selectTransport(dvf.isBrowser)
-      transport = await Transport.create()
+    return withTransport(dvf)(async (transport) => {
       const eth = new Eth(transport)
       const { address } = await eth.getAddress(path)
       const starkPath = dvf.stark.ledger.getPath(address)
@@ -81,9 +102,7 @@ const getTxSignature = async (dvf, tx, path) => {
         tx.type === 'ConditionalTransferRequest' ? tx.factRegistryAddress : null,
         tx.type === 'ConditionalTransferRequest' ? tx.fact : null
       )
-    } finally {
-      await transport.close()
-    }
+    })
   } else {
     const { starkSignature } = await createSignedOrder(
       dvf, path, tx, { returnStarkPublicKey: false }
diff --git a/src/lib/stark/createRegistrationMessage.js b/src/lib/stark/createRegistrationMessage.js
index 59ab8a19..601bd3bb 100644
--- a/src/lib/stark/createRegistrationMessage.js
+++ b/src/lib/stark/createRegistrationMessage.js
@@ -12,13 +12,12 @@ const EC_ORDER = utils.BN('36185027886661312136973227830950701055267437517160874
  * ethAddress: string) => utils.BN}
  */
 module.exports = (dvf, starkHex, ethAddress) => {
-  const starkware = dvf.sw || sw
   if (!ethAddress) {
     throw new Error('ethAddress is required')
   }
 
   if (!starkHex) {
-    throw new Error('ethAddress is required')
+    throw new Error('starkKeyHex is required')
   }
 
   /*
diff --git a/src/lib/stark/signRegistration.js b/src/lib/stark/signRegistration.js
index 140ea1d0..163e49fd 100644
--- a/src/lib/stark/signRegistration.js
+++ b/src/lib/stark/signRegistration.js
@@ -2,39 +2,29 @@ const DVFError = require('../dvf/DVFError')
 const sw = require('@rhino.fi/starkware-crypto')
 
 const pad = (rec) => rec.padStart(64, 0)
+const removePrefix = (input) => input.replace(/^0x/, '')
 
 /**
  * @type {(dvf: ReturnType<import('../dvf/bindApi')>,
  * tradingKey: string,
  * ethAddress: string) => Promise<string>}
  */
-module.exports = async (dvf, tradingKey, ethAddress) => {
-  const starkware = dvf.sw || sw
-  if (!tradingKey) {
-    throw new Error('tradingKey is required')
-  }
-
-  /*
-  uint256 msgHash = uint256(
-      keccak256(abi.encodePacked("UserRegistration:", ethKey, starkKey))
-  ) % ECDSA.EC_ORDER;
-  */
-
+module.exports = async (dvf, ethAddress) => {
   const { dvfStarkProvider } = dvf
   const publicKey = await dvfStarkProvider.getPublicKey()
 
   const message = dvf.stark.createRegistrationMessage(`0x${pad(publicKey.x)}`, ethAddress)
 
   try {
-    const { starkKeyPair } = dvf.stark.createKeyPair(tradingKey)
-
-    const { r, s } = starkware.sign(starkKeyPair, message)
+    const { r, s } = await dvfStarkProvider.sign(message)
 
     const components = [
       r.toString(16),
       s.toString(16),
       publicKey.y
-    ].map(pad)
+    ]
+      .map(removePrefix)
+      .map(pad)
 
     const final = `0x${components.join('')}`