shamir_rs/extensions/
hierarchical.rs

1use crate::{Scheme, Share, SssError};
2use num_bigint::BigUint;
3
4#[derive(Debug, Clone)]
5pub enum Role {
6    President,     // Gets 3 shares
7    VicePresident, // Gets 2 shares
8    Executive,     // Gets 1 share
9}
10
11#[derive(Debug)]
12pub struct HierarchicalShare {
13    pub role: Role,
14    pub shares: Vec<Share>,
15}
16
17pub struct HierarchicalScheme {
18    scheme: Scheme,
19}
20
21impl HierarchicalScheme {
22    pub fn new(prime_modulus: BigUint) -> Result<Self, SssError> {
23        // We use (3,n) scheme as base, where n will be determined by total shares
24        Ok(HierarchicalScheme {
25            scheme: Scheme::new(3, 100, prime_modulus)?, // n=100 as safe upper bound
26        })
27    }
28
29    pub fn split_secret(&self, secret: &BigUint) -> Vec<Share> {
30        self.scheme.split_secret(secret)
31    }
32
33    pub fn assign_shares(&self, all_shares: Vec<Share>, role: Role) -> HierarchicalShare {
34        let num_shares = match role {
35            Role::President => 3,
36            Role::VicePresident => 2,
37            Role::Executive => 1,
38        };
39
40        // Take sequential shares for this role
41        let role_shares: Vec<Share> = all_shares.into_iter().take(num_shares).collect();
42
43        HierarchicalShare {
44            role,
45            shares: role_shares,
46        }
47    }
48
49    pub fn reconstruct_secret(&self, shares: &[Share]) -> Result<BigUint, SssError> {
50        self.scheme.reconstruct_secret(shares)
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57
58    #[test]
59    fn test_president_alone() {
60        let scheme = HierarchicalScheme::new(BigUint::from(257u32)).unwrap();
61        let secret = BigUint::from(123u32);
62        let all_shares = scheme.split_secret(&secret);
63
64        // President gets 3 shares
65        let president = scheme.assign_shares(all_shares, Role::President);
66
67        // President should be able to reconstruct alone
68        let reconstructed = scheme.reconstruct_secret(&president.shares).unwrap();
69        assert_eq!(reconstructed, secret);
70    }
71
72    #[test]
73    fn test_two_vice_presidents() {
74        let scheme = HierarchicalScheme::new(BigUint::from(257u32)).unwrap();
75        let secret = BigUint::from(123u32);
76        let all_shares = scheme.split_secret(&secret);
77
78        // Create shares for two VPs
79        let mut shares_iter = all_shares.into_iter();
80        let vp1_shares = shares_iter.by_ref().take(2).collect();
81        let vp2_shares = shares_iter.by_ref().take(2).collect();
82
83        let vp1 = HierarchicalShare {
84            role: Role::VicePresident,
85            shares: vp1_shares,
86        };
87
88        let vp2 = HierarchicalShare {
89            role: Role::VicePresident,
90            shares: vp2_shares,
91        };
92
93        // Combine some shares from both VPs to reconstruct
94        let mut combined_shares = Vec::new();
95        combined_shares.extend_from_slice(&vp1.shares[0..2]);
96        combined_shares.push(vp2.shares[0].clone());
97
98        let reconstructed = scheme.reconstruct_secret(&combined_shares).unwrap();
99        assert_eq!(reconstructed, secret);
100    }
101
102    #[test]
103    fn test_vp_and_executives() {
104        let scheme = HierarchicalScheme::new(BigUint::from(257u32)).unwrap();
105        let secret = BigUint::from(123u32);
106        let all_shares = scheme.split_secret(&secret);
107
108        let mut shares_iter = all_shares.into_iter();
109
110        // VP gets 2 shares
111        let vp_shares = shares_iter.by_ref().take(2).collect();
112        let vp = HierarchicalShare {
113            role: Role::VicePresident,
114            shares: vp_shares,
115        };
116
117        // Exec gets 1 share
118        let exec_shares = shares_iter.by_ref().take(1).collect();
119        let exec = HierarchicalShare {
120            role: Role::Executive,
121            shares: exec_shares,
122        };
123
124        // VP + Exec should be able to reconstruct
125        let mut combined_shares = Vec::new();
126        combined_shares.extend_from_slice(&vp.shares);
127        combined_shares.extend_from_slice(&exec.shares);
128
129        let reconstructed = scheme.reconstruct_secret(&combined_shares).unwrap();
130        assert_eq!(reconstructed, secret);
131    }
132
133    #[test]
134    fn test_three_executives() {
135        let scheme = HierarchicalScheme::new(BigUint::from(257u32)).unwrap();
136        let secret = BigUint::from(123u32);
137        let all_shares = scheme.split_secret(&secret);
138
139        let mut shares_iter = all_shares.into_iter();
140
141        // Create three executives with one share each
142        let exec1 = HierarchicalShare {
143            role: Role::Executive,
144            shares: shares_iter.by_ref().take(1).collect(),
145        };
146        let exec2 = HierarchicalShare {
147            role: Role::Executive,
148            shares: shares_iter.by_ref().take(1).collect(),
149        };
150        let exec3 = HierarchicalShare {
151            role: Role::Executive,
152            shares: shares_iter.by_ref().take(1).collect(),
153        };
154
155        // Three executives together should be able to reconstruct
156        let mut combined_shares = Vec::new();
157        combined_shares.extend_from_slice(&exec1.shares);
158        combined_shares.extend_from_slice(&exec2.shares);
159        combined_shares.extend_from_slice(&exec3.shares);
160
161        let reconstructed = scheme.reconstruct_secret(&combined_shares).unwrap();
162        assert_eq!(reconstructed, secret);
163    }
164}