1use crate::sqlite_remote::SqliteRemoteRepository;
2use age::secrecy::ExposeSecret;
3use age::x25519;
4use hmac::{Hmac, KeyInit, Mac};
5use kagi_sync::domain::project_token::{ProjectToken, base64_encode_url};
6use sha2::Sha256;
7use std::fs;
8use std::path::Path;
9use std::str::FromStr;
10use std::sync::Arc;
11
12pub type HmacSha256 = Hmac<Sha256>;
13
14#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
15pub struct ServerKeyFile {
16 pub version: u8,
17 pub server_key_id: String,
18 pub age_identity: String,
19 pub token_pepper: String,
20 pub created_at: String,
21}
22
23pub struct AppState {
24 pub repo: SqliteRemoteRepository,
25 pub identity: x25519::Identity,
26 pub server_key_id: String,
27 pub fingerprint: String,
28 pub token_pepper: Vec<u8>,
29}
30
31impl AppState {
32 pub async fn new(db_path: &Path, key_file_path: &Path) -> Result<Arc<Self>, anyhow::Error> {
33 let db_path = if db_path.is_absolute() {
34 db_path.to_path_buf()
35 } else {
36 std::env::current_dir()?.join(db_path)
37 };
38 if let Some(parent) = db_path.parent() {
39 fs::create_dir_all(parent)?;
40 }
41 let repo = SqliteRemoteRepository::new_file(&db_path).await?;
42
43 let key_file = if key_file_path.exists() {
44 let content = fs::read_to_string(key_file_path)?;
45 serde_json::from_str::<ServerKeyFile>(&content)?
46 } else {
47 let identity = x25519::Identity::generate();
48 let pepper: Vec<u8> = (0..32).map(|_| rand::random::<u8>()).collect();
49 let server_key_id = format!(r"kgs_{}", nanoid::nanoid!(12));
50 let key_file = ServerKeyFile {
51 version: 1,
52 server_key_id: server_key_id.clone(),
53 age_identity: identity.to_string().expose_secret().to_string(),
54 token_pepper: base64_encode_url(&pepper),
55 created_at: time::OffsetDateTime::now_utc().to_string(),
56 };
57 if let Some(parent) = key_file_path.parent() {
58 fs::create_dir_all(parent)?;
59 }
60 fs::write(key_file_path, serde_json::to_string_pretty(&key_file)?)?;
61 tracing::info!(
62 "kagi: generated new server key file at {}",
63 key_file_path.display()
64 );
65 #[cfg(unix)]
66 {
67 use std::os::unix::fs::PermissionsExt;
68 fs::set_permissions(key_file_path, fs::Permissions::from_mode(0o600))?;
69 }
70 key_file
71 };
72
73 let identity = x25519::Identity::from_str(&key_file.age_identity)
74 .map_err(|e| anyhow::anyhow!("invalid server identity: {e}"))?;
75 let fingerprint = key_file.server_key_id.clone();
76 let token_pepper = base64_decode_url(&key_file.token_pepper)
77 .map_err(|e| anyhow::anyhow!("invalid token pepper: {e}"))?;
78
79 let has_admin = repo
80 .has_admin_token()
81 .await
82 .map_err(|e| anyhow::anyhow!("failed to check admin token: {e}"))?;
83 if !has_admin {
84 let admin_token = ProjectToken::generate_admin_token(fingerprint.clone());
85 let token_hash = {
86 let mut mac =
87 HmacSha256::new_from_slice(&token_pepper).expect("HMAC key size valid");
88 mac.update(admin_token.full_token.as_bytes());
89 let result = mac.finalize();
90 let hash = result.into_bytes();
91 format!("kh1:{}", base64_encode_url(&hash))
92 };
93 let caps_json = serde_json::to_string(&admin_token.payload.capabilities)
94 .map_err(|e| anyhow::anyhow!("failed to serialize capabilities: {e}"))?;
95 let now = time::OffsetDateTime::now_utc().to_string();
96 repo.create_admin_token(&admin_token.payload.token_id, &token_hash, &caps_json, &now)
97 .await
98 .map_err(|e| anyhow::anyhow!("failed to store admin token: {e}"))?;
99 println!("kagi: generated admin token: {}", admin_token.full_token);
100 println!("kagi: store this in KAGI_ADMIN_TOKEN env var for admin operations");
101 }
102
103 Ok(Arc::new(Self {
104 repo,
105 identity,
106 server_key_id: key_file.server_key_id,
107 fingerprint,
108 token_pepper,
109 }))
110 }
111
112 pub fn hash_token(&self, full_token: &str) -> String {
113 let mut mac = HmacSha256::new_from_slice(&self.token_pepper).expect("HMAC key size valid");
114 mac.update(full_token.as_bytes());
115 let result = mac.finalize();
116 let hash = result.into_bytes();
117 format!("kh1:{}", base64_encode_url(&hash))
118 }
119}
120
121pub fn hash_claim_secret(claim_secret: &str) -> String {
122 use sha2::{Digest, Sha256};
123 let mut hasher = Sha256::new();
124 hasher.update(claim_secret.as_bytes());
125 format!("cs:{}", base64_encode_url(&hasher.finalize()))
126}
127
128fn base64_decode_url(input: &str) -> Result<Vec<u8>, base64::DecodeError> {
129 use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
130 URL_SAFE_NO_PAD.decode(input)
131}
132
133#[cfg(test)]
134mod tests {
135 use super::*;
136 use crate::sqlite_remote::SqliteRemoteRepository;
137
138 async fn test_repo() -> SqliteRemoteRepository {
139 let id = rand::random::<u64>();
140 let path = std::env::temp_dir().join(format!("kagi_state_test_{id}.db"));
141 SqliteRemoteRepository::new_file(path).await.unwrap()
142 }
143
144 #[tokio::test]
145 async fn test_hash_token_deterministic() {
146 let repo = test_repo().await;
147 let pepper = vec![
148 1u8, 2u8, 3u8, 4u8, 5u8, 6u8, 7u8, 8u8, 9u8, 10u8, 11u8, 12u8, 13u8, 14u8, 15u8, 16u8,
149 17u8, 18u8, 19u8, 20u8, 21u8, 22u8, 23u8, 24u8, 25u8, 26u8, 27u8, 28u8, 29u8, 30u8,
150 31u8, 32u8,
151 ];
152 let state = AppState {
153 repo,
154 identity: x25519::Identity::generate(),
155 server_key_id: "kgs_test".into(),
156 fingerprint: "fp_test".into(),
157 token_pepper: pepper.clone(),
158 };
159
160 let hash1 = state.hash_token("my_secret_token");
161 let hash2 = state.hash_token("my_secret_token");
162 assert_eq!(hash1, hash2);
163 assert!(hash1.starts_with("kh1:"));
164
165 let hash3 = state.hash_token("different_token");
166 assert_ne!(hash1, hash3);
167 }
168
169 #[tokio::test]
170 async fn test_hash_token_different_pepper() {
171 let repo1 = test_repo().await;
172 let repo2 = test_repo().await;
173 let pepper1 = vec![0u8; 32];
174 let state1 = AppState {
175 repo: repo1,
176 identity: x25519::Identity::generate(),
177 server_key_id: "kgs_test".into(),
178 fingerprint: "fp_test".into(),
179 token_pepper: pepper1,
180 };
181
182 let pepper2 = vec![1u8; 32];
183 let state2 = AppState {
184 repo: repo2,
185 identity: x25519::Identity::generate(),
186 server_key_id: "kgs_test".into(),
187 fingerprint: "fp_test".into(),
188 token_pepper: pepper2,
189 };
190
191 let hash1 = state1.hash_token("same_token");
192 let hash2 = state2.hash_token("same_token");
193 assert_ne!(hash1, hash2);
194 }
195
196 #[test]
197 fn test_base64_decode_url_roundtrip() {
198 let data = b"hello world";
199 let encoded = base64_encode_url(data);
200 let decoded = base64_decode_url(&encoded).unwrap();
201 assert_eq!(decoded, data);
202 }
203
204 #[test]
205 fn test_base64_decode_url_invalid() {
206 assert!(base64_decode_url("!!!").is_err());
207 }
208}