/
lib.rs
139 lines (116 loc) · 3.79 KB
/
lib.rs
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
mod delegation;
mod hash;
mod id_token;
mod state;
mod types;
mod users;
mod utils;
use candid::Principal;
use ic_cdk::*;
use ic_cdk_timers::set_timer;
use ic_stable_structures::{
memory_manager::{MemoryId, MemoryManager, VirtualMemory},
storable::Blob,
DefaultMemoryImpl, StableBTreeMap,
};
use id_token::IdToken;
use jsonwebtoken_rustcrypto::Algorithm;
use serde_bytes::ByteBuf;
use std::{cell::RefCell, time::Duration};
use crate::{
state::State,
types::{
AuthenticatedResponse, GetDelegationResponse, PrepareDelegationResponse, SessionKey,
Timestamp, UserSub,
},
};
thread_local! {
static MEMORY_MANAGER: RefCell<MemoryManager<DefaultMemoryImpl>> =
RefCell::new(MemoryManager::init(DefaultMemoryImpl::default()));
}
thread_local! {
/* flexible */ static STATE: RefCell<State> = RefCell::new(State::default());
/* stable */static PRINCIPAL_USER_SUB: RefCell<StableBTreeMap<Blob<29>, UserSub, VirtualMemory<DefaultMemoryImpl>>> = RefCell::new(
StableBTreeMap::init(
MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(0))),
)
);
}
#[init]
fn init() {
set_timer(Duration::ZERO, || {
spawn(state::init());
});
}
#[post_upgrade]
fn post_upgrade() {
set_timer(Duration::ZERO, || {
spawn(state::init());
});
}
fn check_authorization(caller: Principal, jwt: String) -> Result<(IdToken, SessionKey), String> {
let token = id_token::decode(&jwt, Algorithm::RS256).map_err(|e| format!("{:?}", e))?;
id_token::validate(&token.claims).map_err(|e| format!("{:?}", e))?;
let nonce = {
let nonce = hex::decode(&token.claims.nonce).map_err(|e| format!("{:?}", e))?;
ByteBuf::from(nonce)
};
let token_principal = Principal::self_authenticating(&nonce);
if caller != token_principal {
return Err("caller and token principal mismatch".to_string());
}
Ok((token, nonce))
}
#[update]
async fn prepare_delegation(jwt: String) -> PrepareDelegationResponse {
let session_principal = caller();
let (token, session_key) = match check_authorization(session_principal, jwt) {
Ok(res) => res,
Err(e) => {
trap(&e);
}
};
let sub = token.claims.clone().sub;
let expiration = id_token::expiration_timestamp_ns(&token.claims);
let user_key = delegation::prepare_delegation(&sub, session_key, expiration).await;
let principal = delegation::get_principal(&sub);
users::register_user(principal, sub);
PrepareDelegationResponse {
user_key,
expiration,
}
}
#[query]
fn get_delegation(jwt: String, expiration: Timestamp) -> GetDelegationResponse {
let session_principal = caller();
let (token, session_key) = match check_authorization(session_principal, jwt) {
Ok(res) => res,
Err(e) => {
trap(&e);
}
};
let sub = &token.claims.sub;
delegation::get_delegation(sub, session_key, expiration)
}
#[query]
fn authenticated() -> AuthenticatedResponse {
let caller = caller();
match users::get_user_sub(caller) {
Some(sub) => {
print(format!("sub: {} principal: {}", sub, caller.to_text(),));
AuthenticatedResponse {
user_sub: sub,
user_principal: caller,
}
}
None => trap("No user found"),
}
}
// In the following, we register a custom getrandom implementation because
// otherwise getrandom (which is a dependency of some packages) fails to compile.
// This is necessary because getrandom by default fails to compile for the
// wasm32-unknown-unknown target (which is required for deploying a canister).
getrandom::register_custom_getrandom!(always_fail);
pub fn always_fail(_buf: &mut [u8]) -> Result<(), getrandom::Error> {
Err(getrandom::Error::UNSUPPORTED)
}