1use candid::{CandidType, Principal};
2use coset::{
3 cwt::{ClaimName, ClaimsSet, Timestamp},
4 iana, Algorithm, CborSerializable, CoseSign1, CoseSign1Builder, HeaderBuilder,
5};
6use ed25519_dalek::{Signature, VerifyingKey};
7use k256::{ecdsa, ecdsa::signature::hazmat::PrehashVerifier};
8use num_traits::ToPrimitive;
9use serde::{Deserialize, Serialize};
10use serde_bytes::{ByteArray, ByteBuf};
11use sha2::Digest;
12
13pub use coset;
14pub use iana::Algorithm::{EdDSA, ES256K};
15
16const CLOCK_SKEW: i64 = 5 * 60; const ALG_ED25519: Algorithm = Algorithm::Assigned(EdDSA);
18const ALG_SECP256K1: Algorithm = Algorithm::Assigned(ES256K);
19
20static SCOPE_NAME: ClaimName = ClaimName::Assigned(iana::CwtClaimName::Scope);
21
22pub static BUCKET_TOKEN_AAD: &[u8] = b"ic_oss_bucket";
23
24#[derive(CandidType, Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
25pub struct Token {
26 pub subject: Principal,
27 pub audience: Principal,
28 pub policies: String,
29}
30
31impl Token {
32 pub fn from_sign1(
33 sign1_token: &[u8],
34 secp256k1_pub_keys: &[ByteBuf],
35 ed25519_pub_keys: &[ByteArray<32>],
36 aad: &[u8],
37 now_sec: i64,
38 ) -> Result<Self, String> {
39 let cs1 = CoseSign1::from_slice(sign1_token)
40 .map_err(|err| format!("invalid COSE sign1 token: {}", err))?;
41
42 match cs1.protected.header.alg {
43 Some(ALG_SECP256K1) => {
44 Self::secp256k1_verify(secp256k1_pub_keys, &cs1.tbs_data(aad), &cs1.signature)?;
45 }
46 Some(ALG_ED25519) => {
47 Self::ed25519_verify(ed25519_pub_keys, &cs1.tbs_data(aad), &cs1.signature)?;
48 }
49 alg => {
50 Err(format!("unsupported algorithm: {:?}", alg))?;
51 }
52 }
53
54 Self::from_cwt_bytes(&cs1.payload.unwrap_or_default(), now_sec)
55 }
56
57 pub fn to_cwt(self, now_sec: i64, expiration_sec: i64) -> ClaimsSet {
58 ClaimsSet {
59 issuer: None,
60 subject: Some(self.subject.to_text()),
61 audience: Some(self.audience.to_text()),
62 expiration_time: Some(Timestamp::WholeSeconds(now_sec + expiration_sec)),
63 not_before: Some(Timestamp::WholeSeconds(now_sec)),
64 issued_at: Some(Timestamp::WholeSeconds(now_sec)),
65 cwt_id: None,
66 rest: vec![(SCOPE_NAME.clone(), self.policies.into())],
67 }
68 }
69
70 fn secp256k1_verify(
71 pub_keys: &[ByteBuf],
72 tbs_data: &[u8],
73 signature: &[u8],
74 ) -> Result<(), String> {
75 let keys: Vec<ecdsa::VerifyingKey> = pub_keys
76 .iter()
77 .map(|key| {
78 ecdsa::VerifyingKey::from_sec1_bytes(key)
79 .map_err(|_| "invalid verifying key".to_string())
80 })
81 .collect::<Result<_, _>>()?;
82 let sig = ecdsa::Signature::try_from(signature).map_err(|_| "invalid signature")?;
83 let digest = sha256(tbs_data);
84 match keys
85 .iter()
86 .any(|key| key.verify_prehash(digest.as_slice(), &sig).is_ok())
87 {
88 true => Ok(()),
89 false => Err("signature verification failed".to_string()),
90 }
91 }
92
93 fn ed25519_verify(
94 pub_keys: &[ByteArray<32>],
95 tbs_data: &[u8],
96 signature: &[u8],
97 ) -> Result<(), String> {
98 let keys: Vec<VerifyingKey> = pub_keys
99 .iter()
100 .map(|key| {
101 VerifyingKey::from_bytes(key).map_err(|_| "invalid verifying key".to_string())
102 })
103 .collect::<Result<_, _>>()?;
104 let sig = Signature::from_slice(signature).map_err(|_| "invalid signature")?;
105
106 match keys
107 .iter()
108 .any(|key| key.verify_strict(tbs_data, &sig).is_ok())
109 {
110 true => Ok(()),
111 false => Err("signature verification failed".to_string()),
112 }
113 }
114
115 fn from_cwt_bytes(data: &[u8], now_sec: i64) -> Result<Self, String> {
116 let claims =
117 ClaimsSet::from_slice(data).map_err(|err| format!("invalid claims: {}", err))?;
118 if let Some(ref exp) = claims.expiration_time {
119 let exp = match exp {
120 Timestamp::WholeSeconds(v) => *v,
121 Timestamp::FractionalSeconds(v) => (*v).to_i64().unwrap_or_default(),
122 };
123 if exp < now_sec - CLOCK_SKEW {
124 return Err("token expired".to_string());
125 }
126 }
127 if let Some(ref nbf) = claims.not_before {
128 let nbf = match nbf {
129 Timestamp::WholeSeconds(v) => *v,
130 Timestamp::FractionalSeconds(v) => (*v).to_i64().unwrap_or_default(),
131 };
132 if nbf > now_sec + CLOCK_SKEW {
133 return Err("token not yet valid".to_string());
134 }
135 }
136 Self::try_from(claims)
137 }
138}
139
140pub fn cose_sign1(
142 cs: ClaimsSet,
143 alg: iana::Algorithm,
144 key_id: Option<Vec<u8>>,
145) -> Result<CoseSign1, String> {
146 let payload = cs.to_vec().map_err(|err| err.to_string())?;
147 let mut protected = HeaderBuilder::new().algorithm(alg);
148 if let Some(key_id) = key_id {
149 protected = protected.key_id(key_id);
150 }
151
152 Ok(CoseSign1Builder::new()
153 .protected(protected.build())
154 .payload(payload)
155 .build())
156}
157
158impl TryFrom<ClaimsSet> for Token {
159 type Error = String;
160
161 fn try_from(claims: ClaimsSet) -> Result<Self, Self::Error> {
162 let scope = claims
163 .rest
164 .iter()
165 .find(|(key, _)| key == &SCOPE_NAME)
166 .ok_or("missing scope")?;
167 let scope = scope.1.as_text().ok_or("invalid scope text")?;
168
169 Ok(Token {
170 subject: Principal::from_text(claims.subject.as_ref().ok_or("missing subject")?)
171 .map_err(|err| format!("invalid subject: {}", err))?,
172 audience: Principal::from_text(claims.audience.as_ref().ok_or("missing audience")?)
173 .map_err(|err| format!("invalid audience: {}", err))?,
174 policies: scope.to_string(),
175 })
176 }
177}
178
179pub fn sha256(data: &[u8]) -> [u8; 32] {
180 let mut hasher = sha2::Sha256::new();
181 hasher.update(data);
182 hasher.finalize().into()
183}
184
185#[cfg(test)]
186mod test {
187 use super::*;
188 use crate::permission::{Operation, Permission, Policies, Policy, Resource, Resources};
189 use ed25519_dalek::Signer;
190
191 #[test]
192 fn test_ed25519_token() {
193 let secret_key = [8u8; 32];
194 let signing_key = ed25519_dalek::SigningKey::from_bytes(&secret_key);
195 let pub_key: &VerifyingKey = signing_key.as_ref();
196 let pub_key = pub_key.to_bytes();
197 let ps = Policies::from([
198 Policy {
199 permission: Permission {
200 resource: Resource::Bucket,
201 operation: Operation::Read,
202 constraint: Some(Resource::All),
203 },
204 resources: Resources::from([]),
205 },
206 Policy {
207 permission: Permission {
208 resource: Resource::Folder,
209 operation: Operation::All,
210 constraint: None,
211 },
212 resources: Resources::from(["1".to_string()]),
213 },
214 ]);
215 let token = Token {
216 subject: Principal::from_text(
217 "z7wjp-v6fe3-kksu5-26f64-dedtw-j7ndj-57onx-qga6c-et5e3-njx53-tae",
218 )
219 .unwrap(),
220 audience: Principal::from_text("mmrxu-fqaaa-aaaap-ahhna-cai").unwrap(),
221 policies: ps.to_string(),
222 };
223 println!("token: {:?}", &token);
224
225 let now_sec = 1720676064;
226 let claims = token.clone().to_cwt(now_sec, 3600);
227 let mut sign1 = cose_sign1(claims, EdDSA, None).unwrap();
228 let tbs_data = sign1.tbs_data(BUCKET_TOKEN_AAD);
229 let sig = signing_key.sign(&tbs_data).to_bytes();
230 sign1.signature = sig.to_vec();
231 let sign1_token = sign1.to_vec().unwrap();
232 println!("principal: {:?}", &Principal::anonymous().to_text());
233 println!("pub_key: {:?}", &pub_key);
234 println!("sign1_token: {:?}", &sign1_token);
235
236 let token2 = Token::from_sign1(
237 &sign1_token,
238 &[],
239 &[pub_key.into()],
240 BUCKET_TOKEN_AAD,
241 now_sec,
242 )
243 .unwrap();
244 assert_eq!(token, token2);
245 }
246}