linera_client/
wallet.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{collections::BTreeMap, iter::IntoIterator};
5
6use linera_base::{
7    crypto::CryptoHash,
8    data_types::{BlockHeight, ChainDescription, Epoch, Timestamp},
9    ensure,
10    identifiers::{AccountOwner, ChainId},
11};
12use linera_core::{client::PendingProposal, data_types::ChainInfo};
13use serde::{Deserialize, Serialize};
14
15use crate::{config::GenesisConfig, error, Error};
16
17#[derive(Serialize, Deserialize)]
18pub struct Wallet {
19    pub chains: BTreeMap<ChainId, UserChain>,
20    pub default: Option<ChainId>,
21    genesis_config: GenesisConfig,
22}
23
24impl Extend<UserChain> for Wallet {
25    fn extend<Chains: IntoIterator<Item = UserChain>>(&mut self, chains: Chains) {
26        for chain in chains.into_iter() {
27            if !self.chains.contains_key(&chain.chain_id) {
28                self.insert(chain);
29            }
30        }
31    }
32}
33
34impl Wallet {
35    pub fn new(genesis_config: GenesisConfig) -> Self {
36        Wallet {
37            chains: BTreeMap::new(),
38            default: None,
39            genesis_config,
40        }
41    }
42
43    pub fn get(&self, chain_id: ChainId) -> Option<&UserChain> {
44        self.chains.get(&chain_id)
45    }
46
47    pub fn insert(&mut self, chain: UserChain) {
48        if self.default.is_none() {
49            self.default = Some(chain.chain_id);
50        }
51
52        self.chains.insert(chain.chain_id, chain);
53    }
54
55    pub fn forget_keys(&mut self, chain_id: &ChainId) -> Result<AccountOwner, Error> {
56        let chain = self
57            .chains
58            .get_mut(chain_id)
59            .ok_or(error::Inner::NonexistentChain(*chain_id))?;
60
61        let owner = chain
62            .owner
63            .take()
64            .ok_or(error::Inner::NonexistentKeypair(*chain_id))?;
65
66        Ok(owner)
67    }
68
69    pub fn forget_chain(&mut self, chain_id: &ChainId) -> Result<UserChain, Error> {
70        let user_chain = self
71            .chains
72            .remove(chain_id)
73            .ok_or::<Error>(error::Inner::NonexistentChain(*chain_id).into())?;
74        Ok(user_chain)
75    }
76
77    pub fn default_chain(&self) -> Option<ChainId> {
78        self.default
79    }
80
81    pub fn first_non_admin_chain(&self) -> Option<ChainId> {
82        self.chain_ids()
83            .into_iter()
84            .find(|chain_id| *chain_id != self.genesis_config.admin_id())
85    }
86
87    pub fn chain_ids(&self) -> Vec<ChainId> {
88        self.chains.keys().copied().collect()
89    }
90
91    /// Returns the list of all chain IDs for which we have a secret key.
92    pub fn owned_chain_ids(&self) -> Vec<ChainId> {
93        self.chains
94            .iter()
95            .filter_map(|(chain_id, chain)| chain.owner.is_some().then_some(*chain_id))
96            .collect()
97    }
98
99    pub fn num_chains(&self) -> usize {
100        self.chains.len()
101    }
102
103    pub fn chains_mut(&mut self) -> impl Iterator<Item = &mut UserChain> {
104        self.chains.values_mut()
105    }
106
107    pub fn assign_new_chain_to_owner(
108        &mut self,
109        owner: AccountOwner,
110        chain_id: ChainId,
111        timestamp: Timestamp,
112        epoch: Epoch,
113    ) -> Result<(), Error> {
114        let user_chain = UserChain {
115            chain_id,
116            owner: Some(owner),
117            block_hash: None,
118            timestamp,
119            next_block_height: BlockHeight(0),
120            pending_proposal: None,
121            epoch: Some(epoch),
122        };
123        self.insert(user_chain);
124        Ok(())
125    }
126
127    pub fn set_default_chain(&mut self, chain_id: ChainId) -> Result<(), Error> {
128        ensure!(
129            self.chains.contains_key(&chain_id),
130            error::Inner::NonexistentChain(chain_id)
131        );
132        self.default = Some(chain_id);
133        Ok(())
134    }
135
136    pub fn update_from_info(
137        &mut self,
138        pending_proposal: Option<PendingProposal>,
139        owner: Option<AccountOwner>,
140        info: &ChainInfo,
141    ) {
142        self.insert(UserChain {
143            chain_id: info.chain_id,
144            owner,
145            block_hash: info.block_hash,
146            next_block_height: info.next_block_height,
147            timestamp: info.timestamp,
148            pending_proposal,
149            epoch: Some(info.epoch),
150        });
151    }
152
153    pub fn genesis_admin_chain(&self) -> ChainId {
154        self.genesis_config.admin_id()
155    }
156
157    pub fn genesis_config(&self) -> &GenesisConfig {
158        &self.genesis_config
159    }
160}
161
162#[derive(Serialize, Deserialize)]
163pub struct UserChain {
164    pub chain_id: ChainId,
165    /// The owner of the chain, if we own it.
166    pub owner: Option<AccountOwner>,
167    pub block_hash: Option<CryptoHash>,
168    pub timestamp: Timestamp,
169    pub next_block_height: BlockHeight,
170    pub pending_proposal: Option<PendingProposal>,
171    pub epoch: Option<Epoch>,
172}
173
174impl Clone for UserChain {
175    fn clone(&self) -> Self {
176        Self {
177            chain_id: self.chain_id,
178            owner: self.owner,
179            block_hash: self.block_hash,
180            timestamp: self.timestamp,
181            next_block_height: self.next_block_height,
182            pending_proposal: self.pending_proposal.clone(),
183            epoch: self.epoch,
184        }
185    }
186}
187
188impl UserChain {
189    /// Create a user chain that we own.
190    pub fn make_initial(
191        owner: AccountOwner,
192        description: ChainDescription,
193        timestamp: Timestamp,
194    ) -> Self {
195        let epoch = description.config().epoch;
196        Self {
197            chain_id: description.into(),
198            owner: Some(owner),
199            block_hash: None,
200            timestamp,
201            next_block_height: BlockHeight::ZERO,
202            pending_proposal: None,
203            epoch: Some(epoch),
204        }
205    }
206
207    /// Creates an entry for a chain that we don't own. The timestamp must be the genesis
208    /// timestamp or earlier. The Epoch is `None`.
209    pub fn make_other(chain_id: ChainId, timestamp: Timestamp) -> Self {
210        Self {
211            chain_id,
212            owner: None,
213            block_hash: None,
214            timestamp,
215            next_block_height: BlockHeight::ZERO,
216            pending_proposal: None,
217            epoch: None,
218        }
219    }
220}