1use blake3;
7use ed25519_dalek::{
8 Signature, Signer as DalekSigner, SigningKey, Verifier as DalekVerifier, VerifyingKey,
9};
10
11use crate::binary::{SignedInvocation, SignedToolDef};
12use crate::SecurityError;
13
14pub struct Signer {
16 signing_key: SigningKey,
17}
18
19impl Signer {
20 pub fn from_seed(seed: &[u8; 32]) -> Self {
22 Self {
23 signing_key: SigningKey::from_bytes(seed),
24 }
25 }
26
27 pub fn generate() -> Self {
29 use ed25519_dalek::SigningKey;
30 let mut rng = rand::thread_rng();
31 Self {
32 signing_key: SigningKey::generate(&mut rng),
33 }
34 }
35
36 pub fn public_key_bytes(&self) -> [u8; 32] {
38 self.signing_key.verifying_key().to_bytes()
39 }
40
41 pub fn sign_tool_def(
43 &self,
44 tool_id: u32,
45 schema_hash: [u8; 32],
46 capabilities: u64,
47 ) -> SignedToolDef {
48 let mut message = Vec::with_capacity(44);
50 message.extend_from_slice(&tool_id.to_le_bytes());
51 message.extend_from_slice(&schema_hash);
52 message.extend_from_slice(&capabilities.to_le_bytes());
53
54 let signature = self.signing_key.sign(&message);
56
57 SignedToolDef {
58 tool_id,
59 schema_hash,
60 capabilities,
61 signature: signature.to_bytes(),
62 public_key: self.public_key_bytes(),
63 }
64 }
65
66 pub fn sign_invocation(
68 &self,
69 tool_id: u32,
70 nonce: u64,
71 timestamp: u64,
72 args: &[u8],
73 ) -> SignedInvocation {
74 let args_hash = *blake3::hash(args).as_bytes();
76
77 let mut message = Vec::with_capacity(52);
79 message.extend_from_slice(&tool_id.to_le_bytes());
80 message.extend_from_slice(&nonce.to_le_bytes());
81 message.extend_from_slice(×tamp.to_le_bytes());
82 message.extend_from_slice(&args_hash);
83
84 let signature = self.signing_key.sign(&message);
86
87 SignedInvocation {
88 tool_id,
89 nonce,
90 timestamp,
91 args_hash,
92 signature: signature.to_bytes(),
93 }
94 }
95}
96
97pub struct Verifier;
99
100impl Verifier {
101 pub fn verify_tool_def(def: &SignedToolDef) -> Result<(), SecurityError> {
103 let mut message = Vec::with_capacity(44);
105 message.extend_from_slice(&def.tool_id.to_le_bytes());
106 message.extend_from_slice(&def.schema_hash);
107 message.extend_from_slice(&def.capabilities.to_le_bytes());
108
109 let verifying_key = VerifyingKey::from_bytes(&def.public_key)
111 .map_err(|_| SecurityError::InvalidSignature)?;
112
113 let signature = Signature::from_bytes(&def.signature);
115
116 verifying_key
118 .verify(&message, &signature)
119 .map_err(|_| SecurityError::InvalidSignature)
120 }
121
122 pub fn verify_invocation(
124 inv: &SignedInvocation,
125 public_key: &[u8; 32],
126 ) -> Result<(), SecurityError> {
127 let mut message = Vec::with_capacity(52);
129 message.extend_from_slice(&inv.tool_id.to_le_bytes());
130 message.extend_from_slice(&inv.nonce.to_le_bytes());
131 message.extend_from_slice(&inv.timestamp.to_le_bytes());
132 message.extend_from_slice(&inv.args_hash);
133
134 let verifying_key =
136 VerifyingKey::from_bytes(public_key).map_err(|_| SecurityError::InvalidSignature)?;
137
138 let signature = Signature::from_bytes(&inv.signature);
140
141 verifying_key
143 .verify(&message, &signature)
144 .map_err(|_| SecurityError::InvalidSignature)
145 }
146
147 pub fn verify_args_hash(inv: &SignedInvocation, args: &[u8]) -> bool {
149 let computed_hash = *blake3::hash(args).as_bytes();
150 computed_hash == inv.args_hash
151 }
152}
153
154#[cfg(test)]
155mod tests {
156 use super::*;
157
158 #[test]
159 fn test_sign_and_verify_tool_def() {
160 let signer = Signer::from_seed(&[42u8; 32]);
161
162 let def = signer.sign_tool_def(123, [0xAB; 32], 0x1234);
163
164 assert_eq!(def.tool_id, 123);
165 assert_eq!(def.schema_hash, [0xAB; 32]);
166 assert_eq!(def.capabilities, 0x1234);
167 assert_eq!(def.public_key, signer.public_key_bytes());
168
169 assert!(Verifier::verify_tool_def(&def).is_ok());
171 }
172
173 #[test]
174 fn test_sign_and_verify_invocation() {
175 let signer = Signer::from_seed(&[42u8; 32]);
176 let args = b"test arguments";
177
178 let inv = signer.sign_invocation(456, 0xDEADBEEF, 1234567890, args);
179
180 assert_eq!(inv.tool_id, 456);
181 assert_eq!(inv.nonce, 0xDEADBEEF);
182 assert_eq!(inv.timestamp, 1234567890);
183
184 let public_key = signer.public_key_bytes();
186 assert!(Verifier::verify_invocation(&inv, &public_key).is_ok());
187
188 assert!(Verifier::verify_args_hash(&inv, args));
190 assert!(!Verifier::verify_args_hash(&inv, b"wrong args"));
191 }
192
193 #[test]
194 fn test_tampered_tool_def_fails() {
195 let signer = Signer::from_seed(&[42u8; 32]);
196
197 let mut def = signer.sign_tool_def(123, [0xAB; 32], 0x1234);
198
199 def.tool_id = 999;
201
202 assert!(Verifier::verify_tool_def(&def).is_err());
204 }
205
206 #[test]
207 fn test_tampered_invocation_fails() {
208 let signer = Signer::from_seed(&[42u8; 32]);
209
210 let mut inv = signer.sign_invocation(456, 0xDEADBEEF, 1234567890, b"args");
211
212 inv.nonce = 0xCAFEBABE;
214
215 let public_key = signer.public_key_bytes();
217 assert!(Verifier::verify_invocation(&inv, &public_key).is_err());
218 }
219
220 #[test]
221 fn test_wrong_public_key_fails() {
222 let signer1 = Signer::from_seed(&[1u8; 32]);
223 let signer2 = Signer::from_seed(&[2u8; 32]);
224
225 let inv = signer1.sign_invocation(456, 0xDEADBEEF, 1234567890, b"args");
226
227 let wrong_key = signer2.public_key_bytes();
229 assert!(Verifier::verify_invocation(&inv, &wrong_key).is_err());
230 }
231
232 #[test]
233 fn test_generate_random_signer() {
234 let signer1 = Signer::generate();
235 let signer2 = Signer::generate();
236
237 assert_ne!(signer1.public_key_bytes(), signer2.public_key_bytes());
239 }
240}