tenant.openagents/omega

1use anyhow::{Context as _, Result};
2use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation};
3use serde::{Deserialize, Serialize};
4use std::{
5    borrow::Cow,
6    time::{Duration, SystemTime, UNIX_EPOCH},
7};
8
9const DEFAULT_TTL: Duration = Duration::from_secs(6 * 60 * 60); // 6 hours
10
11pub trait UnixTimestampSource: Send + Sync {
12    fn unix_timestamp(&self) -> Result<u64>;
13}
14
15pub struct SystemUnixTimestampSource;
16
17impl UnixTimestampSource for SystemUnixTimestampSource {
18    fn unix_timestamp(&self) -> Result<u64> {
19        Ok(SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs())
20    }
21}
22
23#[derive(Default, Serialize, Deserialize)]
24#[serde(rename_all = "camelCase")]
25pub struct ClaimGrants<'a> {
26    pub iss: Cow<'a, str>,
27    pub sub: Option<Cow<'a, str>>,
28    pub iat: u64,
29    pub exp: u64,
30    pub nbf: u64,
31    pub jwtid: Option<Cow<'a, str>>,
32    pub video: VideoGrant<'a>,
33}
34
35#[derive(Default, Serialize, Deserialize)]
36#[serde(rename_all = "camelCase")]
37pub struct VideoGrant<'a> {
38    pub room_create: Option<bool>,
39    pub room_join: Option<bool>,
40    pub room_list: Option<bool>,
41    pub room_record: Option<bool>,
42    pub room_admin: Option<bool>,
43    pub room: Option<Cow<'a, str>>,
44    pub can_publish: Option<bool>,
45    pub can_subscribe: Option<bool>,
46    pub can_publish_data: Option<bool>,
47    pub hidden: Option<bool>,
48    pub recorder: Option<bool>,
49}
50
51impl<'a> VideoGrant<'a> {
52    pub fn to_admin(room: &'a str) -> Self {
53        Self {
54            room_admin: Some(true),
55            room: Some(Cow::Borrowed(room)),
56            ..Default::default()
57        }
58    }
59
60    pub fn to_join(room: &'a str) -> Self {
61        Self {
62            room: Some(Cow::Borrowed(room)),
63            room_join: Some(true),
64            can_publish: Some(true),
65            can_subscribe: Some(true),
66            ..Default::default()
67        }
68    }
69
70    pub fn for_guest(room: &'a str) -> Self {
71        Self {
72            room: Some(Cow::Borrowed(room)),
73            room_join: Some(true),
74            can_publish: Some(false),
75            can_subscribe: Some(true),
76            ..Default::default()
77        }
78    }
79}
80
81pub fn create(
82    api_key: &str,
83    secret_key: &str,
84    identity: Option<&str>,
85    video_grant: VideoGrant,
86) -> Result<String> {
87    create_with_timestamp_source(
88        api_key,
89        secret_key,
90        identity,
91        video_grant,
92        &SystemUnixTimestampSource,
93    )
94}
95
96pub fn create_with_timestamp_source(
97    api_key: &str,
98    secret_key: &str,
99    identity: Option<&str>,
100    video_grant: VideoGrant,
101    timestamp_source: &dyn UnixTimestampSource,
102) -> Result<String> {
103    let issued_at = timestamp_source.unix_timestamp()?;
104    create_with_issued_at(api_key, secret_key, identity, video_grant, issued_at)
105}
106
107fn create_with_issued_at(
108    api_key: &str,
109    secret_key: &str,
110    identity: Option<&str>,
111    video_grant: VideoGrant,
112    issued_at: u64,
113) -> Result<String> {
114    let room_join = video_grant.room_join.unwrap_or(false);
115    if room_join && identity.is_none() {
116        anyhow::bail!("identity is required for room_join grant, but it is none");
117    }
118
119    let expires_at = issued_at
120        .checked_add(DEFAULT_TTL.as_secs())
121        .context("token expiration overflow")?;
122    let not_before = if room_join {
123        // LiveKit Cloud applies participant revocations by comparing the
124        // revocation timestamp to room-join token `nbf`.
125        issued_at
126    } else {
127        0
128    };
129
130    let claims = ClaimGrants {
131        iss: Cow::Borrowed(api_key),
132        sub: identity.map(Cow::Borrowed),
133        iat: issued_at,
134        exp: expires_at,
135        nbf: not_before,
136        jwtid: identity.map(Cow::Borrowed),
137        video: video_grant,
138    };
139    Ok(jsonwebtoken::encode(
140        &Header::default(),
141        &claims,
142        &EncodingKey::from_secret(secret_key.as_ref()),
143    )?)
144}
145
146pub fn validate<'a>(token: &'a str, secret_key: &str) -> Result<ClaimGrants<'a>> {
147    let token = jsonwebtoken::decode(
148        token,
149        &DecodingKey::from_secret(secret_key.as_ref()),
150        &Validation::default(),
151    )?;
152
153    Ok(token.claims)
154}
155
156#[cfg(any(test, feature = "test-support"))]
157pub fn validate_with_timestamp_source<'a>(
158    token: &'a str,
159    secret_key: &str,
160    timestamp_source: &dyn UnixTimestampSource,
161) -> Result<ClaimGrants<'a>> {
162    let mut validation = Validation::default();
163    validation.validate_exp = false;
164    validation.validate_nbf = false;
165    let token: jsonwebtoken::TokenData<ClaimGrants<'_>> = jsonwebtoken::decode(
166        token,
167        &DecodingKey::from_secret(secret_key.as_ref()),
168        &validation,
169    )?;
170    let claims = token.claims;
171    let timestamp = timestamp_source.unix_timestamp()?;
172
173    anyhow::ensure!(claims.nbf <= timestamp, "token is not yet valid");
174    anyhow::ensure!(claims.exp > timestamp, "token has expired");
175
176    Ok(claims)
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182    use crate::{Client as _, LiveKitClient};
183    use std::sync::Arc;
184
185    const ISSUED_AT: u64 = 1_234_567;
186
187    struct FixedUnixTimestampSource(u64);
188
189    impl UnixTimestampSource for FixedUnixTimestampSource {
190        fn unix_timestamp(&self) -> Result<u64> {
191            Ok(self.0)
192        }
193    }
194
195    #[test]
196    fn token_not_before_matches_issue_time() -> Result<()> {
197        let token = create_with_timestamp_source(
198            "api-key",
199            "secret-key",
200            Some("participant"),
201            VideoGrant::to_join("room"),
202            &FixedUnixTimestampSource(ISSUED_AT),
203        )?;
204
205        assert_claims_timestamp(&token, ISSUED_AT, ISSUED_AT)?;
206        Ok(())
207    }
208
209    #[test]
210    fn room_token_not_before_matches_issue_time() -> Result<()> {
211        let client = LiveKitClient::new_with_timestamp_source(
212            "http://livekit.test".into(),
213            "api-key".into(),
214            "secret-key".into(),
215            Arc::new(FixedUnixTimestampSource(ISSUED_AT)),
216        );
217        let token = client.room_token("room", "participant")?;
218
219        let claims = assert_claims_timestamp(&token, ISSUED_AT, ISSUED_AT)?;
220        assert_eq!(claims.video.room_join, Some(true));
221        assert_eq!(claims.video.can_publish, Some(true));
222        assert_eq!(claims.video.can_subscribe, Some(true));
223
224        Ok(())
225    }
226
227    #[test]
228    fn guest_token_not_before_matches_issue_time() -> Result<()> {
229        let client = LiveKitClient::new_with_timestamp_source(
230            "http://livekit.test".into(),
231            "api-key".into(),
232            "secret-key".into(),
233            Arc::new(FixedUnixTimestampSource(ISSUED_AT)),
234        );
235        let token = client.guest_token("room", "participant")?;
236
237        let claims = assert_claims_timestamp(&token, ISSUED_AT, ISSUED_AT)?;
238        assert_eq!(claims.video.room_join, Some(true));
239        assert_eq!(claims.video.can_publish, Some(false));
240        assert_eq!(claims.video.can_subscribe, Some(true));
241
242        Ok(())
243    }
244
245    #[test]
246    fn admin_token_not_before_remains_unset() -> Result<()> {
247        let token = create_with_timestamp_source(
248            "api-key",
249            "secret-key",
250            None,
251            VideoGrant::to_admin("room"),
252            &FixedUnixTimestampSource(ISSUED_AT),
253        )?;
254
255        let claims = assert_claims_timestamp(&token, ISSUED_AT, 0)?;
256        assert_eq!(claims.video.room_admin, Some(true));
257
258        Ok(())
259    }
260
261    fn assert_claims_timestamp(
262        token: &str,
263        issued_at: u64,
264        expected_not_before: u64,
265    ) -> Result<ClaimGrants<'_>> {
266        let claims = validate_with_timestamp_source(
267            token,
268            "secret-key",
269            &FixedUnixTimestampSource(issued_at),
270        )?;
271
272        assert_eq!(claims.iat, issued_at);
273        assert_eq!(claims.nbf, expected_not_before);
274        assert_eq!(claims.exp, issued_at + DEFAULT_TTL.as_secs());
275
276        Ok(claims)
277    }
278}
279