pie 0.2.0

Pie: A Programmable LLM Serving System
Documentation 
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

use anyhow::{Result, anyhow};
use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation, decode, encode};
use serde::{Deserialize, Serialize};
use std::sync::OnceLock;

#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
pub enum Role {
    User,
    Admin,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Claims {
    pub sub: String, // Subject (e.g., user or service account ID)
    pub role: Role,
    pub exp: usize, // Expiration timestamp
}

static SECRET: OnceLock<String> = OnceLock::new();
pub fn init_secret(secret: &str) {
    // This will only set the secret once. Subsequent calls have no effect.
    let _ = SECRET.set(secret.to_string());
}

fn get_secret() -> Result<&'static str> {
    SECRET
        .get()
        .map(|s| s.as_str())
        .ok_or_else(|| anyhow!("JWT secret not initialized"))
}

pub fn create_jwt(user_id: &str, role: Role) -> Result<String> {
    let expiration = chrono::Utc::now()
        .checked_add_signed(chrono::Duration::days(365))
        .expect("valid timestamp")
        .timestamp();

    let claims = Claims {
        sub: user_id.to_owned(),
        role,
        exp: expiration as usize,
    };

    let secret = get_secret()?;

    encode(
        &Header::default(),
        &claims,
        &EncodingKey::from_secret(secret.as_bytes()),
    )
    .map_err(Into::into)
}

pub fn validate_jwt(token: &str) -> Result<Claims> {
    let secret = get_secret()?;

    decode::<Claims>(
        token,
        &DecodingKey::from_secret(secret.as_bytes()),
        &Validation::default(),
    )
    .map(|data| data.claims)
    .map_err(Into::into)
}