tibba_crypto/
key_grip.rs

1// Copyright 2025 Tree xie.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15// Import necessary dependencies for cryptographic operations and error handling
16use super::Error;
17use hex::encode;
18use hmac::{Hmac, Mac};
19use sha2::Sha256;
20use std::sync::Arc;
21use std::sync::RwLock;
22
23/// Custom Result type using the crate's Error type
24type Result<T> = std::result::Result<T, Error>;
25
26/// Type alias for HMAC-SHA256 implementation
27type HmacSha256 = Hmac<Sha256>;
28
29enum KeyStore {
30    Static(Vec<Vec<u8>>),
31    Shared(Arc<RwLock<Vec<Vec<u8>>>>),
32}
33
34/// KeyGrip struct manages a set of cryptographic keys
35/// Provides both thread-safe (RwLock) and non-thread-safe implementations
36pub struct KeyGrip {
37    store: KeyStore,
38}
39
40/// Helper function to create an HMAC-SHA256 signature
41/// Returns the hex-encoded signature string
42fn sign_with_key(data: &[u8], key: &[u8]) -> Result<String> {
43    let mut mac = HmacSha256::new_from_slice(key).map_err(|e| Error::HmacSha256 {
44        message: e.to_string(),
45    })?;
46    mac.update(data);
47    Ok(encode(mac.finalize().into_bytes()))
48}
49
50impl KeyGrip {
51    /// Creates a new KeyGrip instance with non-thread-safe key storage
52    /// Returns error if keys vector is empty
53    pub fn new(keys: Vec<Vec<u8>>) -> Result<Self> {
54        if keys.is_empty() {
55            return Err(Error::KeyGripEmpty);
56        }
57        Ok(KeyGrip {
58            store: KeyStore::Static(keys),
59        })
60    }
61
62    /// Creates a new KeyGrip instance with thread-safe key storage using RwLock
63    pub fn new_with_lock(keys: Vec<Vec<u8>>) -> Result<Self> {
64        Ok(KeyGrip {
65            store: KeyStore::Shared(Arc::new(RwLock::new(keys))),
66        })
67    }
68    fn with_keys<F, R>(&self, f: F) -> R
69    where
70        F: FnOnce(&[Vec<u8>]) -> R,
71    {
72        match &self.store {
73            KeyStore::Static(keys) => f(keys),
74            KeyStore::Shared(lock_keys) => {
75                // it will not fail
76                if let Ok(keys) = lock_keys.read() {
77                    f(&keys)
78                } else {
79                    f(&[])
80                }
81            }
82        }
83    }
84
85    /// Updates the keys in the thread-safe storage
86    /// No-op if using non-thread-safe storage
87    pub fn update_keys(&self, new_keys: Vec<Vec<u8>>) {
88        if let KeyStore::Shared(lock_keys) = &self.store
89            && let Ok(mut keys) = lock_keys.write()
90        {
91            *keys = new_keys;
92        }
93    }
94
95    /// Finds the index of the key that was used to create the given digest
96    /// Returns -1 if no matching key is found
97    fn index(&self, data: &[u8], digest: &str) -> Result<Option<usize>> {
98        // no need to clone
99        self.with_keys(|keys| {
100            for (index, key) in keys.iter().enumerate() {
101                // we must handle the error of sign_with_key
102                match sign_with_key(data, key) {
103                    Ok(signature) if signature == digest => return Ok(Some(index)),
104                    // if the signature does not match, continue
105                    Ok(_) => continue,
106                    // if the signature process itself fails (e.g., invalid key), pass the error up
107                    Err(e) => return Err(e),
108                }
109            }
110            // if the loop ends without finding a match, return Ok(None)
111            Ok(None)
112        })
113    }
114
115    /// Signs the input data using the first key in the key set
116    /// Returns error if no keys are available
117    pub fn sign(&self, data: &[u8]) -> Result<String> {
118        self.with_keys(|keys| {
119            // new() has already guaranteed that keys is not empty
120            sign_with_key(data, &keys[0])
121        })
122    }
123
124    /// Verifies a signature (digest) against the input data
125    /// Returns tuple (is_valid, is_current):
126    /// - is_valid: true if signature matches any key
127    /// - is_current: true if signature matches the current (first) key
128    pub fn verify(&self, data: &[u8], digest: &str) -> Result<(bool, bool)> {
129        match self.index(data, digest)? {
130            Some(0) => Ok((true, true)),  // match the first key, valid and current
131            Some(_) => Ok((true, false)), // match other keys, valid but not current
132            None => Ok((false, false)),   // no match found
133        }
134    }
135}