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
48/// The type byte that prefixes a serialized [`Message::Request`]
49const REQUEST_TAG: u8 = 0;
50/// The type byte that prefixes a serialized [`Message::Response`]
51const RESPONSE_TAG: u8 = 1;
52
53/// A serialized [`Message`] with only its framing parsed: the type byte, and for a response the
54/// request hash. The receiver uses it to route a response to its waiters, or drop it, without
55/// deserializing the body
56#[derive(Clone, Copy, Debug)]
57pub enum MessageFrame<'a> {
58    /// A request. `body` is a serialized [`RequestMessage`]
59    Request { body: &'a [u8] },
60    /// A response to the request with `request_hash`. `body` is a serialized response
61    Response {
62        request_hash: RequestHash,
63        body: &'a [u8],
64    },
65}
66
67impl<'a> MessageFrame<'a> {
68    /// Parse the framing of a serialized [`Message`]
69    ///
70    /// # Errors
71    /// - If the message is empty or has an unknown type byte
72    /// - If a response is too short to hold a request hash
73    pub fn parse(bytes: &'a [u8]) -> Result<Self> {
74        let (&type_byte, body) = bytes.split_first().context("empty message")?;
75
76        match type_byte {
77            REQUEST_TAG => Ok(Self::Request { body }),
78            RESPONSE_TAG => {
79                let (request_hash, body) = split_request_hash(body)?;
80                Ok(Self::Response { request_hash, body })
81            },
82            _ => Err(anyhow::anyhow!("invalid message type")),
83        }
84    }
85}
86
87impl<R: Request, K: SignatureKey> RequestMessage<R, K> {
88    /// Create a new signed request message from a request
89    ///
90    /// # Errors
91    /// - If the request's content cannot be serialized
92    /// - If the request cannot be signed
93    ///
94    /// # Panics
95    /// - If time is not monotonic
96    pub fn new_signed(public_key: &K, private_key: &K::PrivateKey, request: &R) -> Result<Self>
97    where
98        <K as SignatureKey>::SignError: 'static,
99    {
100        // Get the current timestamp
101        let timestamp_unix_seconds = SystemTime::now()
102            .duration_since(UNIX_EPOCH)
103            .expect("time went backwards")
104            .as_secs();
105
106        // Concatenate the content and timestamp
107        let content_to_sign = [
108            request
109                .to_bytes()
110                .with_context(|| "failed to serialize request content")?
111                .as_slice(),
112            timestamp_unix_seconds.to_le_bytes().as_slice(),
113            b"espresso-request-response",
114        ]
115        .concat();
116
117        // Sign the actual request content (+ a namespace) with the private key
118        let signature =
119            K::sign(private_key, &content_to_sign).with_context(|| "failed to sign message")?;
120
121        // Return the newly signed request message
122        Ok(RequestMessage {
123            public_key: public_key.clone(),
124            signature,
125            timestamp_unix_seconds,
126            request: request.clone(),
127        })
128    }
129
130    /// Validate the [`RequestMessage`], checking the signature and the timestamp and
131    /// calling the request's application-specific validation function
132    ///
133    /// # Errors
134    /// - If the request's signature is invalid
135    /// - If the request is too old
136    ///
137    /// # Panics
138    /// - If time is not monotonic
139    pub fn validate(&self, incoming_request_ttl: Duration) -> Result<()> {
140        // Make sure the request is not too old
141        if self
142            .timestamp_unix_seconds
143            .saturating_add(incoming_request_ttl.as_secs())
144            < SystemTime::now()
145                .duration_since(UNIX_EPOCH)
146                .expect("time went backwards")
147                .as_secs()
148        {
149            return Err(anyhow::anyhow!("request is too old"));
150        }
151        // Check the signature over the request content and timestamp
152        if !self.public_key.validate(
153            &self.signature,
154            &[
155                self.request.to_bytes()?,
156                self.timestamp_unix_seconds.to_le_bytes().to_vec(),
157                b"espresso-request-response".to_vec(),
158            ]
159            .concat(),
160        ) {
161            return Err(anyhow::anyhow!("invalid request signature"));
162        }
163
164        // Call the request's application-specific validation function
165        self.request.validate()
166    }
167}
168
169/// A blanket implementation of the [`Serializable`] trait for any [`Message`]
170impl<R: Request, K: SignatureKey> Serializable for Message<R, K> {
171    /// Converts any [`Message`] to bytes if the content is also [`Serializable`]
172    fn to_bytes(&self) -> Result<Vec<u8>> {
173        // Create a buffer for the bytes
174        let mut bytes = Vec::new();
175
176        // Convert the message to bytes based on the type. By default it is just type-prefixed
177        match self {
178            Message::Request(request_message) => {
179                // Write the type (request)
180                bytes.push(REQUEST_TAG);
181
182                // Write the request content
183                bytes.extend_from_slice(request_message.to_bytes()?.as_slice());
184            },
185            Message::Response(response_message) => {
186                // Write the type (response)
187                bytes.push(RESPONSE_TAG);
188
189                // Write the response content
190                bytes.extend_from_slice(response_message.to_bytes()?.as_slice());
191            },
192        };
193
194        Ok(bytes)
195    }
196
197    /// Convert bytes to a [`Message`]
198    fn from_bytes(bytes: &[u8]) -> Result<Self> {
199        Ok(match MessageFrame::parse(bytes)? {
200            MessageFrame::Request { body } => Message::Request(RequestMessage::from_bytes(body)?),
201            MessageFrame::Response { request_hash, body } => Message::Response(ResponseMessage {
202                request_hash,
203                response: R::Response::from_bytes(body)?,
204            }),
205        })
206    }
207}
208
209impl<R: Request, K: SignatureKey> Serializable for RequestMessage<R, K> {
210    fn to_bytes(&self) -> Result<Vec<u8>> {
211        // Create a buffer for the bytes
212        let mut bytes = Vec::new();
213
214        // Write the public key (length-prefixed)
215        write_length_prefixed(&mut bytes, &self.public_key.to_bytes())?;
216
217        // Write the signature (length-prefixed)
218        write_length_prefixed(&mut bytes, &bincode::serialize(&self.signature)?)?;
219
220        // Write the timestamp
221        bytes.write_all(&self.timestamp_unix_seconds.to_le_bytes())?;
222
223        // Write the actual request
224        bytes.write_all(self.request.to_bytes()?.as_slice())?;
225
226        Ok(bytes)
227    }
228
229    fn from_bytes(bytes: &[u8]) -> Result<Self> {
230        // Create a cursor so we can easily read the bytes in order
231        let mut bytes = Cursor::new(bytes);
232
233        // Read the public key (length-prefixed)
234        let public_key = K::from_bytes(&read_length_prefixed(&mut bytes)?)?;
235
236        // Read the signature (length-prefixed)
237        let signature = bincode::deserialize(&read_length_prefixed(&mut bytes)?)?;
238
239        // Read the timestamp as a [`u64`]
240        let timestamp = bytes.read_u64::<LittleEndian>()?;
241
242        // Deserialize the request
243        let request = R::from_bytes(&read_to_end(&mut bytes)?)?;
244
245        Ok(Self {
246            public_key,
247            signature,
248            timestamp_unix_seconds: timestamp,
249            request,
250        })
251    }
252}
253
254impl<R: Request> Serializable for ResponseMessage<R> {
255    fn to_bytes(&self) -> Result<Vec<u8>> {
256        // Create a buffer for the bytes
257        let mut bytes = Vec::new();
258
259        // Write the request hash as bytes
260        bytes.write_all(self.request_hash.as_bytes())?;
261
262        // Write the response content
263        bytes.write_all(self.response.to_bytes()?.as_slice())?;
264
265        Ok(bytes)
266    }
267
268    fn from_bytes(bytes: &[u8]) -> Result<Self> {
269        let (request_hash, body) = split_request_hash(bytes)?;
270        Ok(Self {
271            request_hash,
272            response: R::Response::from_bytes(body)?,
273        })
274    }
275}
276
277/// Split the request hash prefix off a serialized [`ResponseMessage`]
278fn split_request_hash(bytes: &[u8]) -> Result<(RequestHash, &[u8])> {
279    let (hash, body) = bytes
280        .split_at_checked(blake3::OUT_LEN)
281        .context("response message is shorter than a request hash")?;
282    let hash: [u8; blake3::OUT_LEN] = hash.try_into()?;
283    Ok((RequestHash::from(hash), body))
284}
285
286/// A helper function to write a length-prefixed value to a writer
287fn write_length_prefixed<W: Write>(writer: &mut W, value: &[u8]) -> Result<()> {
288    // Write the length of the value as a u32
289    writer.write_u32::<LittleEndian>(
290        u32::try_from(value.len()).with_context(|| "value was too large")?,
291    )?;
292
293    // Write the (already serialized) value
294    writer.write_all(value)?;
295    Ok(())
296}
297
298/// A helper function to read a length-prefixed value from a reader
299fn read_length_prefixed<R: Read>(reader: &mut R) -> Result<Vec<u8>> {
300    // Read the length of the value as a u32
301    let length = u64::from(reader.read_u32::<LittleEndian>()?);
302
303    let mut value = Vec::with_capacity(length.min(64 * 1024) as usize);
304    let read = reader.take(length).read_to_end(&mut value)?;
305    if read as u64 != length {
306        return Err(anyhow::anyhow!("length prefix exceeds available data"));
307    }
308    Ok(value)
309}
310
311/// A helper function to read to the end of the reader
312fn read_to_end<R: Read>(reader: &mut R) -> Result<Vec<u8>> {
313    let mut value = Vec::new();
314    reader.read_to_end(&mut value)?;
315    Ok(value)
316}
317
318#[cfg(test)]
319mod tests {
320    use hotshot_types::signature_key::BLSPubKey;
321    use rand::Rng;
322
323    use super::*;
324
325    // A testing implementation of the [`Serializable`] trait for [`Vec<u8>`]
326    impl Serializable for Vec<u8> {
327        fn to_bytes(&self) -> Result<Vec<u8>> {
328            Ok(self.clone())
329        }
330        fn from_bytes(bytes: &[u8]) -> Result<Self> {
331            Ok(bytes.to_vec())
332        }
333    }
334
335    /// A testing implementation of the [`Request`] trait for [`Vec<u8>`]
336    impl Request for Vec<u8> {
337        type Response = Vec<u8>;
338
339        fn validate(&self) -> Result<()> {
340            Ok(())
341        }
342    }
343
344    /// Tests that properly signed requests are validated correctly and that invalid requests
345    /// (bad timestamp/signature) are rejected
346    #[test]
347    fn test_request_validation() {
348        // Create some RNG
349        let mut rng = rand::thread_rng();
350
351        for _ in 0..100 {
352            // Create a random keypair
353            let (public_key, private_key) =
354                BLSPubKey::generated_from_seed_indexed([1; 32], rng.r#gen::<u64>());
355
356            // Create a valid request with some random content
357            let mut request = RequestMessage::new_signed(
358                &public_key,
359                &private_key,
360                &vec![rng.r#gen::<u8>(); rng.gen_range(1..10000)],
361            )
362            .expect("Failed to create signed request");
363
364            let (should_be_valid, request_ttl) = match rng.gen_range(0..4) {
365                0 => (true, Duration::from_secs(1)),
366
367                1 => {
368                    // Alter the requests's actual content
369                    request.request[0] = !request.request[0];
370
371                    // It should not be valid anymore
372                    (false, Duration::from_secs(1))
373                },
374
375                2 => {
376                    // Alter the timestamp
377                    request.timestamp_unix_seconds += 1000;
378
379                    // It should not be valid anymore
380                    (false, Duration::from_secs(1))
381                },
382
383                3 => {
384                    // Change the request ttl to be 0. This should make the request
385                    // invalid immediately
386                    (true, Duration::from_secs(0))
387                },
388
389                _ => unreachable!(),
390            };
391
392            // Validate the request
393            assert_eq!(request.validate(request_ttl).is_ok(), should_be_valid);
394        }
395    }
396
397    /// Tests that messages are serialized and deserialized correctly
398    #[test]
399    fn test_message_parity() {
400        for _ in 0..100 {
401            // Create some RNG
402            let mut rng = rand::thread_rng();
403
404            // Generate a random message type
405            let is_request = rng.r#gen::<u8>() % 2 == 0;
406
407            // The request content will be a random vector of bytes
408            let request = vec![rng.r#gen::<u8>(); rng.gen_range(0..10000)];
409
410            // Create a message
411            let message = if is_request {
412                // Create a random keypair
413                let (public_key, private_key) =
414                    BLSPubKey::generated_from_seed_indexed([1; 32], rng.r#gen::<u64>());
415
416                // Create a new signed request
417                let request = RequestMessage::new_signed(&public_key, &private_key, &request)
418                    .expect("Failed to create signed request");
419
420                Message::Request(request)
421            } else {
422                // Create a response message
423                Message::Response(ResponseMessage {
424                    request_hash: blake3::hash(&request),
425                    response: vec![rng.r#gen::<u8>(); rng.gen_range(0..10000)],
426                })
427            };
428
429            // Serialize the message
430            let serialized = message.to_bytes().expect("Failed to serialize message");
431
432            // Deserialize the message
433            let deserialized =
434                Message::from_bytes(&serialized).expect("Failed to deserialize message");
435
436            // Assert that the deserialized message is the same as the original message
437            assert_eq!(message, deserialized);
438
439            // The frame alone must agree with the full deserialization
440            let frame = MessageFrame::parse(&serialized).expect("Failed to parse frame");
441            match (&message, frame) {
442                (Message::Request(_), MessageFrame::Request { .. }) => {},
443                (Message::Response(response), MessageFrame::Response { request_hash, .. }) => {
444                    assert_eq!(request_hash, response.request_hash);
445                },
446                (message, frame) => panic!("frame {frame:?} does not match {message:?}"),
447            }
448        }
449    }
450
451    /// Tests that malformed framing is rejected before any body would be deserialized
452    #[test]
453    fn test_frame_rejects_malformed_messages() {
454        assert!(MessageFrame::parse(&[]).is_err(), "empty message");
455        assert!(MessageFrame::parse(&[2]).is_err(), "unknown type byte");
456        assert!(
457            MessageFrame::parse(&[RESPONSE_TAG; 20]).is_err(),
458            "response shorter than a hash"
459        );
460        assert!(matches!(
461            MessageFrame::parse(&[REQUEST_TAG]),
462            Ok(MessageFrame::Request { body: &[] })
463        ));
464    }
465
466    /// Tests that length-prefixed values are read and written correctly
467    #[test]
468    fn test_length_prefix_parity() {
469        // Create some RNG
470        let mut rng = rand::thread_rng();
471
472        for _ in 0..100 {
473            // Create a buffer to test over
474            let mut bytes = Vec::new();
475
476            // Generate the value to test over
477            let value = vec![rng.r#gen::<u8>(); rng.gen_range(0..10000)];
478
479            // Write the length-prefixed value
480            write_length_prefixed(&mut bytes, &value).unwrap();
481
482            // Create a reader from the bytes
483            let mut reader = Cursor::new(bytes);
484
485            // Read the length-prefixed value
486            let value = read_length_prefixed(&mut reader).unwrap();
487            assert_eq!(value, value);
488        }
489    }
490}
request_response/message.rs

request_response/
message.rs
