request_response/
message.rs

1use std::{
2    io::{Cursor, Read, Write},
3    time::{Duration, SystemTime, UNIX_EPOCH},
4};
5
6use anyhow::{Context, Result};
7use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
8use hotshot_types::traits::signature_key::SignatureKey;
9
10use super::{RequestHash, Serializable, request::Request};
11
12/// The outer message type for the request-response protocol. Can either be a request or a response
13#[derive(Clone, Debug)]
14#[cfg_attr(test, derive(PartialEq, Eq))]
15pub enum Message<R: Request, K: SignatureKey> {
16    /// A request
17    Request(RequestMessage<R, K>),
18    /// A response
19    Response(ResponseMessage<R>),
20}
21
22/// A request message, which includes the requester's public key, the request's signature, a timestamp, and the request itself
23#[derive(Clone, Debug)]
24#[cfg_attr(test, derive(PartialEq, Eq))]
25pub struct RequestMessage<R: Request, K: SignatureKey> {
26    /// The requester's public key
27    pub public_key: K,
28    /// The requester's signature over the [the actual request content + timestamp]
29    pub signature: K::PureAssembledSignatureType,
30    /// The timestamp of when the request was sent (in seconds since the Unix epoch). We use this to
31    /// ensure that the request is not old, which is useful for preventing replay attacks.
32    pub timestamp_unix_seconds: u64,
33    /// The actual request data. This is from the application
34    pub request: R,
35}
36
37/// A response message, which includes the hash of the request we're responding to and the response itself.
38#[derive(Clone, Debug)]
39#[cfg_attr(test, derive(PartialEq, Eq))]
40pub struct ResponseMessage<R: Request> {
41    /// The hash of the application-specific request we're responding to. The hash is a free way
42    /// to identify the request and weed out any potential incompatibilities
43    pub request_hash: RequestHash,
44    /// The actual response content
45    pub response: R::Response,
46}
47
48impl<R: Request, K: SignatureKey> RequestMessage<R, K> {
49    /// Create a new signed request message from a request
50    ///
51    /// # Errors
52    /// - If the request's content cannot be serialized
53    /// - If the request cannot be signed
54    ///
55    /// # Panics
56    /// - If time is not monotonic
57    pub fn new_signed(public_key: &K, private_key: &K::PrivateKey, request: &R) -> Result<Self>
58    where
59        <K as SignatureKey>::SignError: 'static,
60    {
61        // Get the current timestamp
62        let timestamp_unix_seconds = SystemTime::now()
63            .duration_since(UNIX_EPOCH)
64            .expect("time went backwards")
65            .as_secs();
66
67        // Concatenate the content and timestamp
68        let content_to_sign = [
69            request
70                .to_bytes()
71                .with_context(|| "failed to serialize request content")?
72                .as_slice(),
73            timestamp_unix_seconds.to_le_bytes().as_slice(),
74            b"espresso-request-response",
75        ]
76        .concat();
77
78        // Sign the actual request content (+ a namespace) with the private key
79        let signature =
80            K::sign(private_key, &content_to_sign).with_context(|| "failed to sign message")?;
81
82        // Return the newly signed request message
83        Ok(RequestMessage {
84            public_key: public_key.clone(),
85            signature,
86            timestamp_unix_seconds,
87            request: request.clone(),
88        })
89    }
90
91    /// Validate the [`RequestMessage`], checking the signature and the timestamp and
92    /// calling the request's application-specific validation function
93    ///
94    /// # Errors
95    /// - If the request's signature is invalid
96    /// - If the request is too old
97    ///
98    /// # Panics
99    /// - If time is not monotonic
100    pub fn validate(&self, incoming_request_ttl: Duration) -> Result<()> {
101        // Make sure the request is not too old
102        if self
103            .timestamp_unix_seconds
104            .saturating_add(incoming_request_ttl.as_secs())
105            < SystemTime::now()
106                .duration_since(UNIX_EPOCH)
107                .expect("time went backwards")
108                .as_secs()
109        {
110            return Err(anyhow::anyhow!("request is too old"));
111        }
112        // Check the signature over the request content and timestamp
113        if !self.public_key.validate(
114            &self.signature,
115            &[
116                self.request.to_bytes()?,
117                self.timestamp_unix_seconds.to_le_bytes().to_vec(),
118                b"espresso-request-response".to_vec(),
119            ]
120            .concat(),
121        ) {
122            return Err(anyhow::anyhow!("invalid request signature"));
123        }
124
125        // Call the request's application-specific validation function
126        self.request.validate()
127    }
128}
129
130/// A blanket implementation of the [`Serializable`] trait for any [`Message`]
131impl<R: Request, K: SignatureKey> Serializable for Message<R, K> {
132    /// Converts any [`Message`] to bytes if the content is also [`Serializable`]
133    fn to_bytes(&self) -> Result<Vec<u8>> {
134        // Create a buffer for the bytes
135        let mut bytes = Vec::new();
136
137        // Convert the message to bytes based on the type. By default it is just type-prefixed
138        match self {
139            Message::Request(request_message) => {
140                // Write the type (request)
141                bytes.push(0);
142
143                // Write the request content
144                bytes.extend_from_slice(request_message.to_bytes()?.as_slice());
145            },
146            Message::Response(response_message) => {
147                // Write the type (response)
148                bytes.push(1);
149
150                // Write the response content
151                bytes.extend_from_slice(response_message.to_bytes()?.as_slice());
152            },
153        };
154
155        Ok(bytes)
156    }
157
158    /// Convert bytes to a [`Message`]
159    fn from_bytes(bytes: &[u8]) -> Result<Self> {
160        // Create a cursor so we can easily read the bytes in order
161        let mut bytes = Cursor::new(bytes);
162
163        // Get the message type
164        let type_byte = bytes.read_u8()?;
165
166        // Deserialize the message based on the type
167        match type_byte {
168            0 => {
169                // Read the `RequestMessage`
170                Ok(Message::Request(RequestMessage::from_bytes(&read_to_end(
171                    &mut bytes,
172                )?)?))
173            },
174            1 => {
175                // Read the `ResponseMessage`
176                Ok(Message::Response(ResponseMessage::from_bytes(
177                    &read_to_end(&mut bytes)?,
178                )?))
179            },
180            _ => Err(anyhow::anyhow!("invalid message type")),
181        }
182    }
183}
184
185impl<R: Request, K: SignatureKey> Serializable for RequestMessage<R, K> {
186    fn to_bytes(&self) -> Result<Vec<u8>> {
187        // Create a buffer for the bytes
188        let mut bytes = Vec::new();
189
190        // Write the public key (length-prefixed)
191        write_length_prefixed(&mut bytes, &self.public_key.to_bytes())?;
192
193        // Write the signature (length-prefixed)
194        write_length_prefixed(&mut bytes, &bincode::serialize(&self.signature)?)?;
195
196        // Write the timestamp
197        bytes.write_all(&self.timestamp_unix_seconds.to_le_bytes())?;
198
199        // Write the actual request
200        bytes.write_all(self.request.to_bytes()?.as_slice())?;
201
202        Ok(bytes)
203    }
204
205    fn from_bytes(bytes: &[u8]) -> Result<Self> {
206        // Create a cursor so we can easily read the bytes in order
207        let mut bytes = Cursor::new(bytes);
208
209        // Read the public key (length-prefixed)
210        let public_key = K::from_bytes(&read_length_prefixed(&mut bytes)?)?;
211
212        // Read the signature (length-prefixed)
213        let signature = bincode::deserialize(&read_length_prefixed(&mut bytes)?)?;
214
215        // Read the timestamp as a [`u64`]
216        let timestamp = bytes.read_u64::<LittleEndian>()?;
217
218        // Deserialize the request
219        let request = R::from_bytes(&read_to_end(&mut bytes)?)?;
220
221        Ok(Self {
222            public_key,
223            signature,
224            timestamp_unix_seconds: timestamp,
225            request,
226        })
227    }
228}
229
230impl<R: Request> Serializable for ResponseMessage<R> {
231    fn to_bytes(&self) -> Result<Vec<u8>> {
232        // Create a buffer for the bytes
233        let mut bytes = Vec::new();
234
235        // Write the request hash as bytes
236        bytes.write_all(self.request_hash.as_bytes())?;
237
238        // Write the response content
239        bytes.write_all(self.response.to_bytes()?.as_slice())?;
240
241        Ok(bytes)
242    }
243
244    fn from_bytes(bytes: &[u8]) -> Result<Self> {
245        // Create a buffer for the bytes
246        let mut bytes = Cursor::new(bytes);
247
248        // Read the request hash as a [`blake3::Hash`]
249        let mut request_hash_bytes = [0; 32];
250        bytes.read_exact(&mut request_hash_bytes)?;
251        let request_hash = RequestHash::from(request_hash_bytes);
252
253        // Read the response content to the end
254        let response = R::Response::from_bytes(&read_to_end(&mut bytes)?)?;
255
256        Ok(Self {
257            request_hash,
258            response,
259        })
260    }
261}
262
263/// A helper function to write a length-prefixed value to a writer
264fn write_length_prefixed<W: Write>(writer: &mut W, value: &[u8]) -> Result<()> {
265    // Write the length of the value as a u32
266    writer.write_u32::<LittleEndian>(
267        u32::try_from(value.len()).with_context(|| "value was too large")?,
268    )?;
269
270    // Write the (already serialized) value
271    writer.write_all(value)?;
272    Ok(())
273}
274
275/// A helper function to read a length-prefixed value from a reader
276fn read_length_prefixed<R: Read>(reader: &mut R) -> Result<Vec<u8>> {
277    // Read the length of the value as a u32
278    let length = reader.read_u32::<LittleEndian>()?;
279
280    // Read the value
281    let mut value = vec![0; length as usize];
282    reader.read_exact(&mut value)?;
283    Ok(value)
284}
285
286/// A helper function to read to the end of the reader
287fn read_to_end<R: Read>(reader: &mut R) -> Result<Vec<u8>> {
288    let mut value = Vec::new();
289    reader.read_to_end(&mut value)?;
290    Ok(value)
291}
292
293#[cfg(test)]
294mod tests {
295    use hotshot_types::signature_key::BLSPubKey;
296    use rand::Rng;
297
298    use super::*;
299
300    // A testing implementation of the [`Serializable`] trait for [`Vec<u8>`]
301    impl Serializable for Vec<u8> {
302        fn to_bytes(&self) -> Result<Vec<u8>> {
303            Ok(self.clone())
304        }
305        fn from_bytes(bytes: &[u8]) -> Result<Self> {
306            Ok(bytes.to_vec())
307        }
308    }
309
310    /// A testing implementation of the [`Request`] trait for [`Vec<u8>`]
311    impl Request for Vec<u8> {
312        type Response = Vec<u8>;
313
314        fn validate(&self) -> Result<()> {
315            Ok(())
316        }
317    }
318
319    /// Tests that properly signed requests are validated correctly and that invalid requests
320    /// (bad timestamp/signature) are rejected
321    #[test]
322    fn test_request_validation() {
323        // Create some RNG
324        let mut rng = rand::thread_rng();
325
326        for _ in 0..100 {
327            // Create a random keypair
328            let (public_key, private_key) =
329                BLSPubKey::generated_from_seed_indexed([1; 32], rng.r#gen::<u64>());
330
331            // Create a valid request with some random content
332            let mut request = RequestMessage::new_signed(
333                &public_key,
334                &private_key,
335                &vec![rng.r#gen::<u8>(); rng.gen_range(1..10000)],
336            )
337            .expect("Failed to create signed request");
338
339            let (should_be_valid, request_ttl) = match rng.gen_range(0..4) {
340                0 => (true, Duration::from_secs(1)),
341
342                1 => {
343                    // Alter the requests's actual content
344                    request.request[0] = !request.request[0];
345
346                    // It should not be valid anymore
347                    (false, Duration::from_secs(1))
348                },
349
350                2 => {
351                    // Alter the timestamp
352                    request.timestamp_unix_seconds += 1000;
353
354                    // It should not be valid anymore
355                    (false, Duration::from_secs(1))
356                },
357
358                3 => {
359                    // Change the request ttl to be 0. This should make the request
360                    // invalid immediately
361                    (true, Duration::from_secs(0))
362                },
363
364                _ => unreachable!(),
365            };
366
367            // Validate the request
368            assert_eq!(request.validate(request_ttl).is_ok(), should_be_valid);
369        }
370    }
371
372    /// Tests that messages are serialized and deserialized correctly
373    #[test]
374    fn test_message_parity() {
375        for _ in 0..100 {
376            // Create some RNG
377            let mut rng = rand::thread_rng();
378
379            // Generate a random message type
380            let is_request = rng.r#gen::<u8>() % 2 == 0;
381
382            // The request content will be a random vector of bytes
383            let request = vec![rng.r#gen::<u8>(); rng.gen_range(0..10000)];
384
385            // Create a message
386            let message = if is_request {
387                // Create a random keypair
388                let (public_key, private_key) =
389                    BLSPubKey::generated_from_seed_indexed([1; 32], rng.r#gen::<u64>());
390
391                // Create a new signed request
392                let request = RequestMessage::new_signed(&public_key, &private_key, &request)
393                    .expect("Failed to create signed request");
394
395                Message::Request(request)
396            } else {
397                // Create a response message
398                Message::Response(ResponseMessage {
399                    request_hash: blake3::hash(&request),
400                    response: vec![rng.r#gen::<u8>(); rng.gen_range(0..10000)],
401                })
402            };
403
404            // Serialize the message
405            let serialized = message.to_bytes().expect("Failed to serialize message");
406
407            // Deserialize the message
408            let deserialized =
409                Message::from_bytes(&serialized).expect("Failed to deserialize message");
410
411            // Assert that the deserialized message is the same as the original message
412            assert_eq!(message, deserialized);
413        }
414    }
415
416    /// Tests that length-prefixed values are read and written correctly
417    #[test]
418    fn test_length_prefix_parity() {
419        // Create some RNG
420        let mut rng = rand::thread_rng();
421
422        for _ in 0..100 {
423            // Create a buffer to test over
424            let mut bytes = Vec::new();
425
426            // Generate the value to test over
427            let value = vec![rng.r#gen::<u8>(); rng.gen_range(0..10000)];
428
429            // Write the length-prefixed value
430            write_length_prefixed(&mut bytes, &value).unwrap();
431
432            // Create a reader from the bytes
433            let mut reader = Cursor::new(bytes);
434
435            // Read the length-prefixed value
436            let value = read_length_prefixed(&mut reader).unwrap();
437            assert_eq!(value, value);
438        }
439    }
440}
request_response/message.rs

request_response/
message.rs
