various fixes and test script

contracts/voting_snapshot/src/lib.rs

@@ -163,6 +163,7 @@ impl Contract {
     /// - User should be eligible
     /// - User should not be registered before
     /// - User should pay for storage
+    #[payable]
     pub fn register_as_nominee(&mut self) {
         let storage = env::storage_usage();
 
@@ -174,10 +175,7 @@ impl Contract {
 
         self.nominees.insert(user);
 
-        require!(
-            finalize_storage_check(storage, SNAPSHOT_RECORD_COST),
-            STORAGE_LIMIT_EXCEEDED
-        );
+        require!(finalize_storage_check(storage, 0), STORAGE_LIMIT_EXCEEDED);
     }
 
     /// *Transaction*: Any user can challenge the snapshot. User can deposit NEAR several times

relayer/src/cryptography/index.ts

@@ -39,29 +39,10 @@ export const verifySignature = (data: string, public_key: string, signature: str
     }
 }
 
-export const createSignature = (data: string, privateStr: string): string | undefined => {
+export const createSignature = (data: string, keyPair: KeyPair): string | undefined => {
     const message = base_decode(data);
-    const [type, private_key] = privateStr.split(':');
 
-    let signature: Uint8Array;
-    try {
-        if (type === "secp256k1") {
-            // secp256k1
-            const privKey = base_decode(private_key);
-            signature = ecdsaSign(message, privKey).signature;
-        } else if (type === "ed25519") {
-            // ed25519
-            const keyPair = KeyPair.fromString(privateStr);
-            signature = keyPair.sign(message).signature;
-        }
-        else {
-            return undefined;
-        }
-    } catch {
-        return undefined;
-    }
-
-    return base_encode(signature);
+    return base_encode(keyPair.sign(message).signature);
 }
 
 export const decrypt = async (vote: EncryptedVotingPackage, privateKey: Uint8Array): Promise<Result<VotingPackage>> => {

relayer/src/utils/vote.ts

@@ -1,6 +1,7 @@
 import { base_encode } from "near-api-js/lib/utils/serialize";
 import { createSignature, encrypt } from "../cryptography";
 import { AccountId, EncryptedVotingPackage, EncryptedVotingPackageWithProof, VotingPackage } from "../cryptography/types";
+import { KeyPair } from "near-api-js";
 
 export default class VotingPackageBuilder {
     private accountId: AccountId;
@@ -22,31 +23,30 @@ export default class VotingPackageBuilder {
     }
 
     /// Builds a voting package with the given encryption key and private key.
-    /// Private key should follow next standard: 'ed25519:bs58private_key' or 'secp256k1:bs58private_key'
-    signPackage(privateKey: string): VotingPackageEncryptor {
+    signPackage(keyPair: KeyPair): VotingPackageEncryptor {
         const votes = Array.from(this.votes.entries()).map(([candidate, weight]) => ({ candidate, weight }));
 
-        const signature = createSignature(base_encode(JSON.stringify({ accountId: this.accountId, votes })), privateKey);
+        const signature = createSignature(base_encode(JSON.stringify({ accountId: this.accountId, votes })), keyPair);
 
         if (!signature) {
             throw new Error("Failed to sign the voting package");
         }
 
-        return new VotingPackageEncryptor(privateKey, {
+        return new VotingPackageEncryptor(keyPair, {
             accountId: this.accountId,
             votes,
             signature,
         }, this.accountId);
     }
 }
 
-class VotingPackageEncryptor {
-    private privateKey: string;
+export class VotingPackageEncryptor {
+    private keyPair: KeyPair;
     private vpackage: VotingPackage;
     private accountId: AccountId;
 
-    constructor(privateKey: string, vpackage: VotingPackage, accountId: AccountId) {
-        this.privateKey = privateKey;
+    constructor(keyPair: KeyPair, vpackage: VotingPackage, accountId: AccountId) {
+        this.keyPair = keyPair;
         this.vpackage = vpackage;
         this.accountId = accountId;
     }
@@ -59,25 +59,25 @@ class VotingPackageEncryptor {
             throw new Error(`Failed to encrypt the voting package: ${encryptedData.error}`);
         }
 
-        return new EncryptedVotingPackageSigner(this.privateKey, encryptedData.data, this.accountId);
+        return new EncryptedVotingPackageSigner(this.keyPair, encryptedData.data, this.accountId);
     }
 
 }
 
-class EncryptedVotingPackageSigner {
-    private privateKey: string;
+export class EncryptedVotingPackageSigner {
+    private keyPair: KeyPair;
     private vpackage: EncryptedVotingPackage;
     private accountId: AccountId;
 
-    constructor(privateKey: string, vpackage: EncryptedVotingPackage, accountId: AccountId) {
-        this.privateKey = privateKey;
+    constructor(keyPair: KeyPair, vpackage: EncryptedVotingPackage, accountId: AccountId) {
+        this.keyPair = keyPair;
         this.vpackage = vpackage;
         this.accountId = accountId;
     }
 
     /// Signs the encrypted voting package
     signPackage(): EncryptedVotingPackageWithProof {
-        const signature = createSignature(this.vpackage.encryptedData + this.vpackage.publicKey, this.privateKey);
+        const signature = createSignature(this.vpackage.encryptedData + this.vpackage.publicKey, this.keyPair);
 
         if (!signature) {
             throw new Error("Failed to sign the encrypted voting package");

relayer/testExample.ts

@@ -0,0 +1,79 @@
+import { randomBytes } from "crypto";
+import VotingPackageBuilder, { VotingPackageEncryptor } from "./src/utils/vote";
+import { keyStores } from "near-api-js";
+import { privateKeyVerify } from "secp256k1";
+import { base_decode } from "near-api-js/lib/utils/serialize";
+
+// you probably need to implement this function more securely
+const generateKeyPair = () => {
+    let secret;
+    do {
+        secret = randomBytes(32)
+    } while (!privateKeyVerify(secret))
+
+    return secret;
+}
+
+const secretPubKey = async (): Promise<Uint8Array> => {
+    const pubkey = await fetch("http://localhost:3000/api/encryption-public-key");
+    const string = await pubkey.text();
+
+    return base_decode(string);
+}
+
+const encryptAndSendPackage = async (votingPackageEncryptor: VotingPackageEncryptor) => {
+    const secretPublicKey = await secretPubKey();
+    const encryptionKey = generateKeyPair();
+
+    const encrypted = (await votingPackageEncryptor.encryptPackage(encryptionKey, secretPublicKey)).signPackage();
+
+    const response = await fetch("http://localhost:3000/api/vote", {
+        method: "POST",
+        body: JSON.stringify(encrypted),
+        headers: {
+            "Content-Type": "application/json"
+        }
+    });
+
+    if (response.status !== 200) {
+        throw new Error(`Error sending vote: ${await response.text()}`);
+    }
+}
+
+const main = async () => {
+    const homedir = require("os").homedir();
+    const CREDENTIALS_DIR = ".near-credentials";
+    const credentialsPath = require("path").join(homedir, CREDENTIALS_DIR);
+    const myKeyStore = new keyStores.UnencryptedFileSystemKeyStore(credentialsPath);
+    const yurtur = await myKeyStore.getKey("testnet", "yurtur.testnet");
+
+    const pack = new VotingPackageBuilder("yurtur.testnet")
+        .addVote("yurtur.testnet", 1)
+        .addVote("absurd-jam.testnet", 5)
+        .signPackage(yurtur);
+
+    // Send first vote
+    await encryptAndSendPackage(pack);
+
+    // Send invalid vote
+    const invalidPack = new VotingPackageBuilder("yurtur.testnet")
+        .addVote("yurtur.testnet", 1)
+        .addVote("absurd-jam.haha", 5)
+        .signPackage(yurtur);
+
+    await encryptAndSendPackage(invalidPack);
+
+    // re-vote for yurtur
+    const revote = new VotingPackageBuilder("yurtur.testnet")
+        .addVote("yurtur.testnet", 6)
+        .signPackage(yurtur);
+
+    await encryptAndSendPackage(revote);
+
+    const absurdJam = await myKeyStore.getKey("testnet", "absurd-jam.testnet");
+    const absurdJampack = new VotingPackageBuilder("absurd-jam.testnet").addVote("yurtur.testnet", 1).addVote("absurd-jam.testnet", 1).signPackage(absurdJam);
+
+    await encryptAndSendPackage(absurdJampack);
+}
+
+main()