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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
use std::borrow::Cow;
use std::fmt::{Debug, Formatter};
use std::net::SocketAddr;
use std::ops::Deref;

use async_trait::async_trait;
use bytecheck::CheckBytes;
use rkyv::validation::validators::DefaultValidator;
use rkyv::{Archive, Deserialize, Serialize};

use crate::view::DataView;
use crate::{Body, Status};

#[async_trait]
/// The deserializer trait for converting the request body into
/// the desired type specified by [Self::Content].
///
/// This trait is automatically implemented for the [Body] type
/// and any type implementing [rkyv]'s (de)serializer traits.
pub trait RequestContents {
    /// The deserialized message type.
    type Content: Send + Sized + 'static;

    async fn from_body(body: Body) -> Result<Self::Content, Status>;
}

#[async_trait]
impl RequestContents for Body {
    type Content = Self;

    async fn from_body(body: Body) -> Result<Self, Status> {
        Ok(body)
    }
}

#[async_trait]
impl<Msg> RequestContents for Msg
where
    Msg: Archive + Send + Sync + 'static,
    Msg::Archived: CheckBytes<DefaultValidator<'static>> + Send + Sync + 'static,
{
    type Content = DataView<Self>;

    async fn from_body(body: Body) -> Result<Self::Content, Status> {
        let bytes = crate::utils::to_aligned(body.0)
            .await
            .map_err(Status::internal)?;

        DataView::using(bytes).map_err(|_| Status::invalid())
    }
}

#[repr(C)]
#[derive(Serialize, Deserialize, Archive, PartialEq)]
#[cfg_attr(test, derive(Debug))]
#[archive_attr(derive(CheckBytes))]
pub struct MessageMetadata {
    #[with(rkyv::with::AsOwned)]
    /// The name of the service being targeted.
    pub(crate) service_name: Cow<'static, str>,
    #[with(rkyv::with::AsOwned)]
    /// The message name/path.
    pub(crate) path: Cow<'static, str>,
}

/// A zero-copy view of the message data and any additional metadata provided
/// by the RPC system.
///
/// The request contains the original request buffer which is used to create
/// the 'view' of the given message type.
pub struct Request<Msg>
where
    Msg: RequestContents,
{
    pub(crate) remote_addr: SocketAddr,

    // A small hack to stop linters miss-guiding users
    // into thinking their messages are `!Sized` when in fact they are.
    // We don't want to box in release mode however.
    #[cfg(debug_assertions)]
    pub(crate) view: Box<Msg::Content>,
    #[cfg(not(debug_assertions))]
    pub(crate) view: Msg::Content,
}

impl<Msg> Debug for Request<Msg>
where
    Msg: RequestContents,
    Msg::Content: Debug,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Request")
            .field("view", &self.view)
            .field("remote_addr", &self.remote_addr)
            .finish()
    }
}

impl<Msg> Deref for Request<Msg>
where
    Msg: RequestContents,
{
    type Target = Msg::Content;

    fn deref(&self) -> &Self::Target {
        &self.view
    }
}

impl<Msg> Request<Msg>
where
    Msg: RequestContents,
{
    pub(crate) fn new(remote_addr: SocketAddr, view: Msg::Content) -> Self {
        Self {
            remote_addr,
            #[cfg(debug_assertions)]
            view: Box::new(view),
            #[cfg(not(debug_assertions))]
            view,
        }
    }

    #[cfg(debug_assertions)]
    /// Consumes the request into the value of the message.
    pub fn into_inner(self) -> Msg::Content {
        *self.view
    }

    #[cfg(not(debug_assertions))]
    /// Consumes the request into the value of the message.
    pub fn into_inner(self) -> Msg::Content {
        self.view
    }

    /// The remote address of the incoming message.
    pub fn remote_addr(&self) -> SocketAddr {
        self.remote_addr
    }
}

#[cfg(feature = "test-utils")]
impl<Msg> Request<Msg>
where
    Msg: RequestContents
        + rkyv::Serialize<
            rkyv::ser::serializers::AllocSerializer<{ crate::SCRATCH_SPACE }>,
        >,
{
    /// A test utility for creating a mocked request.
    ///
    /// This takes the owned value of the msg and acts like the target request.
    ///
    /// This should be used for testing only.
    pub async fn using_owned(msg: Msg) -> Self {
        let bytes = rkyv::to_bytes(&msg).unwrap();
        let contents = Msg::from_body(Body::from(bytes.to_vec())).await.unwrap();

        use std::net::{Ipv4Addr, SocketAddrV4};

        let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::from([127, 0, 0, 1]), 80));
        Self::new(addr, contents)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_metadata() {
        let meta = MessageMetadata {
            service_name: Cow::Borrowed("test"),
            path: Cow::Borrowed("demo"),
        };

        let bytes = rkyv::to_bytes::<_, 1024>(&meta).expect("Serialize");
        let copy: MessageMetadata = rkyv::from_bytes(&bytes).expect("Deserialize");
        assert_eq!(meta, copy, "Deserialized value should match");
    }

    #[test]
    fn test_request() {
        let msg = MessageMetadata {
            service_name: Cow::Borrowed("test"),
            path: Cow::Borrowed("demo"),
        };

        let addr = "127.0.0.1:8000".parse().unwrap();
        let bytes = rkyv::to_bytes::<_, 1024>(&msg).expect("Serialize");
        let view: DataView<MessageMetadata, _> =
            DataView::using(bytes).expect("Create view");
        let req = Request::<MessageMetadata>::new(addr, view);
        assert_eq!(req.remote_addr(), addr, "Remote addr should match.");
        assert_eq!(
            req.to_owned().unwrap(),
            msg,
            "Deserialized value should match."
        );
    }
}