1use base64::engine::general_purpose::URL_SAFE_NO_PAD;
10use base64::Engine;
11use serde::{Deserialize, Serialize};
12
13const PREFIX: &str = "bae:";
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct RestoreCode {
18 pub v: u8,
20 pub lid: String,
22 pub ek: String,
24 pub name: String,
26 pub provider: RestoreProvider,
28 pub sk: String,
30}
31
32#[derive(Debug, Clone, Serialize, Deserialize)]
35#[serde(tag = "t")]
36pub enum RestoreProvider {
37 #[serde(rename = "s3")]
38 S3 {
39 bucket: String,
40 region: String,
41 #[serde(default, skip_serializing_if = "Option::is_none")]
42 endpoint: Option<String>,
43 #[serde(default, skip_serializing_if = "Option::is_none")]
44 key_prefix: Option<String>,
45 access_key: String,
46 secret_key: String,
47 },
48 #[serde(rename = "ck")]
49 CloudKit,
50 #[serde(rename = "gd")]
51 GoogleDrive { folder_id: String },
52 #[serde(rename = "db")]
53 Dropbox { folder_path: String },
54 #[serde(rename = "od")]
55 OneDrive { drive_id: String, folder_id: String },
56 #[serde(rename = "bc")]
57 HttpProxy { url: String },
58}
59
60#[derive(Debug, thiserror::Error)]
61pub enum RestoreCodeError {
62 #[error("missing 'bae:' prefix")]
63 MissingPrefix,
64 #[error("invalid base64url encoding")]
65 InvalidBase64,
66 #[error("invalid restore code payload: {0}")]
67 InvalidJson(String),
68 #[error("unsupported version: {0}")]
69 UnsupportedVersion(u8),
70}
71
72pub fn encode_restore_code(code: &RestoreCode) -> String {
74 let json = serde_json::to_string(code).expect("RestoreCode serialization cannot fail");
75 let b64 = URL_SAFE_NO_PAD.encode(json.as_bytes());
76 format!("{PREFIX}{b64}")
77}
78
79pub fn decode_restore_code(s: &str) -> Result<RestoreCode, RestoreCodeError> {
81 let trimmed = s.trim();
82 let payload = trimmed
83 .strip_prefix(PREFIX)
84 .ok_or(RestoreCodeError::MissingPrefix)?;
85 let bytes = URL_SAFE_NO_PAD
86 .decode(payload)
87 .map_err(|_| RestoreCodeError::InvalidBase64)?;
88 let code: RestoreCode =
89 serde_json::from_slice(&bytes).map_err(|e| RestoreCodeError::InvalidJson(e.to_string()))?;
90 if code.v != 1 {
91 return Err(RestoreCodeError::UnsupportedVersion(code.v));
92 }
93 Ok(code)
94}
95
96pub fn provider_needs_oauth(provider: &RestoreProvider) -> bool {
98 matches!(
99 provider,
100 RestoreProvider::GoogleDrive { .. }
101 | RestoreProvider::Dropbox { .. }
102 | RestoreProvider::OneDrive { .. }
103 )
104}
105
106pub struct RestoreCodeInfo {
108 pub library_id: String,
109 pub library_name: String,
110 pub cloud_provider: crate::config::CloudProvider,
111 pub needs_oauth: bool,
112 pub signing_key: Vec<u8>,
114}
115
116pub fn decode_restore_code_info(code: &str) -> Result<RestoreCodeInfo, RestoreCodeError> {
118 let parsed = decode_restore_code(code)?;
119
120 let cloud_provider = match &parsed.provider {
121 RestoreProvider::S3 { .. } => crate::config::CloudProvider::S3,
122 RestoreProvider::CloudKit => crate::config::CloudProvider::CloudKit,
123 RestoreProvider::GoogleDrive { .. } => crate::config::CloudProvider::GoogleDrive,
124 RestoreProvider::Dropbox { .. } => crate::config::CloudProvider::Dropbox,
125 RestoreProvider::OneDrive { .. } => crate::config::CloudProvider::OneDrive,
126 RestoreProvider::HttpProxy { .. } => crate::config::CloudProvider::HttpProxy,
127 };
128
129 let signing_key = URL_SAFE_NO_PAD
130 .decode(&parsed.sk)
131 .map_err(|e| RestoreCodeError::InvalidJson(format!("Invalid signing key encoding: {e}")))?;
132
133 if signing_key.len() != 64 {
134 return Err(RestoreCodeError::InvalidJson(format!(
135 "Signing key must be 64 bytes, got {}",
136 signing_key.len()
137 )));
138 }
139
140 Ok(RestoreCodeInfo {
141 library_id: parsed.lid,
142 library_name: parsed.name,
143 cloud_provider,
144 needs_oauth: provider_needs_oauth(&parsed.provider),
145 signing_key,
146 })
147}
148
149#[cfg(test)]
150mod tests {
151 use super::*;
152
153 fn test_sk() -> String {
154 URL_SAFE_NO_PAD.encode([0xAB_u8; 64])
155 }
156
157 fn sample_s3_code() -> RestoreCode {
158 RestoreCode {
159 v: 1,
160 lid: "550e8400-e29b-41d4-a716-446655440000".to_string(),
161 ek: "aa".repeat(32),
162 name: "Test Library".to_string(),
163 provider: RestoreProvider::S3 {
164 bucket: "my-bucket".to_string(),
165 region: "us-east-1".to_string(),
166 endpoint: Some("https://s3.example.com".to_string()),
167 key_prefix: None,
168 access_key: "AKIAIOSFODNN7EXAMPLE".to_string(),
169 secret_key: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".to_string(),
170 },
171 sk: test_sk(),
172 }
173 }
174
175 #[test]
176 fn roundtrip_s3() {
177 let code = sample_s3_code();
178 let encoded = encode_restore_code(&code);
179 assert!(encoded.starts_with("bae:"));
180
181 let decoded = decode_restore_code(&encoded).unwrap();
182 assert_eq!(decoded.v, 1);
183 assert_eq!(decoded.lid, code.lid);
184 assert_eq!(decoded.ek, code.ek);
185 assert_eq!(decoded.sk, code.sk);
186 assert_eq!(decoded.name, "Test Library");
187 match &decoded.provider {
188 RestoreProvider::S3 {
189 bucket,
190 region,
191 endpoint,
192 key_prefix,
193 access_key,
194 secret_key,
195 } => {
196 assert_eq!(bucket, "my-bucket");
197 assert_eq!(region, "us-east-1");
198 assert_eq!(endpoint.as_deref(), Some("https://s3.example.com"));
199 assert!(key_prefix.is_none());
200 assert_eq!(access_key, "AKIAIOSFODNN7EXAMPLE");
201 assert_eq!(secret_key, "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY");
202 }
203 _ => panic!("expected S3 provider"),
204 }
205 }
206
207 #[test]
208 fn roundtrip_cloudkit() {
209 let code = RestoreCode {
210 v: 1,
211 lid: "lib-123".to_string(),
212 ek: "bb".repeat(32),
213 name: "CloudKit Library".to_string(),
214 provider: RestoreProvider::CloudKit,
215 sk: test_sk(),
216 };
217 let encoded = encode_restore_code(&code);
218 let decoded = decode_restore_code(&encoded).unwrap();
219 assert_eq!(decoded.name, "CloudKit Library");
220 assert!(matches!(decoded.provider, RestoreProvider::CloudKit));
221 }
222
223 #[test]
224 fn roundtrip_google_drive() {
225 let code = RestoreCode {
226 v: 1,
227 lid: "lib-456".to_string(),
228 ek: "cc".repeat(32),
229 name: "GDrive Library".to_string(),
230 provider: RestoreProvider::GoogleDrive {
231 folder_id: "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs".to_string(),
232 },
233 sk: test_sk(),
234 };
235 let decoded = decode_restore_code(&encode_restore_code(&code)).unwrap();
236 match &decoded.provider {
237 RestoreProvider::GoogleDrive { folder_id } => {
238 assert_eq!(folder_id, "1BxiMVs0XRA5nFMdKvBdBZjgmUUqptlbs");
239 }
240 _ => panic!("expected GoogleDrive provider"),
241 }
242 }
243
244 #[test]
245 fn roundtrip_dropbox() {
246 let code = RestoreCode {
247 v: 1,
248 lid: "lib-789".to_string(),
249 ek: "dd".repeat(32),
250 name: "Dropbox Library".to_string(),
251 provider: RestoreProvider::Dropbox {
252 folder_path: "/Apps/bae/My Library".to_string(),
253 },
254 sk: test_sk(),
255 };
256 let decoded = decode_restore_code(&encode_restore_code(&code)).unwrap();
257 match &decoded.provider {
258 RestoreProvider::Dropbox { folder_path } => {
259 assert_eq!(folder_path, "/Apps/bae/My Library");
260 }
261 _ => panic!("expected Dropbox provider"),
262 }
263 }
264
265 #[test]
266 fn roundtrip_onedrive() {
267 let code = RestoreCode {
268 v: 1,
269 lid: "lib-abc".to_string(),
270 ek: "ee".repeat(32),
271 name: "OneDrive Library".to_string(),
272 provider: RestoreProvider::OneDrive {
273 drive_id: "drive-id-123".to_string(),
274 folder_id: "folder-id-456".to_string(),
275 },
276 sk: test_sk(),
277 };
278 let decoded = decode_restore_code(&encode_restore_code(&code)).unwrap();
279 match &decoded.provider {
280 RestoreProvider::OneDrive {
281 drive_id,
282 folder_id,
283 } => {
284 assert_eq!(drive_id, "drive-id-123");
285 assert_eq!(folder_id, "folder-id-456");
286 }
287 _ => panic!("expected OneDrive provider"),
288 }
289 }
290
291 #[test]
292 fn roundtrip_bae_cloud() {
293 let code = RestoreCode {
294 v: 1,
295 lid: "lib-def".to_string(),
296 ek: "ff".repeat(32),
297 name: "Cloud Library".to_string(),
298 provider: RestoreProvider::HttpProxy {
299 url: "https://cloud.bae.fm/lib/abc".to_string(),
300 },
301 sk: test_sk(),
302 };
303 let decoded = decode_restore_code(&encode_restore_code(&code)).unwrap();
304 match &decoded.provider {
305 RestoreProvider::HttpProxy { url } => {
306 assert_eq!(url, "https://cloud.bae.fm/lib/abc");
307 }
308 _ => panic!("expected HttpProxy provider"),
309 }
310 }
311
312 #[test]
313 fn missing_prefix() {
314 let code = sample_s3_code();
315 let encoded = encode_restore_code(&code);
316 let without_prefix = &encoded[4..];
318 assert!(matches!(
319 decode_restore_code(without_prefix),
320 Err(RestoreCodeError::MissingPrefix)
321 ));
322 }
323
324 #[test]
325 fn invalid_base64() {
326 assert!(matches!(
327 decode_restore_code("bae:not-valid!!!"),
328 Err(RestoreCodeError::InvalidBase64)
329 ));
330 }
331
332 #[test]
333 fn invalid_json() {
334 let b64 = URL_SAFE_NO_PAD.encode(b"not json");
335 let code = format!("bae:{b64}");
336 assert!(matches!(
337 decode_restore_code(&code),
338 Err(RestoreCodeError::InvalidJson(_))
339 ));
340 }
341
342 #[test]
343 fn unsupported_version() {
344 let mut code = sample_s3_code();
345 code.v = 99;
346 let encoded = encode_restore_code(&code);
347 assert!(matches!(
348 decode_restore_code(&encoded),
349 Err(RestoreCodeError::UnsupportedVersion(99))
350 ));
351 }
352
353 #[test]
354 fn whitespace_trimmed() {
355 let code = sample_s3_code();
356 let encoded = encode_restore_code(&code);
357 let padded = format!(" {encoded} \n");
358 let decoded = decode_restore_code(&padded).unwrap();
359 assert_eq!(decoded.lid, code.lid);
360 }
361
362 #[test]
363 fn optional_fields_omitted_in_json() {
364 let code = RestoreCode {
365 v: 1,
366 lid: "lib-1".to_string(),
367 ek: "aa".repeat(32),
368 name: "Test Library".to_string(),
369 provider: RestoreProvider::S3 {
370 bucket: "b".to_string(),
371 region: "r".to_string(),
372 endpoint: None,
373 key_prefix: None,
374 access_key: "ak".to_string(),
375 secret_key: "sk-cred".to_string(),
376 },
377 sk: test_sk(),
378 };
379 let json = serde_json::to_string(&code).unwrap();
380 assert!(!json.contains("endpoint"));
382 assert!(!json.contains("key_prefix"));
383 assert!(json.contains("name"));
385 }
386
387 #[test]
388 fn needs_oauth() {
389 assert!(!provider_needs_oauth(&RestoreProvider::S3 {
390 bucket: String::new(),
391 region: String::new(),
392 endpoint: None,
393 key_prefix: None,
394 access_key: String::new(),
395 secret_key: String::new(),
396 }));
397 assert!(!provider_needs_oauth(&RestoreProvider::CloudKit));
398 assert!(provider_needs_oauth(&RestoreProvider::GoogleDrive {
399 folder_id: String::new(),
400 }));
401 assert!(provider_needs_oauth(&RestoreProvider::Dropbox {
402 folder_path: String::new(),
403 }));
404 assert!(provider_needs_oauth(&RestoreProvider::OneDrive {
405 drive_id: String::new(),
406 folder_id: String::new(),
407 }));
408 assert!(!provider_needs_oauth(&RestoreProvider::HttpProxy {
409 url: String::new(),
410 }));
411 }
412}