1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171

use candid::{CandidType, Nat};
use ex3_canister_named_canisters::CanisterNames;
use ex3_canister_types::range::ReportRange;
use ex3_canister_types::report::ShardingReportSubmitter;
use ex3_crypto::sha256;
use ex3_node_types::settings::CandidWalletRegistrySetting;
use ex3_node_types::transaction::EncodedTransaction;
use ex3_node_types::CanisterId;
use ex3_node_types::{CandidBlockHeight, MerkleRoot};
use ex3_serde::bincode;
use ex3_timestamp::TimeInNs;
use ic_stable_structures::storable::Bound;
use ic_stable_structures::Storable;
use rs_merkle::algorithms::Sha256;
use rs_merkle::MerkleProof;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

pub mod error;

#[derive(CandidType, Deserialize, Debug, Clone)]
pub struct WalletRegistryInitArgs {
    #[cfg(feature = "dev_env")]
    pub dev_named_canister_ids: HashMap<CanisterNames, CanisterId>,
    /// Wallet registry seq id
    pub seq_id: u64,
    /// The start of wallet id that can be served by this wallet registry canister.
    pub wallet_start_id: Nat,
    /// Maximum number of wallets that can be served by this wallet registry container
    pub max_wallet_count: Nat,
}

impl From<WalletRegistryInitArgs> for CandidWalletRegistrySetting {
    fn from(args: WalletRegistryInitArgs) -> Self {
        CandidWalletRegistrySetting {
            wallet_start_id: args.wallet_start_id,
            max_wallet_count: args.max_wallet_count,
        }
    }
}

/// Wallet registry tx sharding sub report
///
/// Note: This is a sub report of wallet registry tx sharding report, contains register txs
///       and wallet asset relation operation txs.
#[derive(CandidType, Debug, Default, Clone, Deserialize, Serialize, PartialEq, Eq)]
pub struct WalletTxShardingSubReport {
    /// merkle proof of the register txs
    pub merkle_proof: Vec<u8>,
    pub range: ReportRange,
    pub txs: Vec<EncodedTransaction>,
    pub tx_indexes: Vec<u32>,
}

impl Storable for WalletTxShardingSubReport {
    fn to_bytes(&self) -> std::borrow::Cow<[u8]> {
        bincode::serialize(self).unwrap().into()
    }

    fn from_bytes(bytes: std::borrow::Cow<[u8]>) -> Self {
        bincode::deserialize(bytes.as_ref()).unwrap()
    }

    const BOUND: Bound = Bound::Unbounded;
}

impl WalletTxShardingSubReport {
    pub fn merkle_verify(&self, merkle_root: &[u8; 32]) -> bool {
        let merkle_proof: MerkleProof<Sha256> = match MerkleProof::from_bytes(&self.merkle_proof) {
            Ok(proof) => proof,
            Err(_) => {
                return false;
            }
        };

        let leaves_to_prove: Vec<[u8; 32]> =
            self.txs.iter().map(|tx| sha256(tx.as_ref())).collect();

        let leaf_indexes: Vec<usize> = self
            .tx_indexes
            .iter()
            .map(|index| *index as usize)
            .collect();
        merkle_proof.verify(
            merkle_root.clone(),
            &leaf_indexes,
            &leaves_to_prove,
            self.range.total_count as usize,
        )
    }
}

#[derive(CandidType, Debug, Default, Clone, Deserialize, Serialize, PartialEq, Eq)]
pub struct WalletRegistryDataIntegrityShardingSubReport {
    pub merkle_proof: Vec<u8>,

    /// The data digest of the sharding report
    pub data_digest: [u8; 32],

    /// The merkle node count of the data integrity merkle tree
    pub merkle_node_count: u32,

    /// Index of wallet registry data integrity in merkle tree
    pub index: u32,

    /// Total count wallet register tx in wallet registry sharding sub report
    pub total_register_tx_count: u32,
}

impl Storable for WalletRegistryDataIntegrityShardingSubReport {
    fn to_bytes(&self) -> std::borrow::Cow<[u8]> {
        bincode::serialize(self).unwrap().into()
    }

    fn from_bytes(bytes: std::borrow::Cow<[u8]>) -> Self {
        bincode::deserialize(bytes.as_ref()).unwrap()
    }

    const BOUND: Bound = Bound::Unbounded;
}

impl WalletRegistryDataIntegrityShardingSubReport {
    pub fn merkle_verify(&self, merkle_root: &MerkleRoot) -> bool {
        let merkle_proof: Option<MerkleProof<Sha256>> =
            match MerkleProof::from_bytes(&self.merkle_proof) {
                Ok(merkle_proof) => Some(merkle_proof),
                Err(_) => None,
            };

        match merkle_proof {
            Some(merkle_proof) => merkle_proof.verify(
                merkle_root.clone(),
                vec![self.index as usize].as_ref(),
                vec![self.data_digest].as_ref(),
                self.merkle_node_count as usize,
            ),
            None => false,
        }
    }
}

#[derive(CandidType, Debug, Default, Clone, Deserialize, Eq, PartialEq)]
pub struct ReportStatus {
    /// if submitted sharding report of wallet register txs
    pub submitted_register_tx_report: bool,

    /// sharding range is the full data of the sharding block
    /// report range is the current progress
    pub register_tx_report_progress: Option<ReportRange>,

    /// if submitted register id cursor report
    pub submitted_data_integrity_report: bool,

    pub validated: bool,
}

#[derive(CandidType, Debug, Clone, Deserialize)]
pub struct ShardingReportProgressInfo {
    pub current_block_height: CandidBlockHeight,
    pub in_progress: bool,
    pub report_status: ReportStatus,
    pub major_submitter: Option<ShardingReportSubmitter>,
    pub secondary_submitter: Option<ShardingReportSubmitter>,
    pub now: TimeInNs,
}

#[derive(CandidType, Debug, Clone, Deserialize, Serialize)]
pub enum ShardingSubReport {
    WalletRegister(WalletTxShardingSubReport),
    DataIntegrity(WalletRegistryDataIntegrityShardingSubReport),
}