squigit_auth/security/
crypto.rs

1// Copyright 2026 a7mddra
2// SPDX-License-Identifier: Apache-2.0
3
4use aes_gcm::{
5    aead::{Aead, KeyInit, Payload},
6    Aes256Gcm, Nonce,
7};
8use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
9use hkdf::Hkdf;
10use hmac::{Hmac, Mac};
11use rand::{rngs::OsRng, RngCore};
12use sha2::Sha256;
13use squigit_storage::{EncryptedKeyRecord, ProfileStore, RecordCipher, RecordKdf};
14use std::sync::{OnceLock, RwLock};
15use zeroize::{Zeroize, Zeroizing};
16
17use crate::{ByokErrorCode, ProfileError, Result};
18
19use super::vault::{
20    OsSecretVault, SecretVault, VaultKey, CAS_BINDING_KEY_ACCOUNT, RECORD_ENCRYPTION_MASTER_ACCOUNT,
21};
22use super::{validate_api_key, ApiKeyProvider};
23
24const RECORD_KEY_DOMAIN: &str = "squigit/byok/v1/record-key";
25const RECORD_AAD_DOMAIN: &str = "squigit/byok/v1/record-aad";
26const RUNTIME_CREDENTIAL_DOMAIN: &str = "squigit/cas/v1/runtime-credential";
27const OBJECT_REMOTE_DOMAIN: &str = "squigit/cas/v1/object-remote";
28const SESSION_CREDENTIAL_DOMAIN: &str = "squigit/session/v1/runtime-credential";
29const SESSION_OBJECT_REMOTE_DOMAIN: &str = "squigit/session/v1/object-remote";
30const AES_256_GCM: &str = "aes-256-gcm";
31const HKDF_SHA256: &str = "hkdf-sha256";
32
33type HmacSha256 = Hmac<Sha256>;
34
35#[derive(Default)]
36struct SessionApiKeys {
37    active: bool,
38    google_ai_studio: Option<SecretString>,
39    imgbb: Option<SecretString>,
40}
41
42impl SessionApiKeys {
43    fn get(&self, provider: ApiKeyProvider) -> Option<&SecretString> {
44        match provider {
45            ApiKeyProvider::GoogleAiStudio => self.google_ai_studio.as_ref(),
46            ApiKeyProvider::ImgBb => self.imgbb.as_ref(),
47        }
48    }
49}
50
51static SESSION_API_KEYS: OnceLock<RwLock<SessionApiKeys>> = OnceLock::new();
52
53fn session_api_keys() -> &'static RwLock<SessionApiKeys> {
54    SESSION_API_KEYS.get_or_init(|| RwLock::new(SessionApiKeys::default()))
55}
56
57pub struct SecretString(Zeroizing<String>);
58
59impl std::fmt::Debug for SecretString {
60    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
61        formatter.write_str("SecretString([REDACTED])")
62    }
63}
64
65impl SecretString {
66    pub fn new(value: String) -> Self {
67        Self(Zeroizing::new(value))
68    }
69
70    pub fn expose(&self) -> &str {
71        self.0.as_str()
72    }
73
74    pub fn into_inner(mut self) -> String {
75        std::mem::take(&mut *self.0)
76    }
77}
78
79pub struct CredentialDigest(Zeroizing<[u8; 32]>);
80
81impl std::fmt::Debug for CredentialDigest {
82    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        formatter.write_str("CredentialDigest([REDACTED])")
84    }
85}
86
87impl CredentialDigest {
88    pub fn matches(&self, other: &Self) -> bool {
89        let Ok(mut mac) = <HmacSha256 as Mac>::new_from_slice(&self.0[..]) else {
90            return false;
91        };
92        mac.update(b"credential-digest-comparison");
93        mac.verify_slice(
94            &<HmacSha256 as Mac>::new_from_slice(&other.0[..])
95                .expect("HMAC-SHA256 accepts a 32-byte key")
96                .chain_update(b"credential-digest-comparison")
97                .finalize()
98                .into_bytes(),
99        )
100        .is_ok()
101    }
102}
103
104pub struct DecryptedApiKey {
105    pub api_key: SecretString,
106    pub runtime_digest: CredentialDigest,
107}
108
109impl std::fmt::Debug for DecryptedApiKey {
110    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
111        formatter
112            .debug_struct("DecryptedApiKey")
113            .field("api_key", &"[REDACTED]")
114            .field("runtime_digest", &"[REDACTED]")
115            .finish()
116    }
117}
118
119fn frame(fields: &[&str]) -> Vec<u8> {
120    let capacity = fields
121        .iter()
122        .map(|field| std::mem::size_of::<u32>() + field.len())
123        .sum();
124    let mut framed = Vec::with_capacity(capacity);
125    for field in fields {
126        let bytes = field.as_bytes();
127        let length = u32::try_from(bytes.len()).expect("BYOK frame field exceeds u32");
128        framed.extend_from_slice(&length.to_be_bytes());
129        framed.extend_from_slice(bytes);
130    }
131    framed
132}
133
134fn canonicalize_api_key(provider: ApiKeyProvider, plaintext: &str) -> Result<SecretString> {
135    let canonical = plaintext.trim();
136    if canonical.is_empty() {
137        return Err(ProfileError::byok(
138            ByokErrorCode::InvalidCredential,
139            "An empty API key cannot be saved. Use explicit deletion instead.",
140        ));
141    }
142    validate_api_key(provider, canonical)?;
143    Ok(SecretString::new(canonical.to_owned()))
144}
145
146/// Replace the process-only API keys used by developer shells.
147///
148/// These credentials are validated and zeroized in memory. They are never
149/// written to the profile key store or the operating-system vault.
150pub fn set_session_api_keys(google_ai_studio: Option<&str>, imgbb: Option<&str>) -> Result<()> {
151    let google_ai_studio = google_ai_studio
152        .map(|value| canonicalize_api_key(ApiKeyProvider::GoogleAiStudio, value))
153        .transpose()?;
154    let imgbb = imgbb
155        .map(|value| canonicalize_api_key(ApiKeyProvider::ImgBb, value))
156        .transpose()?;
157    let mut keys = session_api_keys()
158        .write()
159        .map_err(|_| ProfileError::Auth("Process-only API-key state is unavailable.".into()))?;
160    *keys = SessionApiKeys {
161        active: true,
162        google_ai_studio,
163        imgbb,
164    };
165    Ok(())
166}
167
168/// Return whether process-only credentials replace the persistent key store.
169pub fn session_api_keys_active() -> bool {
170    session_api_keys().read().is_ok_and(|keys| keys.active)
171}
172
173/// Return the character width of a process-only credential, when configured.
174pub fn session_api_key_width(provider: ApiKeyProvider) -> Option<u32> {
175    session_api_keys().read().ok().and_then(|keys| {
176        keys.get(provider)
177            .map(|key| key.expose().chars().count() as u32)
178    })
179}
180
181fn session_api_key(provider: ApiKeyProvider) -> Option<SecretString> {
182    session_api_keys().read().ok().and_then(|keys| {
183        keys.get(provider)
184            .map(|key| SecretString::new(key.expose().to_owned()))
185    })
186}
187
188fn random_vault_key() -> [u8; 32] {
189    let mut bytes = [0u8; 32];
190    OsRng.fill_bytes(&mut bytes);
191    bytes
192}
193
194fn create_and_verify_vault_key<V: SecretVault>(vault: &V, account: &str) -> Result<VaultKey> {
195    let mut generated = Zeroizing::new(random_vault_key());
196    vault.set(account, &generated)?;
197    let read_back = match vault.get(account) {
198        Ok(Some(secret)) => secret,
199        Ok(None) => {
200            let _ = vault.delete(account);
201            return Err(ProfileError::byok(
202                ByokErrorCode::VaultUnavailable,
203                "The OS vault did not return the secret it just stored.",
204            ));
205        }
206        Err(error) => {
207            let _ = vault.delete(account);
208            return Err(error);
209        }
210    };
211    if generated.as_ref() != read_back.expose() {
212        let _ = vault.delete(account);
213        return Err(ProfileError::byok(
214            ByokErrorCode::VaultUnavailable,
215            "The OS vault did not preserve the stored secret.",
216        ));
217    }
218    generated.zeroize();
219    Ok(read_back)
220}
221
222fn required_vault_key<V: SecretVault>(
223    vault: &V,
224    account: &str,
225    missing_code: ByokErrorCode,
226) -> Result<VaultKey> {
227    vault.get(account)?.ok_or_else(|| {
228        ProfileError::byok(
229            missing_code,
230            format!("Required OS-vault entry '{account}' is missing."),
231        )
232    })
233}
234
235fn derive_record_key(
236    master: &VaultKey,
237    salt: &[u8],
238    profile_id: &str,
239    provider: ApiKeyProvider,
240) -> Result<Zeroizing<[u8; 32]>> {
241    let info = frame(&[RECORD_KEY_DOMAIN, profile_id, provider.storage_key_name()]);
242    let hkdf = Hkdf::<Sha256>::new(Some(salt), master.expose());
243    let mut key = Zeroizing::new([0u8; 32]);
244    hkdf.expand(&info, &mut *key).map_err(|_| {
245        ProfileError::byok(
246            ByokErrorCode::EncryptionFailed,
247            "Failed to derive an API-key record encryption key.",
248        )
249    })?;
250    Ok(key)
251}
252
253fn record_aad(profile_id: &str, provider: ApiKeyProvider) -> Vec<u8> {
254    frame(&[
255        RECORD_AAD_DOMAIN,
256        profile_id,
257        provider.storage_key_name(),
258        AES_256_GCM,
259        HKDF_SHA256,
260    ])
261}
262
263fn encrypt_record(
264    master: &VaultKey,
265    profile_id: &str,
266    provider: ApiKeyProvider,
267    plaintext: &SecretString,
268) -> Result<EncryptedKeyRecord> {
269    let mut salt = [0u8; 32];
270    let mut nonce_bytes = [0u8; 12];
271    OsRng.fill_bytes(&mut salt);
272    OsRng.fill_bytes(&mut nonce_bytes);
273
274    let record_key = derive_record_key(master, &salt, profile_id, provider)?;
275    let cipher =
276        Aes256Gcm::new_from_slice(&record_key[..]).expect("HKDF-SHA256 produces an AES-256 key");
277    let nonce = Nonce::from(nonce_bytes);
278    let ciphertext = cipher
279        .encrypt(
280            &nonce,
281            Payload {
282                msg: plaintext.expose().as_bytes(),
283                aad: &record_aad(profile_id, provider),
284            },
285        )
286        .map_err(|_| {
287            ProfileError::byok(
288                ByokErrorCode::EncryptionFailed,
289                "Failed to encrypt the API-key record.",
290            )
291        })?;
292
293    Ok(EncryptedKeyRecord {
294        cipher: RecordCipher::Aes256Gcm,
295        width: plaintext.expose().chars().count() as u32,
296        kdf: RecordKdf::HkdfSha256,
297        salt: URL_SAFE_NO_PAD.encode(salt),
298        nonce: URL_SAFE_NO_PAD.encode(nonce_bytes),
299        ciphertext: URL_SAFE_NO_PAD.encode(ciphertext),
300    })
301}
302
303fn decode_record_field(encoded: &str, expected_length: Option<usize>) -> Result<Vec<u8>> {
304    let decoded = URL_SAFE_NO_PAD.decode(encoded).map_err(|_| {
305        ProfileError::byok(
306            ByokErrorCode::MalformedKeyStore,
307            "The encrypted API-key store contains invalid base64url.",
308        )
309    })?;
310    if expected_length.is_some_and(|length| decoded.len() != length)
311        || URL_SAFE_NO_PAD.encode(&decoded) != encoded
312    {
313        return Err(ProfileError::byok(
314            ByokErrorCode::MalformedKeyStore,
315            "The encrypted API-key store contains non-canonical data.",
316        ));
317    }
318    Ok(decoded)
319}
320
321fn decrypt_record(
322    master: &VaultKey,
323    profile_id: &str,
324    provider: ApiKeyProvider,
325    record: &EncryptedKeyRecord,
326) -> Result<SecretString> {
327    let salt = decode_record_field(&record.salt, Some(32))?;
328    let nonce = decode_record_field(&record.nonce, Some(12))?;
329    let mut ciphertext = decode_record_field(&record.ciphertext, None)?;
330    if ciphertext.len() < 16 {
331        ciphertext.zeroize();
332        return Err(ProfileError::byok(
333            ByokErrorCode::MalformedKeyStore,
334            "The encrypted API-key record does not contain an authentication tag.",
335        ));
336    }
337
338    let record_key = derive_record_key(master, &salt, profile_id, provider)?;
339    let cipher =
340        Aes256Gcm::new_from_slice(&record_key[..]).expect("HKDF-SHA256 produces an AES-256 key");
341    let nonce = Nonce::from(
342        <[u8; 12]>::try_from(nonce.as_slice())
343            .expect("encrypted record nonces are validated as 12 bytes"),
344    );
345    let plaintext = cipher
346        .decrypt(
347            &nonce,
348            Payload {
349                msg: &ciphertext,
350                aad: &record_aad(profile_id, provider),
351            },
352        )
353        .map_err(|_| {
354            ProfileError::byok(
355                ByokErrorCode::MalformedKeyStore,
356                "The encrypted API-key record failed authentication.",
357            )
358        })?;
359    ciphertext.zeroize();
360    let plaintext = String::from_utf8(plaintext).map_err(|error| {
361        let mut bytes = error.into_bytes();
362        bytes.zeroize();
363        ProfileError::byok(
364            ByokErrorCode::MalformedKeyStore,
365            "The decrypted API-key record is not valid UTF-8.",
366        )
367    })?;
368    let secret = SecretString::new(plaintext);
369    validate_api_key(provider, secret.expose()).map_err(|_| {
370        ProfileError::byok(
371            ByokErrorCode::MalformedKeyStore,
372            "The decrypted API-key record has an invalid provider format.",
373        )
374    })?;
375    Ok(secret)
376}
377
378fn hmac_digest(key: &VaultKey, fields: &[&str]) -> CredentialDigest {
379    let mut mac = <HmacSha256 as Mac>::new_from_slice(key.expose())
380        .expect("HMAC-SHA256 accepts a 32-byte key");
381    mac.update(&frame(fields));
382    let mut digest = [0u8; 32];
383    digest.copy_from_slice(&mac.finalize().into_bytes());
384    CredentialDigest(Zeroizing::new(digest))
385}
386
387fn runtime_digest(
388    binding_key: &VaultKey,
389    provider: ApiKeyProvider,
390    api_key: &SecretString,
391) -> CredentialDigest {
392    hmac_digest(
393        binding_key,
394        &[
395            RUNTIME_CREDENTIAL_DOMAIN,
396            provider.storage_key_name(),
397            api_key.expose(),
398        ],
399    )
400}
401
402fn session_runtime_digest(provider: ApiKeyProvider, api_key: &SecretString) -> CredentialDigest {
403    let mut mac = <HmacSha256 as Mac>::new_from_slice(SESSION_CREDENTIAL_DOMAIN.as_bytes())
404        .expect("HMAC-SHA256 accepts the session credential domain");
405    mac.update(&frame(&[provider.storage_key_name(), api_key.expose()]));
406    let mut digest = [0u8; 32];
407    digest.copy_from_slice(&mac.finalize().into_bytes());
408    CredentialDigest(Zeroizing::new(digest))
409}
410
411fn is_session_api_key(provider: ApiKeyProvider, api_key: &SecretString) -> bool {
412    session_api_key(provider).is_some_and(|session_key| {
413        session_runtime_digest(provider, &session_key)
414            .matches(&session_runtime_digest(provider, api_key))
415    })
416}
417
418fn session_object_remote_id(
419    provider: ApiKeyProvider,
420    lowercase_object_hash: &str,
421    api_key: &SecretString,
422) -> String {
423    let mut mac = <HmacSha256 as Mac>::new_from_slice(api_key.expose().as_bytes())
424        .expect("HMAC-SHA256 accepts API keys of any length");
425    mac.update(&frame(&[
426        SESSION_OBJECT_REMOTE_DOMAIN,
427        provider.storage_key_name(),
428        lowercase_object_hash,
429    ]));
430    hex::encode(mac.finalize().into_bytes())
431}
432
433pub fn object_remote_id(
434    provider: ApiKeyProvider,
435    lowercase_object_hash: &str,
436    api_key: &SecretString,
437) -> Result<String> {
438    if lowercase_object_hash.len() != 64
439        || !lowercase_object_hash
440            .bytes()
441            .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
442    {
443        return Err(ProfileError::byok(
444            ByokErrorCode::MalformedKeyStore,
445            "The CAS object hash is not canonical lowercase hexadecimal.",
446        ));
447    }
448    if is_session_api_key(provider, api_key) {
449        return Ok(session_object_remote_id(
450            provider,
451            lowercase_object_hash,
452            api_key,
453        ));
454    }
455    let vault = OsSecretVault;
456    let binding_key = required_vault_key(
457        &vault,
458        CAS_BINDING_KEY_ACCOUNT,
459        ByokErrorCode::CasBindingKeyMissing,
460    )?;
461    let digest = hmac_digest(
462        &binding_key,
463        &[
464            OBJECT_REMOTE_DOMAIN,
465            provider.storage_key_name(),
466            lowercase_object_hash,
467            api_key.expose(),
468        ],
469    );
470    Ok(hex::encode(&digest.0[..]))
471}
472
473pub fn get_decrypted_api_key(
474    store: &ProfileStore,
475    provider: ApiKeyProvider,
476    profile_id: &str,
477) -> Result<Option<DecryptedApiKey>> {
478    if session_api_keys_active() {
479        return Ok(session_api_key(provider).map(|api_key| {
480            let runtime_digest = session_runtime_digest(provider, &api_key);
481            DecryptedApiKey {
482                api_key,
483                runtime_digest,
484            }
485        }));
486    }
487    get_decrypted_api_key_with_vault(store, provider, profile_id, &OsSecretVault)
488}
489
490fn get_decrypted_api_key_with_vault<V: SecretVault>(
491    store: &ProfileStore,
492    provider: ApiKeyProvider,
493    profile_id: &str,
494    vault: &V,
495) -> Result<Option<DecryptedApiKey>> {
496    let Some(record) = store.load_encrypted_key_record(profile_id, provider.storage_key_name())?
497    else {
498        return Ok(None);
499    };
500    let master = required_vault_key(
501        vault,
502        RECORD_ENCRYPTION_MASTER_ACCOUNT,
503        ByokErrorCode::MasterKeyMissing,
504    )?;
505    let binding_key = required_vault_key(
506        vault,
507        CAS_BINDING_KEY_ACCOUNT,
508        ByokErrorCode::CasBindingKeyMissing,
509    )?;
510    let api_key = decrypt_record(&master, profile_id, provider, &record)?;
511    let runtime_digest = runtime_digest(&binding_key, provider, &api_key);
512    Ok(Some(DecryptedApiKey {
513        api_key,
514        runtime_digest,
515    }))
516}
517
518pub fn reveal_api_key(
519    store: &ProfileStore,
520    provider: ApiKeyProvider,
521    profile_id: &str,
522) -> Result<Option<SecretString>> {
523    let secret =
524        get_decrypted_api_key(store, provider, profile_id)?.map(|credential| credential.api_key);
525    Ok(secret)
526}
527
528pub fn get_api_key_status(
529    store: &ProfileStore,
530    provider: ApiKeyProvider,
531    profile_id: &str,
532) -> Result<bool> {
533    if session_api_keys_active() {
534        return Ok(session_api_key_width(provider).is_some());
535    }
536    Ok(store
537        .load_encrypted_key_record(profile_id, provider.storage_key_name())?
538        .is_some())
539}
540
541pub fn encrypt_and_save_api_key(
542    store: &ProfileStore,
543    profile_id: &str,
544    provider: ApiKeyProvider,
545    plaintext: &str,
546) -> Result<()> {
547    encrypt_and_save_api_key_with_vault(store, profile_id, provider, plaintext, &OsSecretVault)
548}
549
550fn encrypt_and_save_api_key_with_vault<V: SecretVault>(
551    store: &ProfileStore,
552    profile_id: &str,
553    provider: ApiKeyProvider,
554    plaintext: &str,
555    vault: &V,
556) -> Result<()> {
557    if store.get_profile(profile_id)?.is_none() {
558        return Err(ProfileError::ProfileNotFound(profile_id.to_owned()));
559    }
560    let plaintext = canonicalize_api_key(provider, plaintext)?;
561
562    store.with_key_store_transaction(|transaction| {
563        let mut keys = transaction.load()?;
564        let store_was_empty = keys.profiles.is_empty();
565        let mut created_master = false;
566        let mut created_binding = false;
567
568        let result = (|| {
569            let master = match vault.get(RECORD_ENCRYPTION_MASTER_ACCOUNT)? {
570                Some(key) => key,
571                None if store_was_empty => {
572                    created_master = true;
573                    create_and_verify_vault_key(vault, RECORD_ENCRYPTION_MASTER_ACCOUNT)?
574                }
575                None => {
576                    return Err(ProfileError::byok(
577                        ByokErrorCode::MasterKeyMissing,
578                        "A populated key store is missing its OS-vault encryption master.",
579                    ));
580                }
581            };
582
583            let _binding_key = match vault.get(CAS_BINDING_KEY_ACCOUNT)? {
584                Some(key) => key,
585                None => {
586                    if squigit_storage::ThreadStorage::with_config_root(store.base_dir().clone())?
587                        .has_object_remotes()?
588                    {
589                        return Err(ProfileError::byok(
590                            ByokErrorCode::CasBindingKeyMissing,
591                            "CAS remotes exist but their OS-vault binding key is missing. Restore the vault entry or remove the affected object metadata before saving credentials.",
592                        ));
593                    }
594                    created_binding = true;
595                    create_and_verify_vault_key(vault, CAS_BINDING_KEY_ACCOUNT)?
596                }
597            };
598
599            let record = encrypt_record(&master, profile_id, provider, &plaintext)?;
600            keys.profiles
601                .entry(profile_id.to_owned())
602                .or_default()
603                .insert(provider.storage_key_name(), record)
604                .map_err(|message| {
605                    ProfileError::byok(ByokErrorCode::MalformedKeyStore, message)
606                })?;
607            transaction.save(&keys)?;
608            Ok(())
609        })();
610
611        if result.is_err() {
612            if created_binding {
613                let _ = vault.delete(CAS_BINDING_KEY_ACCOUNT);
614            }
615            if created_master {
616                let _ = vault.delete(RECORD_ENCRYPTION_MASTER_ACCOUNT);
617            }
618        }
619        result
620    })
621}
622
623pub fn delete_api_key(
624    store: &ProfileStore,
625    profile_id: &str,
626    provider: ApiKeyProvider,
627) -> Result<bool> {
628    delete_api_key_with_vault(store, profile_id, provider, &OsSecretVault)
629}
630
631fn delete_api_key_with_vault<V: SecretVault>(
632    store: &ProfileStore,
633    profile_id: &str,
634    provider: ApiKeyProvider,
635    vault: &V,
636) -> Result<bool> {
637    store.with_key_store_transaction(|transaction| {
638        let previous = transaction.load()?;
639        let mut keys = previous.clone();
640        let mut changed = false;
641        let mut remove_profile = false;
642        if let Some(profile_keys) = keys.profiles.get_mut(profile_id) {
643            changed = profile_keys.remove(provider.storage_key_name()).is_some();
644            remove_profile = profile_keys.is_empty();
645        }
646        if remove_profile {
647            keys.profiles.remove(profile_id);
648        }
649        if !changed {
650            return Ok(false);
651        }
652
653        transaction.save(&keys)?;
654        if keys.profiles.is_empty() {
655            if let Err(error) = vault.delete(RECORD_ENCRYPTION_MASTER_ACCOUNT) {
656                transaction.save(&previous)?;
657                return Err(error);
658            }
659        }
660        Ok(true)
661    })
662}
squigit_auth/security/crypto.rs

squigit_auth/security/
crypto.rs
