From 4bea5b3dc0b974f3112e6a2819390103ffa4efcf Mon Sep 17 00:00:00 2001
From: kent-3 <100624004+kent-3@users.noreply.github.com>
Date: Sat, 31 Aug 2024 11:59:46 -0400
Subject: [PATCH] almost

---
 TNLS-Gateways/secret/src/contract.rs | 131 +++++++++++----------------
 1 file changed, 55 insertions(+), 76 deletions(-)

diff --git a/TNLS-Gateways/secret/src/contract.rs b/TNLS-Gateways/secret/src/contract.rs
index e009e2c..b76a176 100644
--- a/TNLS-Gateways/secret/src/contract.rs
+++ b/TNLS-Gateways/secret/src/contract.rs
@@ -254,8 +254,8 @@ fn pre_execution(deps: DepsMut, _env: Env, msg: PreExecutionMsg) -> StdResult<Re
     let input_hash = sha_256(&[input_values.as_bytes(), new_task.task_id.as_bytes(),&unsafe_payload_bytes].concat());
 
     // create a task information store
-    let task_info = TaskInfo { 
-        payload: msg.payload, // storing the payload
+    let task_info = TaskInfo {
+        payload: msg.payload, // store the payload
         payload_hash: msg.payload_hash,
         payload_signature: msg.payload_signature,
         decrypted_payload_data: input_values.clone(),
@@ -264,8 +264,8 @@ fn pre_execution(deps: DepsMut, _env: Env, msg: PreExecutionMsg) -> StdResult<Re
         user_pubkey: msg.user_pubkey,
         handle: msg.handle.clone(),
         nonce: msg.nonce,
-        unsafe_payload: unsafe_payload, //store the unsafe_payload flag for later checks
-        input_hash, // storing the input_values hashed together with task
+        unsafe_payload, // store the unsafe_payload flag for later checks
+        input_hash,     // store the input_values hashed together with task
         source_network: msg.source_network,
         user_address: payload.user_address.clone(),
         user_key: payload.user_key.clone(),
@@ -281,21 +281,9 @@ fn pre_execution(deps: DepsMut, _env: Env, msg: PreExecutionMsg) -> StdResult<Re
     let mut signing_key_bytes = [0u8; 32];
     signing_key_bytes.copy_from_slice(config.signing_keys.sk.as_slice());
 
-    // used in production to create signature
-    #[cfg(target_arch = "wasm32")]
     let signature = deps.api.secp256k1_sign(&input_hash, &signing_key_bytes)
         .map_err(|err| StdError::generic_err(err.to_string()))?;
 
-    // used only in unit testing to create signatures
-    #[cfg(not(target_arch = "wasm32"))]
-    let signature = {
-        let secp = secp256k1::Secp256k1::signing_only();
-        let sk = secp256k1::SecretKey::from_slice(&signing_key_bytes).unwrap();
-        let message = secp256k1::Message::from_slice(&input_hash)
-            .map_err(|err| StdError::generic_err(err.to_string()))?;
-        secp.sign_ecdsa(&message, &sk).serialize_compact().to_vec()
-    };
-
     // construct the message to send to the destination contract
     let private_contract_msg = SecretMsg::Input {
         message: PrivContractHandleMsg {
@@ -358,43 +346,26 @@ fn post_execution(deps: DepsMut, env: Env, msg: PostExecutionMsg) -> StdResult<R
 
     // load this gateway's signing key
     let private_key = CONFIG.load(deps.storage)?.signing_keys.sk;
-    let mut signing_key_bytes = [0u8; 32];
-    signing_key_bytes.copy_from_slice(private_key.as_slice());
+    println!("{:?}", private_key);
 
-    // used in production to create signature
     // NOTE: api.secp256k1_sign() will perform an additional sha_256 hash operation on the given data
-    #[cfg(target_arch = "wasm32")]
     let packet_signature = {
         deps.api
-            .secp256k1_sign(&packet_hash, &signing_key_bytes)
+            .secp256k1_sign(&packet_hash, &private_key.as_slice())
             .map_err(|err| StdError::generic_err(err.to_string()))?
     };
+    println!("{:?}", packet_signature);
 
-    // used only in unit testing to create signature
-    #[cfg(not(target_arch = "wasm32"))]
-    let packet_signature = {
-        let secp = secp256k1::Secp256k1::signing_only();
-        let sk = secp256k1::SecretKey::from_slice(&signing_key_bytes).unwrap();
-
-        let packet_message = secp256k1::Message::from_slice(&sha_256(&packet_hash))
-            .map_err(|err| StdError::generic_err(err.to_string()))?;
-
-        secp.sign_ecdsa(&packet_message, &sk).serialize_compact()
-    };
-
-    // convert the hashes and signatures into hex byte strings
-    // NOTE: we need to perform the additional sha_256 because that is what the secret network API method does
     // NOTE: The result_signature and packet_signature are both missing the recovery ID (v = 0 or 1), due to a Ethereum bug (v = 27 or 28).
     // we need to either manually check both recovery IDs (v = 27 && v = 28) in the solidity contract (I've leaved this the hard way.)
-
     // or we can find out the two recovery IDs inside of this contract here and keep the solidity contract slim (which is probably the better way when it comes to gas costs):
 
     // Load the original public key from storage
     let public_key_data = CONFIG.load(deps.storage)?.signing_keys.pk;
- 
+
     // Deserialize the uncompressed public key
     let uncompressed_public_key = PublicKey::parse(&public_key_data).map_err(|err| StdError::generic_err(err.to_string()))?;
- 
+
     // Compress the public key
     let compressed_public_key = uncompressed_public_key.serialize_compressed();
 
@@ -432,7 +403,7 @@ fn post_execution(deps: DepsMut, env: Env, msg: PostExecutionMsg) -> StdResult<R
         packet_signature: packet_signature,
         callback_address: callback_address,
         callback_selector: callback_selector,
-        callback_gas_limit: callback_gas_limit
+        callback_gas_limit: callback_gas_limit,
     };
 
     RESULT_MAP.insert(deps.storage, &msg.task, &result_info)?;
@@ -470,8 +441,8 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
 fn query_execution_result(deps: Deps, task: Task) -> StdResult<Binary> {
 
     let task_info = RESULT_MAP
-    .get(deps.storage, &task)
-    .ok_or_else(|| StdError::generic_err("task not found"))?;
+        .get(deps.storage, &task)
+        .ok_or_else(|| StdError::generic_err("task not found"))?;
 
     to_binary(&ResultInfo {
         source_network: task_info.source_network,
@@ -523,7 +494,7 @@ mod tests {
     use super::*;
     use crate::types::*;
     use cosmwasm_std::testing::{mock_dependencies, mock_env, mock_info};
-    use cosmwasm_std::{from_binary, Addr, Binary, Empty};
+    use cosmwasm_std::{from_binary, Addr, Api, Binary, Empty};
 
     use chacha20poly1305::aead::{Aead, NewAead};
     use chacha20poly1305::{ChaCha20Poly1305, Key, Nonce};
@@ -533,13 +504,11 @@ mod tests {
     const SOMEBODY: &str = "somebody";
 
     #[track_caller]
-    fn setup_test_case(deps: DepsMut) -> Result<Response<Empty>, StdError> {
+    fn setup_test_case(deps: DepsMut) -> Result<Response, StdError> {
         // Instantiate a contract with entropy
         let admin = Some(Addr::unchecked(OWNER.to_owned()));
 
-        let init_msg = InstantiateMsg {
-            admin,
-        };
+        let init_msg = InstantiateMsg { admin };
         instantiate(deps, mock_env(), mock_info(OWNER, &[]), init_msg)
     }
 
@@ -562,16 +531,14 @@ mod tests {
     }
 
     #[test]
-    #[ignore = "need new callback fields"]
     fn test_init() {
         let mut deps = mock_dependencies();
 
-        let response = setup_test_case(deps.as_mut()).unwrap();
-        assert_eq!(1, response.messages.len());
+        let response = setup_test_case(deps.as_mut());
+        assert!(response.is_ok());
     }
 
     #[test]
-    #[ignore = "need new callback fields"]
     fn test_query() {
         let mut deps = mock_dependencies();
         let env = mock_env();
@@ -589,7 +556,6 @@ mod tests {
     }
 
     #[test]
-    #[ignore = "need new callback fields"]
     fn test_pre_execution() {
         let mut deps = mock_dependencies();
         let env = mock_env();
@@ -632,8 +598,8 @@ mod tests {
             routing_code_hash: routing_code_hash.clone(),
             user_address: user_address.clone(),
             user_key: user_key.clone(),
-            callback_address: todo!(),
-            callback_selector: todo!(),
+            callback_address: b"public gateway address".into(),
+            callback_selector: b"0xfaef40fe".into(),
             callback_gas_limit: 300_000u32,
         };
         let serialized_payload = to_binary(&payload).unwrap();
@@ -648,9 +614,19 @@ mod tests {
             .unwrap();
 
         // sign the payload
-        let payload_hash = sha_256(serialized_payload.as_slice());
-        let message = Message::from_slice(&payload_hash).unwrap();
-        let payload_signature = secp.sign_ecdsa(&message, &secret_key);
+        let prefix = "\x19Ethereum Signed Message:\n32".as_bytes();
+        let mut hasher = Keccak256::new();
+
+        // NOTE: hmmm shouldn't this be a hash of the non-encrypted payload?
+        hasher.update(encrypted_payload.as_slice());
+        let payload_hash_tmp = hasher.finalize_reset();
+        hasher.update([prefix, &payload_hash_tmp].concat());
+        let payload_hash = hasher.finalize();
+
+        // let message = Message::from_slice(&payload_hash).unwrap();
+        // let payload_signature = secp.sign_ecdsa(&message, &secret_key);
+
+        let payload_signature = deps.api.secp256k1_sign(&payload_hash, secret_key.as_ref()).unwrap();
 
         // mock wrong payload (encrypted with a key that does not match the one inside the payload)
         let wrong_user_address = Addr::unchecked("wrong eth address".to_string());
@@ -662,8 +638,8 @@ mod tests {
             routing_code_hash: routing_code_hash.clone(),
             user_address: wrong_user_address.clone(),
             user_key: wrong_user_key.clone(),
-            callback_address: todo!(),
-            callback_selector: todo!(),
+            callback_address: b"public gateway address".into(),
+            callback_selector: b"0xfaef40fe".into(),
             callback_gas_limit: 300_000u32,
         };
         let wrong_serialized_payload = to_binary(&wrong_payload).unwrap();
@@ -681,7 +657,7 @@ mod tests {
             routing_code_hash: routing_code_hash.clone(),
             payload: Binary(wrong_encrypted_payload.clone()),
             payload_hash: Binary(payload_hash.to_vec()),
-            payload_signature: Binary(payload_signature.serialize_compact().to_vec()),
+            payload_signature: Binary(payload_signature.to_vec()),
             user_address: user_address.clone(),
             user_key: user_key.clone(),
             user_pubkey: user_pubkey.clone(),
@@ -709,7 +685,7 @@ mod tests {
             routing_code_hash: routing_code_hash.clone(),
             payload: Binary(encrypted_payload.clone()),
             payload_hash: Binary(payload_hash.to_vec()),
-            payload_signature: Binary(payload_signature.serialize_compact().to_vec()),
+            payload_signature: Binary(payload_signature.to_vec()),
             user_address: user_address.clone(),
             user_key: user_key.clone(),
             user_pubkey: user_pubkey.clone(),
@@ -731,7 +707,7 @@ mod tests {
             routing_code_hash,
             payload: Binary(encrypted_payload),
             payload_hash: Binary(payload_hash.to_vec()),
-            payload_signature: Binary(payload_signature.serialize_compact().to_vec()),
+            payload_signature: Binary(payload_signature.to_vec()),
             user_address,
             user_key,
             user_pubkey,
@@ -754,7 +730,6 @@ mod tests {
     }
 
     #[test]
-    #[ignore = "need new callback fields"]
     fn test_post_execution() {
         let mut deps = mock_dependencies();
         let env = mock_env();
@@ -792,8 +767,8 @@ mod tests {
             routing_code_hash: routing_code_hash.clone(),
             user_address: user_address.clone(),
             user_key: user_key.clone(),
-            callback_address: todo!(),
-            callback_selector: todo!(),
+            callback_address: b"public gateway address".into(),
+            callback_selector: b"0xfaef40fe".into(),
             callback_gas_limit: 300_000u32,
         };
         let serialized_payload = to_binary(&payload).unwrap();
@@ -805,12 +780,19 @@ mod tests {
         let nonce = Nonce::from_slice(b"unique nonce"); // 12-bytes; unique per message
         let encrypted_payload = cipher
             .encrypt(nonce, serialized_payload.as_slice())
-            .expect("encryption failure!"); // NOTE: handle this error to avoid panics!
+            .expect("encryption failure!");
 
         // sign the payload
-        let payload_hash = sha_256(serialized_payload.as_slice());
-        let message = Message::from_slice(&payload_hash).unwrap();
-        let payload_signature = secp.sign_ecdsa(&message, &secret_key);
+        let prefix = "\x19Ethereum Signed Message:\n32".as_bytes();
+        let mut hasher = Keccak256::new();
+
+        // NOTE: shouldn't this be a hash of the non-encrypted payload?
+        hasher.update(encrypted_payload.as_slice());
+        let payload_hash_tmp = hasher.finalize_reset();
+        hasher.update([prefix, &payload_hash_tmp].concat());
+        let payload_hash = hasher.finalize();
+
+        let payload_signature = deps.api.secp256k1_sign(&payload_hash, secret_key.as_ref()).unwrap();
 
         // execute input handle
         let pre_execution_msg = PreExecutionMsg {
@@ -820,7 +802,7 @@ mod tests {
             routing_code_hash,
             payload: Binary(encrypted_payload),
             payload_hash: Binary(payload_hash.to_vec()),
-            payload_signature: Binary(payload_signature.serialize_compact().to_vec()),
+            payload_signature: Binary(payload_signature.to_vec()),
             user_address,
             user_key,
             user_pubkey: user_pubkey.clone(),
@@ -835,7 +817,7 @@ mod tests {
 
         // test incorrect input_hash
         let wrong_post_execution_msg = PostExecutionMsg {
-            result: "{\"answer\": 42}".to_string(),
+            result: base64::encode("{\"answer\": 42}".to_string()),
             task: Task { network: "ethereum".to_string(), task_id: "1".to_string() },
             input_hash: Binary(sha_256("wrong data".as_bytes()).to_vec()),
         };
@@ -843,17 +825,14 @@ mod tests {
             outputs: wrong_post_execution_msg,
         };
         let err = execute(deps.as_mut(), env.clone(), info.clone(), handle_msg).unwrap_err();
-        assert_eq!(
-            err,
-            StdError::generic_err("input hash does not match task id")
-        );
+        assert_eq!(err, StdError::generic_err("input hash does not match task"));
 
         // test output handle
         let post_execution_msg = PostExecutionMsg {
-            result: "{\"answer\": 42}".to_string(),
+            result: base64::encode("{\"answer\": 42}".to_string()),
             task: Task { network: "ethereum".to_string(), task_id: "1".to_string() },
             input_hash: Binary(
-                sha_256(&[data.as_bytes(), 1u64.to_le_bytes().as_ref()].concat()).to_vec(),
+                sha_256(&[data.as_bytes(), "1".to_string().as_bytes(), &[0u8]].concat()).to_vec(),
             ),
         };