Skip to main content

sh_layer2/checkpoint_system/
checksum.rs

1//! # Checksum Utilities
2//!
3//! 校验和计算和验证工具。
4
5use sha2::{Digest, Sha256};
6
7/// 检查点版本
8const CHECKPOINT_VERSION: &str = "1.0";
9const CHECKSUM_FIELD: &str = "_checksum";
10const VERSION_FIELD: &str = "_version";
11
12/// 校验和工具
13pub struct ChecksumUtils;
14
15impl ChecksumUtils {
16    /// 计算数据的 SHA-256 校验和
17    ///
18    /// # Arguments
19    /// * `data` - 要计算校验和的数据
20    ///
21    /// # Returns
22    /// 十六进制格式的校验和字符串
23    pub fn compute_checksum(data: &serde_json::Value) -> String {
24        // 创建规范 JSON(排序键,无空白)
25        let canonical = Self::canonicalize_json(data);
26
27        // 计算 SHA-256
28        let mut hasher = Sha256::new();
29        hasher.update(canonical.as_bytes());
30        format!("{:x}", hasher.finalize())
31    }
32
33    /// 添加校验和到数据
34    ///
35    /// # Arguments
36    /// * `data` - 检查点数据
37    ///
38    /// # Returns
39    /// 带有校验和和版本的数据
40    pub fn add_checksum(mut data: serde_json::Value) -> serde_json::Value {
41        // 移除现有校验和
42        if let Some(obj) = data.as_object_mut() {
43            obj.remove(CHECKSUM_FIELD);
44            obj.remove(VERSION_FIELD);
45        }
46
47        let checksum = Self::compute_checksum(&data);
48
49        if let Some(obj) = data.as_object_mut() {
50            obj.insert(
51                CHECKSUM_FIELD.to_string(),
52                serde_json::Value::String(checksum),
53            );
54            obj.insert(
55                VERSION_FIELD.to_string(),
56                serde_json::Value::String(CHECKPOINT_VERSION.to_string()),
57            );
58        }
59
60        data
61    }
62
63    /// 验证校验和
64    ///
65    /// # Arguments
66    /// * `data` - 带有校验和的数据
67    ///
68    /// # Returns
69    /// 元组:(是否有效, 错误信息)
70    pub fn verify_checksum(data: &serde_json::Value) -> (bool, Option<String>) {
71        let obj = match data.as_object() {
72            Some(o) => o,
73            None => return (false, Some("Data is not an object".to_string())),
74        };
75
76        // 检查校验和字段
77        let expected_checksum = match obj.get(CHECKSUM_FIELD) {
78            Some(v) => v.as_str().unwrap_or("").to_string(),
79            None => return (false, Some("Missing checksum field".to_string())),
80        };
81
82        // 检查版本字段
83        let version = match obj.get(VERSION_FIELD) {
84            Some(v) => v.as_str().unwrap_or(""),
85            None => return (false, Some("Missing version field".to_string())),
86        };
87
88        if version != CHECKPOINT_VERSION {
89            return (
90                false,
91                Some(format!(
92                    "Version mismatch: expected {}, got {}",
93                    CHECKPOINT_VERSION, version
94                )),
95            );
96        }
97
98        // 计算实际校验和
99        let mut data_copy = data.clone();
100        if let Some(obj) = data_copy.as_object_mut() {
101            obj.remove(CHECKSUM_FIELD);
102            obj.remove(VERSION_FIELD);
103        }
104
105        let actual_checksum = Self::compute_checksum(&data_copy);
106
107        if expected_checksum != actual_checksum {
108            return (
109                false,
110                Some(format!(
111                    "Checksum mismatch: expected {}..., got {}...",
112                    &expected_checksum[..16.min(expected_checksum.len())],
113                    &actual_checksum[..16.min(actual_checksum.len())]
114                )),
115            );
116        }
117
118        (true, None)
119    }
120
121    /// 规范化 JSON(排序键,紧凑格式)
122    fn canonicalize_json(data: &serde_json::Value) -> String {
123        // 递归排序对象键
124        let sorted = Self::sort_json_keys(data);
125        serde_json::to_string(&sorted).unwrap_or_default()
126    }
127
128    /// 递归排序 JSON 对象的键
129    fn sort_json_keys(value: &serde_json::Value) -> serde_json::Value {
130        match value {
131            serde_json::Value::Object(map) => {
132                let mut sorted_map = serde_json::Map::new();
133                let mut keys: Vec<_> = map.keys().collect();
134                keys.sort();
135
136                for key in keys {
137                    if let Some(val) = map.get(key) {
138                        sorted_map.insert(key.clone(), Self::sort_json_keys(val));
139                    }
140                }
141
142                serde_json::Value::Object(sorted_map)
143            }
144            serde_json::Value::Array(arr) => {
145                serde_json::Value::Array(arr.iter().map(Self::sort_json_keys).collect())
146            }
147            other => other.clone(),
148        }
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155
156    #[test]
157    fn test_compute_checksum() {
158        let data = serde_json::json!({"test": "value"});
159        let checksum = ChecksumUtils::compute_checksum(&data);
160
161        assert!(!checksum.is_empty());
162        assert_eq!(checksum.len(), 64); // SHA-256 = 64 hex chars
163    }
164
165    #[test]
166    fn test_add_and_verify_checksum() {
167        let data = serde_json::json!({"session_id": "test123"});
168        let data_with_checksum = ChecksumUtils::add_checksum(data);
169
170        let (valid, error) = ChecksumUtils::verify_checksum(&data_with_checksum);
171        assert!(valid);
172        assert!(error.is_none());
173    }
174
175    #[test]
176    fn test_verify_missing_checksum() {
177        let data = serde_json::json!({"session_id": "test123"});
178
179        let (valid, error) = ChecksumUtils::verify_checksum(&data);
180        assert!(!valid);
181        assert!(error.is_some());
182    }
183
184    #[test]
185    fn test_deterministic_checksum() {
186        let data1 = serde_json::json!({"b": 2, "a": 1});
187        let data2 = serde_json::json!({"a": 1, "b": 2});
188
189        let checksum1 = ChecksumUtils::compute_checksum(&data1);
190        let checksum2 = ChecksumUtils::compute_checksum(&data2);
191
192        // 键顺序不同,但校验和应该相同
193        assert_eq!(checksum1, checksum2);
194    }
195}