Skip to main content

ic_dbms_api/
principal.rs

1use std::fmt;
2
3use candid::CandidType;
4use serde::{Deserialize, Serialize};
5use wasm_dbms_api::prelude::{
6    CustomDataType, CustomValue, DEFAULT_ALIGNMENT, DataSize, DataType, DecodeError, Encode, MSize,
7    MemoryError, MemoryResult, PageOffset, Value,
8};
9
10/// Principal data type for the IC DBMS.
11///
12/// This is an IC-specific custom data type that wraps [`candid::Principal`].
13/// Use with `#[custom_type]` annotation in table definitions.
14#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
15pub struct Principal(pub candid::Principal);
16
17impl Default for Principal {
18    fn default() -> Self {
19        Self(candid::Principal::anonymous())
20    }
21}
22
23impl CandidType for Principal {
24    fn _ty() -> candid::types::Type {
25        candid::types::Type(std::rc::Rc::new(candid::types::TypeInner::Principal))
26    }
27
28    fn idl_serialize<S>(&self, serializer: S) -> Result<(), S::Error>
29    where
30        S: candid::types::Serializer,
31    {
32        candid::Principal::idl_serialize(&self.0, serializer)
33    }
34}
35
36impl Encode for Principal {
37    const SIZE: DataSize = DataSize::Dynamic;
38
39    const ALIGNMENT: PageOffset = DEFAULT_ALIGNMENT;
40
41    fn encode(&'_ self) -> std::borrow::Cow<'_, [u8]> {
42        let principal_bytes = self.0.as_slice();
43        let mut bytes = Vec::with_capacity(1 + principal_bytes.len());
44        let len = principal_bytes.len() as u8;
45        bytes.push(len);
46        bytes.extend_from_slice(principal_bytes);
47        std::borrow::Cow::Owned(bytes)
48    }
49
50    fn decode(data: std::borrow::Cow<[u8]>) -> MemoryResult<Self>
51    where
52        Self: Sized,
53    {
54        if data.is_empty() {
55            return Err(MemoryError::DecodeError(DecodeError::TooShort));
56        }
57
58        let buf_len = data[0] as usize;
59
60        if data.len() < 1 + buf_len {
61            return Err(MemoryError::DecodeError(DecodeError::TooShort));
62        }
63
64        let principal = candid::Principal::try_from_slice(&data[1..1 + buf_len]).map_err(|e| {
65            MemoryError::DecodeError(DecodeError::IdentityDecodeError(e.to_string()))
66        })?;
67
68        Ok(Self(principal))
69    }
70
71    fn size(&self) -> MSize {
72        1 + self.0.as_slice().len() as MSize
73    }
74}
75
76impl fmt::Display for Principal {
77    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
78        write!(f, "{}", self.0)
79    }
80}
81
82impl DataType for Principal {}
83
84impl CustomDataType for Principal {
85    const TYPE_TAG: &'static str = "principal";
86}
87
88impl From<Principal> for Value {
89    fn from(val: Principal) -> Value {
90        Value::Custom(CustomValue::new(&val))
91    }
92}
93
94#[cfg(test)]
95mod tests {
96
97    use super::*;
98
99    #[test]
100    fn test_principal_encode_decode() {
101        let original = Principal(
102            candid::Principal::from_text("ryjl3-tyaaa-aaaaa-aaaba-cai").expect("invalid principal"),
103        );
104        let encoded = original.encode();
105        let decoded = Principal::decode(encoded).unwrap();
106        assert_eq!(original, decoded);
107    }
108
109    #[test]
110    fn test_principal_encode_decode_anonymous() {
111        let original = Principal(candid::Principal::anonymous());
112        let encoded = original.encode();
113        let decoded = Principal::decode(encoded).unwrap();
114        assert_eq!(original, decoded);
115    }
116
117    #[test]
118    fn test_should_candid_encode_decode() {
119        let src = Principal(
120            candid::Principal::from_text("ryjl3-tyaaa-aaaaa-aaaba-cai").expect("invalid principal"),
121        );
122        let buf = candid::encode_one(&src).expect("Candid encoding failed");
123        let decoded: Principal = candid::decode_one(&buf).expect("Candid decoding failed");
124        assert_eq!(src, decoded);
125    }
126
127    #[test]
128    fn test_decode_empty_data_returns_too_short() {
129        let result = Principal::decode(std::borrow::Cow::Borrowed(&[]));
130        assert!(result.is_err());
131        assert!(matches!(
132            result.unwrap_err(),
133            MemoryError::DecodeError(DecodeError::TooShort)
134        ));
135    }
136
137    #[test]
138    fn test_decode_truncated_data_returns_too_short() {
139        // Length prefix says 10 bytes, but only 2 bytes follow
140        let data = vec![10, 0x01, 0x02];
141        let result = Principal::decode(std::borrow::Cow::Owned(data));
142        assert!(result.is_err());
143        assert!(matches!(
144            result.unwrap_err(),
145            MemoryError::DecodeError(DecodeError::TooShort)
146        ));
147    }
148
149    #[test]
150    fn test_size_returns_correct_value() {
151        let principal = Principal(
152            candid::Principal::from_text("ryjl3-tyaaa-aaaaa-aaaba-cai").expect("invalid principal"),
153        );
154        let size = principal.size();
155        // 1 byte for length prefix + principal bytes length
156        assert_eq!(size, 1 + principal.0.as_slice().len() as MSize);
157    }
158
159    #[test]
160    fn test_size_anonymous_principal() {
161        let principal = Principal(candid::Principal::anonymous());
162        let size = principal.size();
163        assert_eq!(size, 1 + principal.0.as_slice().len() as MSize);
164    }
165
166    #[test]
167    fn test_display() {
168        let inner =
169            candid::Principal::from_text("ryjl3-tyaaa-aaaaa-aaaba-cai").expect("invalid principal");
170        let principal = Principal(inner);
171        let display = format!("{principal}");
172        assert_eq!(display, "ryjl3-tyaaa-aaaaa-aaaba-cai");
173    }
174
175    #[test]
176    fn test_default_is_anonymous() {
177        let principal = Principal::default();
178        assert_eq!(principal.0, candid::Principal::anonymous());
179    }
180
181    #[test]
182    fn test_from_principal_to_value() {
183        let principal = Principal(
184            candid::Principal::from_text("ryjl3-tyaaa-aaaaa-aaaba-cai").expect("invalid principal"),
185        );
186        let value: Value = principal.into();
187        assert!(matches!(value, Value::Custom(_)));
188    }
189
190    #[test]
191    fn test_custom_data_type_tag() {
192        assert_eq!(Principal::TYPE_TAG, "principal");
193    }
194}