ic_vetkd_notes/
lib.rs

1use candid::{CandidType, Decode, Deserialize, Encode, Principal};
2use custom_debug::Debug;
3use ic_stable_structures::{storable::Bound, Storable};
4use std::{
5    borrow::Cow,
6    collections::{HashMap, HashSet},
7    hash::Hash,
8};
9
10#[derive(Clone, Debug, CandidType, Deserialize, Eq, PartialEq, Hash)]
11pub struct PrincipalRule {
12    when: Option<u64>,
13    was_read: bool,
14}
15
16impl PrincipalRule {
17    pub fn when(&self) -> Option<u64> {
18        self.when
19    }
20    pub fn was_read(&self) -> bool {
21        self.was_read
22    }
23}
24
25pub const EVERYONE: &str = "everyone";
26pub type NoteId = u128;
27
28#[derive(Clone, Debug, CandidType, Deserialize, Eq, PartialEq, Hash)]
29pub struct HistoryEntry {
30    action: String,
31    user: String,
32    rule: Option<(String, Option<u64>)>,
33    labels: Vec<String>,
34    created_at: u64,
35}
36
37impl HistoryEntry {
38    pub fn action(&self) -> String {
39        self.action.clone()
40    }
41    pub fn user(&self) -> String {
42        self.user.clone()
43    }
44    pub fn rule(&self) -> Option<(String, Option<u64>)> {
45        self.rule.clone()
46    }
47    pub fn created_at(&self) -> u64 {
48        self.created_at
49    }
50    pub fn labels(&self) -> Vec<String> {
51        self.labels.clone()
52    }
53}
54
55#[derive(Clone, Debug, CandidType, Deserialize, Eq, PartialEq)]
56pub struct EncryptedNote {
57    id: NoteId,
58    #[debug(skip)]
59    encrypted_text: String,
60    data: String,
61    owner: String,
62    /// Principals with whom this note is shared. Does not include the owner.
63    /// Needed to be able to efficiently show in the UI with whom this note is shared.
64    users: HashMap<String, PrincipalRule>,
65
66    locked: bool,
67    read_by: HashSet<String>,
68    created_at: u64,
69    updated_at: u64,
70    history: Vec<HistoryEntry>,
71}
72
73impl Default for EncryptedNote {
74    fn default() -> Self {
75        EncryptedNote {
76            id: 0,
77            encrypted_text: "".to_string(),
78            data: "".to_string(),
79            owner: "".to_string(),
80            users: HashMap::new(),
81            locked: false,
82            created_at: ic_cdk::api::time(),
83            updated_at: ic_cdk::api::time(),
84            history: vec![],
85            read_by: HashSet::new(),
86        }
87    }
88}
89
90impl EncryptedNote {
91    pub fn create(id: NoteId) -> Self {
92        let user = &caller().to_text();
93        EncryptedNote {
94            id,
95            owner: user.clone(),
96            data: String::new(),
97            users: HashMap::new(),
98            encrypted_text: String::new(),
99            locked: false,
100            read_by: HashSet::new(),
101            created_at: ic_cdk::api::time(),
102            updated_at: ic_cdk::api::time(),
103            history: vec![HistoryEntry {
104                action: "created".to_string(),
105                labels: vec![],
106                user: user.clone(),
107                rule: None,
108                created_at: ic_cdk::api::time(),
109            }],
110        }
111    }
112    pub fn id(&self) -> NoteId {
113        self.id
114    }
115    pub fn data(&self) -> String {
116        self.data.clone()
117    }
118    pub fn read_by(&self) -> HashSet<String> {
119        self.read_by.clone()
120    }
121    pub fn encrypted_text(&self) -> String {
122        self.encrypted_text.clone()
123    }
124    pub fn owner(&self) -> String {
125        self.owner.clone()
126    }
127    pub fn users(&self) -> HashMap<String, PrincipalRule> {
128        self.users.clone()
129    }
130    pub fn locked(&self) -> bool {
131        self.locked
132    }
133    pub fn created_at(&self) -> u64 {
134        self.created_at
135    }
136    pub fn updated_at(&self) -> u64 {
137        self.updated_at
138    }
139    pub fn history(&self) -> Vec<HistoryEntry> {
140        self.history.clone()
141    }
142    // Check if the user is owner or has access to the note as of right now
143    pub fn is_authorized(&self) -> bool {
144        let user = &caller().to_text();
145        if user == &self.owner {
146            return true;
147        }
148        // once a non-owner reads a note it's locked and can no longer
149        // be updated
150        if let Some(r) = self.users.get(user) {
151            if r.when.is_none() || r.when.unwrap() <= ic_cdk::api::time() {
152                return true;
153            }
154        } else if let Some(r) = self.users.get(EVERYONE) {
155            if r.when.is_none() || r.when.unwrap() <= ic_cdk::api::time() {
156                return true;
157            }
158        }
159        false
160    }
161    // Same as above but mark it as being read by that user
162    pub fn lock_authorized(&mut self) -> bool {
163        let user = &caller().to_text();
164        if user == &self.owner {
165            let id = self.id;
166            ic_cdk::println!("note not locked with ID {id} as {self:#?} retrieving owner {user}");
167            return true;
168        }
169        // once a non-owner reads a note it's locked and can no longer
170        // be updated
171        if let Some(r) = self.users.get_mut(user) {
172            if r.when.is_none() || r.when.unwrap() <= ic_cdk::api::time() {
173                r.was_read = true;
174                if !self.read_by.contains(user) {
175                    self.history.push(HistoryEntry {
176                        action: "read".to_string(),
177                        user: user.to_string(),
178                        labels: if self.locked {
179                            vec![]
180                        } else {
181                            vec!["locked".to_string()]
182                        },
183                        rule: Some((user.clone(), r.when)),
184                        created_at: ic_cdk::api::time(),
185                    });
186                }
187                self.locked = true;
188                self.read_by.insert(user.to_string());
189
190                let id = self.id;
191                ic_cdk::println!("locked note with ID {id} as {self:#?} retrieving from {user}");
192
193                return true;
194            }
195        } else if let Some(r) = self.users.get_mut(EVERYONE) {
196            if r.when.is_none() || r.when.unwrap() <= ic_cdk::api::time() {
197                r.was_read = true;
198                if !self.read_by.contains(user) {
199                    self.history.push(HistoryEntry {
200                        action: "read".to_string(),
201                        user: user.to_string(),
202                        labels: if self.locked {
203                            vec![]
204                        } else {
205                            vec!["locked".to_string()]
206                        },
207                        rule: Some((EVERYONE.to_string(), r.when)),
208                        created_at: ic_cdk::api::time(),
209                    });
210                }
211                self.read_by.insert(user.to_string());
212                self.locked = true;
213
214                let id = self.id;
215                ic_cdk::println!(
216                    "locked note with ID {id} as {self:#?} retrieving from {user} as everyone"
217                );
218
219                return true;
220            }
221        }
222        false
223    }
224    // add a new reader to the note
225    pub fn add_reader(&mut self, user: &Option<String>, when: Option<u64>) -> bool {
226        if self.locked && (user.is_none() || self.read_by.contains(&user.clone().unwrap())) {
227            // If this note is locked and the user has already read it then this doesn't seem useful.
228            return false;
229        }
230        let user_name = user.clone().unwrap_or_else(|| EVERYONE.to_string());
231        self.history.append(&mut vec![HistoryEntry {
232            action: "share".to_string(),
233            labels: vec![],
234            user: user_name.clone(),
235            rule: Some((user_name.clone(), when)),
236            created_at: ic_cdk::api::time(),
237        }]);
238        self.users.insert(
239            user_name,
240            PrincipalRule {
241                was_read: false,
242                when,
243            },
244        );
245        true
246    }
247    // Was the note ever read by that user
248    pub fn user_read(&self, user: &String) -> bool {
249        self.read_by.contains(user)
250    }
251    // Remove a reader (will return false if the note was already read by the user)
252    pub fn remove_reader(&mut self, user: &Option<String>) -> bool {
253        if self.locked {
254            if user.iter().any(|u| self.read_by.contains(u)) {
255                return false;
256            } else if let Some(r) = self
257                .users
258                .get(&user.clone().unwrap_or(EVERYONE.to_string()))
259            {
260                if r.was_read {
261                    return false;
262                }
263            }
264        }
265        let user_name = user.clone().unwrap_or_else(|| EVERYONE.to_string());
266        if self.users.contains_key(&user_name) {
267            self.users.remove(user_name.as_str());
268            self.history.push(HistoryEntry {
269                action: "unshare".to_string(),
270                labels: vec![],
271                user: user_name.clone(),
272                rule: None,
273                created_at: ic_cdk::api::time(),
274            });
275
276            true
277        } else {
278            false
279        }
280    }
281    // Update the data. This is only allowed by the owner before the note was locked
282    pub fn set_data(&mut self, data: String) -> bool {
283        let user = caller().to_text();
284        if self.locked && user != self.owner {
285            return false;
286        }
287        self.data = data;
288        self.updated_at = ic_cdk::api::time();
289        self.history.push(HistoryEntry {
290            action: "updated".to_string(),
291            labels: vec!["data".to_string()],
292            user: user.clone(),
293            rule: None,
294            created_at: ic_cdk::api::time(),
295        });
296        true
297    }
298    pub fn set_encrypted_text(&mut self, encrypted_text: String) -> bool {
299        let user = caller().to_text();
300        if self.locked && user != self.owner {
301            return false;
302        }
303        self.encrypted_text = encrypted_text;
304        self.updated_at = ic_cdk::api::time();
305        self.history.push(HistoryEntry {
306            action: "updated".to_string(),
307            labels: vec!["encrypted_text".to_string()],
308            user: user.clone(),
309            rule: None,
310            created_at: ic_cdk::api::time(),
311        });
312        true
313    }
314    pub fn set_data_and_encrypted_text(&mut self, data: String, encrypted_text: String) -> bool {
315        let user = caller().to_text();
316        if self.locked && user != self.owner {
317            return false;
318        }
319        self.data = data;
320        self.encrypted_text = encrypted_text;
321        self.updated_at = ic_cdk::api::time();
322        self.history.push(HistoryEntry {
323            action: "updated".to_string(),
324            labels: vec!["data".to_string(), "encrypted_text".to_string()],
325            user: user.clone(),
326            rule: None,
327            created_at: ic_cdk::api::time(),
328        });
329        true
330    }
331    // Is the note shared at all?
332    pub fn is_shared(&self) -> bool {
333        !self.users.is_empty()
334    }
335    // Has any reader read it?
336    pub fn is_locked(&self) -> bool {
337        self.locked
338    }
339}
340
341impl Storable for EncryptedNote {
342    fn to_bytes(&self) -> Cow<[u8]> {
343        Cow::Owned(Encode!(self).unwrap())
344    }
345    fn from_bytes(bytes: Cow<[u8]>) -> Self {
346        Decode!(bytes.as_ref(), Self).unwrap()
347    }
348    const BOUND: Bound = Bound::Unbounded;
349}
350
351/// Unlike Motoko, the caller identity is not built into Rust.
352/// Thus, we use the ic_cdk::caller() method inside this wrapper function.
353/// The wrapper prevents the use of the anonymous identity. Forbidding anonymous
354/// interactions is the recommended default behavior for IC canisters.
355fn caller() -> Principal {
356    let caller = ic_cdk::caller();
357    // The anonymous principal is not allowed to interact with the
358    // encrypted notes canister.
359    if caller == Principal::anonymous() {
360        panic!("Anonymous principal not allowed to make calls.")
361    }
362    caller
363}
364
365mod vetkd_types;
366
367pub const VETKD_SYSTEM_API_CANISTER_ID: &str = "nn664-2iaaa-aaaao-a3tqq-cai";
368
369use vetkd_types::{
370    CanisterId, VetKDCurve, VetKDEncryptedKeyReply, VetKDEncryptedKeyRequest, VetKDKeyId,
371    VetKDPublicKeyReply, VetKDPublicKeyRequest,
372};
373
374pub async fn symmetric_key_verification_key_for_note() -> String {
375    let request = VetKDPublicKeyRequest {
376        canister_id: None,
377        derivation_path: vec![b"note_symmetric_key".to_vec()],
378        key_id: bls12_381_test_key_1(),
379    };
380
381    let (response,): (VetKDPublicKeyReply,) = ic_cdk::call(
382        vetkd_system_api_canister_id(),
383        "vetkd_public_key",
384        (request,),
385    )
386    .await
387    .expect("call to vetkd_public_key failed");
388
389    hex::encode(response.public_key)
390}
391
392// Be careful this routine will not be able to actully determine whether the
393// current user is permitted to read the note. I.e. anyone can ask for any
394// derivation path.
395pub async fn encrypted_symmetric_key_for_note(
396    note_id: NoteId,
397    owner: &String,
398    encryption_public_key: Vec<u8>,
399) -> String {
400    let request = VetKDEncryptedKeyRequest {
401        derivation_id: {
402            let mut buf = vec![];
403            buf.extend_from_slice(&note_id.to_be_bytes()); // fixed-size encoding
404            buf.extend_from_slice(owner.as_bytes());
405            buf // prefix-free
406        },
407        public_key_derivation_path: vec![b"note_symmetric_key".to_vec()],
408        key_id: bls12_381_test_key_1(),
409        encryption_public_key,
410    };
411
412    let (response,): (VetKDEncryptedKeyReply,) = ic_cdk::call(
413        vetkd_system_api_canister_id(),
414        "vetkd_derive_encrypted_key",
415        (request,),
416    )
417    .await
418    .expect("call to vetkd_derive_encrypted_key failed");
419
420    hex::encode(response.encrypted_key)
421}
422
423fn bls12_381_test_key_1() -> VetKDKeyId {
424    VetKDKeyId {
425        curve: VetKDCurve::Bls12_381,
426        name: "test_key_1".to_string(),
427    }
428}
429
430pub fn vetkd_system_api_canister_id() -> CanisterId {
431    use std::str::FromStr;
432    CanisterId::from_str(VETKD_SYSTEM_API_CANISTER_ID).expect("failed to create canister ID")
433}
ic_vetkd_notes/lib.rs

ic_vetkd_notes/
lib.rs
