sh_layer2/checkpoint_system/
checksum.rs1use sha2::{Digest, Sha256};
6
7const CHECKPOINT_VERSION: &str = "1.0";
9const CHECKSUM_FIELD: &str = "_checksum";
10const VERSION_FIELD: &str = "_version";
11
12pub struct ChecksumUtils;
14
15impl ChecksumUtils {
16 pub fn compute_checksum(data: &serde_json::Value) -> String {
24 let canonical = Self::canonicalize_json(data);
26
27 let mut hasher = Sha256::new();
29 hasher.update(canonical.as_bytes());
30 format!("{:x}", hasher.finalize())
31 }
32
33 pub fn add_checksum(mut data: serde_json::Value) -> serde_json::Value {
41 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 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 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 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 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 fn canonicalize_json(data: &serde_json::Value) -> String {
123 let sorted = Self::sort_json_keys(data);
125 serde_json::to_string(&sorted).unwrap_or_default()
126 }
127
128 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); }
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 assert_eq!(checksum1, checksum2);
194 }
195}