Save and Read wallet

rust-lightclient/.gitignore

@@ -2,3 +2,4 @@ target/
 Cargo.lock
 .vscode/
 history.txt
+wallet.dat

rust-lightclient/Cargo.toml

@@ -23,30 +23,28 @@ rustyline = "5.0.2"
 byteorder = "1"
 
 [dependencies.bellman]
-git = "https://github.com/adityapk00/librustzcash.git"
-branch = "lightclient-work"
+path = "../../librustzcash/bellman"
 default-features = false
 features = ["groth16"]
 
 [dependencies.pairing]
-git = "https://github.com/adityapk00/librustzcash.git"
-branch = "lightclient-work"
+path = "../../librustzcash/pairing"
 
 [dependencies.zcash_client_backend]
-git = "https://github.com/adityapk00/librustzcash.git"
-branch = "lightclient-work"
+path = "../../librustzcash/zcash_client_backend"
 default-features = false
 
 [dependencies.zcash_primitives]
-git = "https://github.com/adityapk00/librustzcash.git"
-branch = "lightclient-work"
+path = "../../librustzcash/zcash_primitives"
 default-features = false
 
 [dependencies.zcash_proofs]
-git = "https://github.com/adityapk00/librustzcash.git"
-branch = "lightclient-work"
+path = "../../librustzcash/zcash_proofs"
 default-features = false
 
+[dependencies.ff]
+path = "../../librustzcash/ff"
+features = ["ff_derive"]
 
 [build-dependencies]
 tower-grpc-build = { git = "https://github.com/tower-rs/tower-grpc",  features = ["tower-hyper"]  }

rust-lightclient/src/commands.rs

@@ -7,7 +7,7 @@ pub trait Command {
 
     fn short_help(&self) -> String;
 
-    fn exec(&self, args: &[String], lightclient: &LightClient);
+    fn exec(&self, args: &[String], lightclient: &mut LightClient);
 }
 
 struct SyncCommand {}
@@ -21,7 +21,7 @@ impl Command for SyncCommand {
         "Download CompactBlocks and sync to the server".to_string()
     }
 
-    fn exec(&self, args: &[String], lightclient: &LightClient) {
+    fn exec(&self, args: &[String], lightclient: &mut LightClient) {
         lightclient.do_sync();
     }
 }
@@ -37,7 +37,7 @@ impl Command for HelpCommand {
         "Lists all available commands".to_string()
     }
 
-    fn exec(&self, args: &[String], _: &LightClient) {
+    fn exec(&self, args: &[String], _: &mut LightClient) {
         // Print a list of all commands
         get_commands().iter().for_each(| (cmd, obj) | {
             println!("{} - {}", cmd, obj.short_help());
@@ -55,7 +55,7 @@ impl Command for InfoCommand {
         "Get the lightwalletd server's info".to_string()
     }
 
-    fn exec(&self, args: &[String], lightclient: &LightClient) {
+    fn exec(&self, args: &[String], lightclient: &mut LightClient) {
         lightclient.do_info();
     }
 }
@@ -70,7 +70,7 @@ impl Command for AddressCommand {
         "List all current addresses".to_string()
     }
 
-    fn exec(&self, args: &[String], lightclient: &LightClient) {
+    fn exec(&self, args: &[String], lightclient: &mut LightClient) {
         lightclient.do_address();
     }
 }
@@ -85,7 +85,7 @@ impl Command for SendCommand {
         "Send ZEC to the given address".to_string()
     }
 
-    fn exec(&self, args: &[String], lightclient: &LightClient) {
+    fn exec(&self, args: &[String], lightclient: &mut LightClient) {
         lightclient.do_send(
             "ztestsapling1x65nq4dgp0qfywgxcwk9n0fvm4fysmapgr2q00p85ju252h6l7mmxu2jg9cqqhtvzd69jwhgv8d".to_string(), 
             1500000, 
@@ -93,6 +93,36 @@ impl Command for SendCommand {
     }
 }
 
+struct SaveCommand {}
+impl Command for SaveCommand {
+    fn help(&self) {
+        println!("Save wallet to disk");
+    }
+
+    fn short_help(&self) -> String {
+        "Save wallet file to disk".to_string()
+    }
+
+    fn exec(&self, args: &[String], lightclient: &mut LightClient) {
+        lightclient.do_save();
+    }
+}
+
+struct ReadCommand {}
+impl Command for ReadCommand {
+        fn help(&self) {
+        println!("Read wallet from disk");
+    }
+
+    fn short_help(&self) -> String {
+        "Read wallet file from disk".to_string()
+    }
+
+    fn exec(&self, args: &[String], lightclient: &mut LightClient) {
+        lightclient.do_read();
+    }
+}
+
 pub fn get_commands() -> Box<HashMap<String, Box<dyn Command>>> {
     let mut map: HashMap<String, Box<dyn Command>> = HashMap::new();
 
@@ -101,11 +131,13 @@ pub fn get_commands() -> Box<HashMap<String, Box<dyn Command>>> {
     map.insert("address".to_string(),   Box::new(AddressCommand{}));
     map.insert("info".to_string(),      Box::new(InfoCommand{}));
     map.insert("send".to_string(),      Box::new(SendCommand{}));
+    map.insert("save".to_string(),      Box::new(SaveCommand{}));
+    map.insert("read".to_string(),      Box::new(ReadCommand{}));
 
     Box::new(map)
 }
 
-pub fn do_user_command(cmd: String, lightclient: &LightClient) {
+pub fn do_user_command(cmd: String, lightclient: &mut LightClient) {
     match get_commands().get(&cmd) {
         Some(cmd) => cmd.exec(&[], lightclient),
         None      => {

rust-lightclient/src/lightclient.rs

@@ -1,11 +1,12 @@
 use crate::lightwallet::LightWallet;
 
+use std::fs::File;
+use std::io::prelude::*;
+
 use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
 
 use std::error::Error;
-use std::io::prelude::*;
-use std::fs::File;
 
 use zcash_primitives::transaction::{TxId, Transaction};
 use zcash_primitives::note_encryption::Memo;
@@ -60,6 +61,20 @@ impl LightClient {
         println!("Balance: {}", self.wallet.balance());
     }
 
+    pub fn do_read(&mut self) {
+        let mut file_buffer = File::open("wallet.dat").unwrap();
+        
+        let lw = LightWallet::read(&mut file_buffer).unwrap();
+        self.wallet = Arc::new(lw);
+    }
+
+    pub fn do_save(&self) {
+        let mut file_buffer = File::create("wallet.dat").unwrap();
+        
+        self.wallet.write(&mut file_buffer).unwrap();
+    }
+
+
     pub fn do_info(&self) {
         let uri: http::Uri = format!("http://127.0.0.1:9067").parse().unwrap();
 

rust-lightclient/src/lightwallet.rs

@@ -1,10 +1,13 @@
+pub extern crate ff;
+
 use std::time::SystemTime;
+
 use std::io::{self, Read, Write};
 
-use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 
-use pairing::bls12_381::Bls12;
-use zcash_primitives::primitives::{Diversifier, Note, PaymentAddress};
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+use pairing::bls12_381::{Bls12, Fr, FrRepr};
+use zcash_primitives::primitives::{Diversifier, Note, PaymentAddress, /*read_note */ };
 use std::cmp;
 use std::collections::HashMap;
 use std::sync::{Arc, RwLock};
@@ -13,10 +16,12 @@ use zcash_client_backend::{
     constants::testnet::HRP_SAPLING_PAYMENT_ADDRESS, encoding::encode_payment_address,
     proto::compact_formats::CompactBlock, welding_rig::scan_block,
 };
+use ff::{PrimeField, PrimeFieldRepr};
 use zcash_primitives::{
     block::BlockHash,
     merkle_tree::{CommitmentTree, IncrementalWitness},
     sapling::Node,
+    serialize::{Vector, Optional},
     transaction::{
         builder::{Builder},
         components::Amount, components::amount::DEFAULT_FEE,
@@ -26,6 +31,9 @@ use zcash_primitives::{
     zip32::{ExtendedFullViewingKey, ExtendedSpendingKey},
     JUBJUB,
 };
+use zcash_primitives::jubjub::{edwards, FixedGenerators, JubjubEngine, JubjubParams, PrimeOrder,
+    fs::{Fs, FsRepr},    
+};
 
 use crate::address;
 use crate::prover;
@@ -65,6 +73,12 @@ impl BlockData {
 
         let tree = CommitmentTree::<Node>::read(&mut reader)?;
 
+        let endtag = reader.read_u64::<LittleEndian>()?;
+        if endtag != 11 {
+            println!("End tag for blockdata {}", endtag);        
+        }
+
+
         Ok(BlockData{
             height, 
             hash: BlockHash{ 0: hash_bytes }, 
@@ -75,12 +89,15 @@ impl BlockData {
     pub fn write<W: Write>(&self, mut writer: W) -> io::Result<()> {
         writer.write_i32::<LittleEndian>(self.height)?;
         writer.write_all(&self.hash.0)?;
-        self.tree.write(writer)
+        self.tree.write(&mut writer)?;
+        writer.write_u64::<LittleEndian>(11)
     }
 }
 
+
 pub struct SaplingNoteData {
     account: usize,
+    extfvk: ExtendedFullViewingKey, // Technically, this should be recoverable from the account number, but we're going to refactor this in the future, so I'll write it again here. 
     diversifier: Diversifier,
     note: Note<Bls12>,
     witnesses: Vec<IncrementalWitness<Node>>,
@@ -89,6 +106,36 @@ pub struct SaplingNoteData {
     pub memo:  Option<Memo>
 }
 
+
+/// Reads an FsRepr from [u8] of length 32
+/// This will panic (abort) if length provided is
+/// not correct
+/// TODO: This is duplicate from rustzcash.rs
+fn read_fs(from: &[u8]) -> FsRepr {
+    assert_eq!(from.len(), 32);
+
+    let mut f = <<Bls12 as JubjubEngine>::Fs as PrimeField>::Repr::default();
+    f.read_le(from).expect("length is 32 bytes");
+
+    f
+}
+
+// Reading a note also needs the corresponding address to read from.
+pub fn read_note<R: Read>(mut reader: R) -> io::Result<(u64, Fs)> {
+    let value = reader.read_u64::<LittleEndian>()?;
+
+    let mut r_bytes: [u8; 32] = [0; 32];
+    reader.read_exact(&mut r_bytes)?;
+
+    let r = match Fs::from_repr(read_fs(&r_bytes)) {
+        Ok(r) => r,
+        Err(_) => return Err(io::Error::new(
+            io::ErrorKind::InvalidInput, "Couldn't parse randomness"))
+    };
+
+    Ok((value, r))
+}
+
 impl SaplingNoteData {
     fn new(
         extfvk: &ExtendedFullViewingKey,
@@ -105,9 +152,9 @@ impl SaplingNoteData {
             nf
         };
 
-
         SaplingNoteData {
             account: output.account,
+            extfvk: extfvk.clone(),
             diversifier: output.to.diversifier,
             note: output.note,
             witnesses: vec![witness],
@@ -117,15 +164,131 @@ impl SaplingNoteData {
         }
     }
 
-    fn print_note(&self) {
-        
+    // Reading a note also needs the corresponding address to read from.
+    pub fn read<R: Read>(mut reader: R) -> io::Result<Self> {
+        println!("Trying to read sapling note data");
+        // Read the version number first
+        let version = reader.read_u64::<LittleEndian>()?;
+        println!("SaplingNoteData version {}", version);
+
+        let account = reader.read_u64::<LittleEndian>()? as usize;
+
+        println!("Trying to read extfvk");
+        let extfvk = ExtendedFullViewingKey::read(&mut reader)?;
+
+        println!("Trying to read diversifier");
+        let mut diversifier_bytes = [0u8; 11];
+        reader.read_exact(&mut diversifier_bytes)?;
+        let diversifier = Diversifier{0: diversifier_bytes};
+
+        // To recover the note, read the value and r, and then use the payment address
+        // to recreate the note
+        println!("Reading value and r");
+        let (value, r) = read_note(&mut reader)?; // TODO: This method is in a different package, because of some fields that are private
+
+        println!("Reading note");
+        let maybe_note = extfvk.fvk.vk.into_payment_address(diversifier, &JUBJUB).unwrap().create_note(value, r, &JUBJUB);
+
+        let note = match maybe_note {
+            Some(n)  => Ok(n),
+            None     => Err(io::Error::new(io::ErrorKind::InvalidInput, "Couldn't create the note for the address"))
+        }?;
+
+        let witnesses = Vector::read(&mut reader, |r| IncrementalWitness::<Node>::read(r))?;
+
+        let mut nullifier = [0u8; 32];
+        reader.read_exact(&mut nullifier)?;
+
+        let spent = Optional::read(&mut reader, |r| {
+            let mut txid_bytes = [0u8; 32];
+            r.read_exact(&mut txid_bytes)?;
+            Ok(TxId{0: txid_bytes})
+        })?;
+
+        let memo = Optional::read(&mut reader, |r| {
+            let mut memo_bytes = [0u8; 512];
+            r.read_exact(&mut memo_bytes)?;
+            match Memo::from_bytes(&memo_bytes) {
+                Some(m) => Ok(m),
+                None    => Err(io::Error::new(io::ErrorKind::InvalidInput, "Couldn't create the memo"))
+            }
+        })?;
+
+        Ok(SaplingNoteData {
+            account,
+            extfvk,
+            diversifier,
+            note,
+            witnesses,
+            nullifier,
+            spent,
+            memo
+        })
     }
 
+    pub fn write<W: Write>(&self, mut writer: W) -> io::Result<()> {
+        // Write a version number first, so we can later upgrade this if needed.
+        writer.write_u64::<LittleEndian>(1)?;
+
+        writer.write_u64::<LittleEndian>(self.account as u64)?;
+
+        println!("Writing extfvk {:?}", self.extfvk);
+        self.extfvk.write(&mut writer)?;
+
+        writer.write_all(&self.diversifier.0)?;
+
+        // Writing the note means writing the note.value and note.r. The Note is recoverable
+        // from these 2 values and the Payment address. 
+        writer.write_u64::<LittleEndian>(self.note.value)?;
+
+        let mut rcm = [0; 32];
+        self.note.r.into_repr().write_le(&mut rcm[..])?;
+        writer.write_all(&rcm)?;
+
+        Vector::write(&mut writer, &self.witnesses, |wr, wi| wi.write(wr) )?;
+
+        writer.write_all(&self.nullifier)?;
+        Optional::write(&mut writer, &self.spent, |w, t| w.write_all(&t.0))?;
+
+        Optional::write(&mut writer, &self.memo, |w, m| w.write_all(m.as_bytes()))?;
+
+        Ok(())
+    }
+}
+
 pub struct WalletTx {
     block: i32,
     pub notes: Vec<SaplingNoteData>,
 }
 
+impl WalletTx {
+    
+    pub fn read<R: Read>(mut reader: R) -> io::Result<Self> {
+        let version = reader.read_u64::<LittleEndian>()?;
+        println!("Wallet version {}", version);
+        // TODO Assert version?
+
+        let block = reader.read_i32::<LittleEndian>()?;
+
+        let notes = Vector::read(&mut reader, |r| SaplingNoteData::read(r))?;
+
+        Ok(WalletTx{
+            block,
+            notes
+        })
+    }
+
+    pub fn write<W: Write>(&self, mut writer: W) -> io::Result<()> {
+        writer.write_u64::<LittleEndian>(1)?;
+
+        writer.write_i32::<LittleEndian>(self.block)?;
+
+        Vector::write(&mut writer, &self.notes, |w, nd| nd.write(w))?;
+
+        Ok(())
+    }
+}
+
 struct SpendableNote {
     txid: TxId,
     nullifier: [u8; 32],
@@ -167,6 +330,8 @@ impl LightWallet {
         let extfvk = ExtendedFullViewingKey::from(&extsk);
         let address = extfvk.default_address().unwrap().1;
 
+        println!("extfvk is {:?}", extfvk);
+
         LightWallet {
             extsks: [extsk],
             extfvks: [extfvk],
@@ -176,6 +341,51 @@ impl LightWallet {
         }
     }
 
+    pub fn read<R: Read>(mut reader: R) -> io::Result<Self> {
+        let version = reader.read_u64::<LittleEndian>()?;
+        println!("LightWallet version {}", version);
+
+        let blocks = Vector::read(&mut reader, |r| BlockData::read(r))?;
+
+        let txs_tuples = Vector::read(&mut reader, |r| {
+            let mut txid_bytes = [0u8; 32];
+            r.read_exact(&mut txid_bytes)?;
+
+            Ok((TxId{0: txid_bytes}, WalletTx::read(r).unwrap()))
+        })?;
+        let txs = txs_tuples.into_iter().collect::<HashMap<TxId, WalletTx>>();
+
+
+        let extsk = ExtendedSpendingKey::master(&[1; 32]);  // New key
+        let extfvk = ExtendedFullViewingKey::from(&extsk);
+        let address = extfvk.default_address().unwrap().1;
+
+        Ok(LightWallet{
+            extsks: [extsk],
+            extfvks: [extfvk],
+            address,
+            blocks: Arc::new(RwLock::new(blocks)),
+            txs: Arc::new(RwLock::new(txs))
+        })
+    }
+
+    pub fn write<W: Write>(&self, mut writer: W) -> io::Result<()> {
+        // Write the version
+        writer.write_u64::<LittleEndian>(1)?;
+
+        // TODO: Write the keys properly. Right now, they're just hardcoded
+
+        Vector::write(&mut writer, &self.blocks.read().unwrap(), |w, b| b.write(w))?;
+                
+        // The hashmap, write as a set of tuples
+        Vector::write(&mut writer, &self.txs.read().unwrap().iter().collect::<Vec<(&TxId, &WalletTx)>>(),
+                        |w, (k, v)| {
+                            w.write_all(&k.0)?;
+                            v.write(w)
+                        })?;
+        Ok(())
+    }
+
     pub fn set_initial_block(&self, height: i32, hash: &str, sapling_tree: &str) -> bool {
         let mut blocks = self.blocks.write().unwrap();
         if !blocks.is_empty() {

rust-lightclient/src/main.rs

@@ -16,7 +16,7 @@ pub mod grpc_client {
 
 
 pub fn main() {
-    let light_client = LightClient::new();
+    let mut light_client = LightClient::new();
 
     // `()` can be used when no completer is required
     let mut rl = Editor::<()>::new();
@@ -28,7 +28,7 @@ pub fn main() {
         match readline {
             Ok(line) => {
                 rl.add_history_entry(line.as_str());
-                commands::do_user_command(line, &light_client);
+                commands::do_user_command(line, &mut light_client);
             },
             Err(ReadlineError::Interrupted) => {
                 println!("CTRL-C");