Merge pull request #157 from athenavm/feature/encoding-verify

Fix and test wallet `verify` method

examples/wallet/program/src/main.rs

@@ -56,32 +56,20 @@ impl WalletProgram for Wallet {
 }
 
 impl VerifiableTemplate for Wallet {
-  fn verify(&self, tx: &[u8], signature: &[u8; 64]) -> bool {
+  fn verify(&self, tx: alloc::vec::Vec<u8>, signature: [u8; 64]) -> bool {
     // Check that the transaction is signed by the owner
     let public_key = VerifyingKey::from_bytes(&self.owner.0).unwrap();
-    let signature = Signature::from_bytes(signature);
+    let signature = Signature::from_bytes(&signature);
     public_key.verify(&tx, &signature).is_ok()
   }
 }
 
-#[cfg(all(
-  any(target_arch = "riscv32", target_arch = "riscv64"),
-  target_feature = "e",
-  test
-))]
-mod test {
-  use super::*;
-
-  #[test]
-  fn test_wallet() {
-    let owner = Pubkey([0u8; 32]);
-    let wallet = Wallet::new(owner);
-
-    Wallet::spawn(owner);
-    let send_arguments = SpendArguments {
-      recipient: [0u8; 24],
-      amount: 0,
-    };
-    wallet.spend(send_arguments);
-  }
+#[link_section = ".text.athexp.verify"]
+#[no_mangle]
+pub unsafe extern "C" fn athexp_verify() {
+  athena_vm::program::Method::call_method(Wallet::verify, &mut athena_vm::io::Io::default());
+  syscall_halt(0);
 }
+#[used]
+#[link_section = ".init_array"]
+static DUMMY_VERIFY: unsafe extern "C" fn() = athexp_verify;

examples/wallet/script/Cargo.toml

@@ -17,5 +17,9 @@ hex = "0.4.3"
 parity-scale-codec = "3.6.12"
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 
+[dev-dependencies]
+ed25519-dalek = { version = "2.1.1", features = ["rand_core"] }
+rand = "0.8.5"
+
 [build-dependencies]
 athena-helper = { path = "../../../helper" }

examples/wallet/script/src/bin/execute.rs

@@ -125,13 +125,20 @@ mod tests {
   };
   use athena_sdk::{AthenaStdin, ExecutionClient};
   use athena_vm_sdk::Pubkey;
+  use ed25519_dalek::ed25519::signature::Signer;
+  use ed25519_dalek::SigningKey;
+  use rand::rngs::OsRng;
 
-  #[test]
-  fn deploy_template() {
-    tracing_subscriber::fmt()
+  fn setup_logger() {
+    let _ = tracing_subscriber::fmt()
       .with_test_writer()
       .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
-      .init();
+      .try_init();
+  }
+
+  #[test]
+  fn deploy_template() {
+    setup_logger();
 
     let mut host = MockHost::new_with_context(
       HostStaticContext::new(ADDRESS_ALICE, 0, ADDRESS_ALICE),
@@ -162,4 +169,61 @@ mod tests {
     let template = host.template(&address);
     assert_eq!(template, Some(&code));
   }
+
+  #[test]
+  fn verifying_tx() {
+    setup_logger();
+
+    let mut host = MockHost::new_with_context(
+      HostStaticContext::new(ADDRESS_ALICE, 0, ADDRESS_ALICE),
+      HostDynamicContext::new([0u8; 24], ADDRESS_ALICE),
+    );
+
+    let signing_key = SigningKey::generate(&mut OsRng);
+    let owner = Pubkey(signing_key.verifying_key().to_bytes());
+    let address = super::spawn(&mut host, &owner).unwrap();
+    let wallet_state = host.get_program(&address).unwrap().clone();
+
+    let tx = b"some really bad tx";
+
+    // First try with invalid signature
+    {
+      let mut stdin = AthenaStdin::new();
+      stdin.write_vec(wallet_state.clone());
+      stdin.write_vec(tx.as_slice().encode());
+      stdin.write_vec([0; 64].encode());
+
+      let result = ExecutionClient::new().execute_function(
+        super::ELF,
+        &MethodSelector::from("athexp_verify"),
+        stdin.clone(),
+        Some(&mut host),
+        Some(25000000),
+        None,
+      );
+      let (mut result, _) = result.unwrap();
+      let valid = result.read::<bool>();
+      assert!(!valid);
+    }
+
+    // Now with valid signature
+    let signature = signing_key.sign(tx);
+
+    let mut stdin = AthenaStdin::new();
+    stdin.write_vec(wallet_state.clone());
+    stdin.write_vec(tx.as_slice().encode());
+    stdin.write_vec(signature.to_bytes().encode());
+
+    let result = ExecutionClient::new().execute_function(
+      super::ELF,
+      &MethodSelector::from("athexp_verify"),
+      stdin.clone(),
+      Some(&mut host),
+      Some(25000000),
+      None,
+    );
+    let (mut result, _) = result.unwrap();
+    let valid = result.read::<bool>();
+    assert!(valid);
+  }
 }

ffi/athcon/bindings/go/athcon.go

@@ -82,6 +82,7 @@ func LoadLibrary(path string) (*Library, error) {
 	purego.RegisterLibFunc(&lib.create, libHandle, "athcon_create_"+vmName)
 	purego.RegisterLibFunc(&lib.encodeTxSpawn, libHandle, "athcon_encode_tx_spawn")
 	purego.RegisterLibFunc(&lib.encodeTxSpend, libHandle, "athcon_encode_tx_spend")
+
 	purego.RegisterLibFunc(&lib.freeBytes, libHandle, "athcon_free_bytes")
 	return lib, nil
 }

ffi/athcon/bindings/go/go.mod

@@ -6,11 +6,12 @@ require (
 	github.com/ChainSafe/gossamer v0.9.0
 	github.com/ebitengine/purego v0.7.1
 	github.com/stretchr/testify v1.9.0
-	golang.org/x/crypto v0.28.0
+	github.com/zeebo/blake3 v0.2.4
 )
 
 require (
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
+	github.com/klauspost/cpuid/v2 v2.2.6 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	golang.org/x/sys v0.26.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect

ffi/athcon/bindings/go/go.sum

@@ -6,6 +6,8 @@ github.com/ebitengine/purego v0.7.1 h1:6/55d26lG3o9VCZX8lping+bZcmShseiqlh2bnUDi
 github.com/ebitengine/purego v0.7.1/go.mod h1:ah1In8AOtksoNK6yk5z1HTJeUkC1Ez4Wk2idgGslMwQ=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/klauspost/cpuid/v2 v2.2.6 h1:ndNyv040zDGIDh8thGkXYjnFtiN02M1PVVF+JE/48xc=
+github.com/klauspost/cpuid/v2 v2.2.6/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws=
 github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
@@ -19,8 +21,13 @@ github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjR
 github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-golang.org/x/crypto v0.28.0 h1:GBDwsMXVQi34v5CCYUm2jkJvu4cbtru2U4TN2PSyQnw=
-golang.org/x/crypto v0.28.0/go.mod h1:rmgy+3RHxRZMyY0jjAJShp2zgEdOqj2AO7U0pYmeQ7U=
+github.com/zeebo/assert v1.1.0 h1:hU1L1vLTHsnO8x8c9KAR5GmM5QscxHg5RNU5z5qbUWY=
+github.com/zeebo/assert v1.1.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0=
+github.com/zeebo/blake3 v0.2.4 h1:KYQPkhpRtcqh0ssGYcKLG1JYvddkEA8QwCM/yBqhaZI=
+github.com/zeebo/blake3 v0.2.4/go.mod h1:7eeQ6d2iXWRGF6npfaxl2CU+xy2Fjo2gxeyZGCRUjcE=
+github.com/zeebo/pcg v1.0.1 h1:lyqfGeWiv4ahac6ttHs+I5hwtH/+1mrhlCtVNQM2kHo=
+github.com/zeebo/pcg v1.0.1/go.mod h1:09F0S9iiKrwn9rlI5yjLkmrug154/YRW6KnnXVDM/l4=
+golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.26.0 h1:KHjCJyddX0LoSTb3J+vWpupP9p0oznkqVk/IfjymZbo=
 golang.org/x/sys v0.26.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

ffi/athcon/bindings/go/host_test.go

@@ -2,6 +2,7 @@ package athcon
 
 import (
 	"bytes"
+	"crypto/ed25519"
 	"crypto/rand"
 	_ "embed"
 	"encoding/binary"
@@ -201,3 +202,67 @@ func TestSpend(t *testing.T) {
 	require.Equal(t, host.balances[recipient], uint64(100))
 	require.Equal(t, host.balances[principal], uint64(900))
 }
+
+func TestVerify(t *testing.T) {
+	vm, _ := Load(modulePath)
+	defer vm.Destroy()
+
+	host := newHost()
+	// Step 1: Spawn wallet
+	principal := randomAddress()
+	pubkey, privkey, err := ed25519.GenerateKey(nil)
+	var walletAddress Address
+	{
+		pubkey := Bytes32(pubkey)
+		executionPayload := EncodedExecutionPayload(nil, vm.Lib.EncodeTxSpawn(pubkey))
+
+		result, err := vm.Execute(host, Frontier, Call, 1, 10000000, principal, principal, executionPayload, 0, WALLET_TEST)
+		require.NoError(t, err)
+		require.Len(t, result.Output, 24)
+
+		walletAddress = Address(result.Output)
+	}
+	// Step 2: Try verify TX with invalid signature
+	tx := make([]byte, 100)
+	_, err = rand.Read(tx)
+	require.NoError(t, err)
+	txEncoded, err := scale.Marshal(tx)
+	require.NoError(t, err)
+
+	var signature [64]byte
+	signatureEncoded, err := scale.Marshal(signature)
+	require.NoError(t, err)
+
+	selector, err := FromString("athexp_verify")
+	require.NoError(t, err)
+	payload := Payload{
+		Selector: &selector,
+		Input:    append(txEncoded, signatureEncoded...),
+	}
+	payloadEncoded, err := scale.Marshal(payload)
+	require.NoError(t, err)
+
+	executionPayload := EncodedExecutionPayload(host.programs[walletAddress], payloadEncoded)
+
+	result, err := vm.Execute(host, Frontier, Call, 1, 1000000000, principal, principal, executionPayload, 0, WALLET_TEST)
+	require.NoError(t, err)
+	require.Zero(t, result.Output[0])
+
+	// Step 3: Try verify TX with valid signature
+	copy(signature[:], ed25519.Sign(privkey, tx))
+	signatureEncoded, err = scale.Marshal(signature)
+	require.NoError(t, err)
+
+	payload = Payload{
+		Selector: &selector,
+		Input:    append(txEncoded, signatureEncoded...),
+	}
+	payloadEncoded, err = scale.Marshal(payload)
+	require.NoError(t, err)
+
+	executionPayload = EncodedExecutionPayload(host.programs[walletAddress], payloadEncoded)
+
+	result, err = vm.Execute(host, Frontier, Call, 1, 1000000000, principal, principal, executionPayload, 0, WALLET_TEST)
+	require.NoError(t, err)
+	require.Equal(t, uint8(1), result.Output[0])
+}

ffi/athcon/bindings/go/types.go

@@ -5,7 +5,7 @@ import (
 	"fmt"
 
 	"github.com/ChainSafe/gossamer/pkg/scale"
-	"golang.org/x/crypto/blake2b"
+	"github.com/zeebo/blake3"
 )
 
 // Address represents the 24 bytes address of an Athena account.
@@ -22,14 +22,11 @@ type MethodSelector [MethodSelectorLength]byte
 
 // FromString converts a string to a MethodSelector, similar to the Rust From<&str> implementation.
 func FromString(value string) (MethodSelector, error) {
-	hash, err := blake2b.New256(nil)
-	if err != nil {
-		return MethodSelector{}, err
-	}
-	hash.Write([]byte(value))
-	hashBytes := hash.Sum(nil)
 	var selector MethodSelector
-	copy(selector[:], hashBytes[:MethodSelectorLength])
+	hasher := blake3.New()
+	hasher.Write([]byte(value))
+	hasher.Digest().Read(selector[:])
+
 	return selector, nil
 }
 

ffi/athcon/bindings/rust/athcon-vm/src/encode_tx.rs

@@ -72,6 +72,7 @@ mod tests {
     assert_eq!(tx.selector, Some(MethodSelector::from("athexp_spawn")));
     assert_eq!(Pubkey::decode(&mut tx.input.as_slice()).unwrap(), pubkey);
   }
+
   #[test]
   fn encoding_spend_tx() {
     let address = [0x1C; 24];