axol 0.2.0

Axol Web Framework
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

use axol_http::{Body, request::RequestPartsRef, response::Response};
use prost::Message;

use crate::{Error, FromRequest, FromRequestParts, IntoResponse, Result, Typed, grpc::Status};

use super::{Encoding, GrpcContentType};

pub struct Grpc<T: Message> {
    pub encoding: Encoding,
    pub content_type: GrpcContentType,
    pub body: T,
}

#[async_trait::async_trait]
impl<'a, T: Default + Message + 'a> FromRequest<'a> for Grpc<T> {
    async fn from_request(request: RequestPartsRef<'a>, body: Body) -> Result<Self> {
        let content_type = Typed::<GrpcContentType>::from_request_parts(request)
            .await?
            .0;
        let encoding = Option::<Typed<Encoding>>::from_request_parts(request)
            .await?
            .map(|x| x.0)
            .unwrap_or(Encoding::Identity);

        if content_type == GrpcContentType::Invalid {
            return Err(Error::UnsupportedMediaType);
        } else if content_type != GrpcContentType::Proto {
            return Err(Error::GrpcMessage(
                Status::Internal,
                "only proto subtype is supported".to_string(),
            ));
        }

        if encoding != Encoding::Identity {
            //TODO: support more encodings
            return Err(Error::GrpcMessage(
                Status::Internal,
                "only identity encoding is supported".to_string(),
            ));
        }

        let body = body.collect().await?;

        if body.len() < 5 {
            return Err(Error::GrpcMessage(
                Status::Internal,
                "truncated body".to_string(),
            ));
        }
        let is_compressed = match body[0] {
            0 => false,
            1 => true,
            _ => {
                return Err(Error::GrpcMessage(
                    Status::Internal,
                    "invalid compression flag".to_string(),
                ));
            }
        };
        if is_compressed != (encoding != Encoding::Identity) {
            return Err(Error::GrpcMessage(
                Status::Internal,
                "compression flag mismatch with headers".to_string(),
            ));
        }
        let length = u32::from_be_bytes((&body[1..5]).try_into().unwrap());
        //TODO: should we assert length == body.len() - 5?
        let body: T = T::decode(&body[5..5 + length as usize])
            .map_err(|e| Error::GrpcMessage(Status::Internal, format!("decode failure: {e}")))?;
        Ok(Self {
            encoding,
            content_type,
            body,
        })
    }
}

impl<T: Message + Default> IntoResponse for Grpc<T> {
    fn into_response(self) -> Result<Response> {
        let is_compressed = self.encoding != Encoding::Identity;

        let mut out = vec![is_compressed as u8, 0, 0, 0, 0];
        self.body
            .encode(&mut out)
            .map_err(|e| Error::GrpcMessage(Status::Internal, format!("encode failure: {e}")))?;
        let out_len = (out.len() - 5) as u32;
        out[1..5].copy_from_slice(&out_len.to_be_bytes()[..]);
        // must return a Body::Bytes or we will break downstream stuff
        let mut response = Response::new(Body::Bytes(out));
        response.headers.insert_typed(&self.content_type);
        response.headers.insert_typed(&self.encoding);
        Ok(response)
    }
}