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}