1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

use bytes::BytesMut;
use http::{header::CONTENT_TYPE, HeaderValue};
use prost::Message;
use warp::{reply::Response, Reply};

use super::status::Status;

/// If a type fails to be encoded as Protobuf, the error is logged at the
/// `error` level, and the returned `impl Reply` will be an empty
/// `500 Internal Server Error` response.
pub fn protobuf<T>(val: T) -> Protobuf
where
    T: Message,
{
    let mut buf = BytesMut::with_capacity(1024);
    Protobuf {
        inner: val.encode(&mut buf).map(|_| buf.to_vec()).map_err(|err| {
            error!("Failed to encode value: {}", err);
        }),
    }
}

/// A Protobuf formatted reply.
#[allow(missing_debug_implementations)]
pub struct Protobuf {
    inner: Result<Vec<u8>, ()>,
}

impl Reply for Protobuf {
    #[inline]
    fn into_response(self) -> Response {
        match self.inner {
            Ok(body) => {
                let mut res = Response::new(body.into());
                res.headers_mut().insert(
                    CONTENT_TYPE,
                    HeaderValue::from_static("application/x-protobuf"),
                );
                res
            }
            Err(()) => {
                let status = Status {
                    message: "Failed to encode error message".into(),
                    ..Default::default()
                };
                let mut res = Response::new(status.encode_to_vec().into());
                res.headers_mut().insert(
                    CONTENT_TYPE,
                    HeaderValue::from_static("application/x-protobuf"),
                );
                res
            }
        }
    }
}