ntex_grpc/client/
transport.rs

1use std::{convert::TryFrom, str::FromStr};
2
3use ntex_bytes::{Buf, BufMut, BytePages};
4use ntex_error::Error;
5use ntex_h2::{self as h2};
6use ntex_http::{HeaderMap, Method, header};
7
8use super::{Client, ClientError, Transport, request::RequestContext, request::Response};
9use crate::{DecodeError, GrpcStatus, Message, consts, service::MethodDef, utils::Data};
10
11impl<T: MethodDef> Transport<T> for Client {
12    type Error = Error<ClientError>;
13
14    #[inline]
15    async fn request(
16        &self,
17        val: &T::Input,
18        ctx: &mut RequestContext,
19    ) -> Result<Response<T>, Self::Error> {
20        Transport::request(&self.0, val, ctx).await
21    }
22}
23
24impl<T: MethodDef> Transport<T> for h2::client::Client {
25    type Error = Error<ClientError>;
26
27    #[inline]
28    async fn request(
29        &self,
30        val: &T::Input,
31        ctx: &mut RequestContext,
32    ) -> Result<Response<T>, Self::Error> {
33        Transport::request(
34            &self.client().await.map_err(|e| e.map(ClientError::from))?,
35            val,
36            ctx,
37        )
38        .await
39    }
40}
41
42impl<T: MethodDef> Transport<T> for h2::client::SimpleClient {
43    type Error = Error<ClientError>;
44
45    #[allow(clippy::too_many_lines)]
46    async fn request(
47        &self,
48        val: &T::Input,
49        ctx: &mut RequestContext,
50    ) -> Result<Response<T>, Self::Error> {
51        let len = val.encoded_len();
52        let mut buf = BytePages::default();
53        buf.put_u8(0); // compression
54        buf.put_u32(len as u32); // length
55        val.write(&mut buf);
56        let req_size = buf.len();
57
58        let mut hdrs = HeaderMap::new();
59        hdrs.append(header::CONTENT_TYPE, consts::HDRV_CT_GRPC);
60        hdrs.append(header::USER_AGENT, consts::HDRV_USER_AGENT);
61        hdrs.insert(header::TE, consts::HDRV_TRAILERS);
62        hdrs.insert(consts::GRPC_ENCODING, consts::IDENTITY);
63        hdrs.insert(consts::GRPC_ACCEPT_ENCODING, consts::IDENTITY);
64        for (key, val) in ctx.headers() {
65            hdrs.insert(key.clone(), val.clone());
66        }
67
68        // send request
69        let (snd_stream, rcv_stream) = self
70            .send(Method::POST, T::PATH, hdrs, false)
71            .await
72            .map_err(|e| e.map(ClientError::from))?;
73        if ctx.get_disconnect_on_drop() {
74            snd_stream.disconnect_on_drop();
75        }
76        snd_stream
77            .send_pages(buf, true)
78            .await
79            .map_err(|e| e.map(ClientError::from))?;
80
81        // read response
82        let mut status = None;
83        let mut hdrs = HeaderMap::default();
84        let mut trailers = HeaderMap::default();
85        let mut payload = Data::Empty;
86
87        async {
88            loop {
89                let Some(msg) = rcv_stream.recv().await else {
90                    return Err(Error::from(ClientError::UnexpectedEof(status, hdrs)));
91                };
92
93                match msg.kind {
94                    h2::MessageKind::Headers {
95                        headers,
96                        pseudo,
97                        eof,
98                    } => {
99                        if eof {
100                            // check grpc status
101                            match check_grpc_status(&headers) {
102                                Some(Ok(GrpcStatus::DeadlineExceeded)) => {
103                                    return Err(Error::from(ClientError::DeadlineExceeded(hdrs)));
104                                }
105                                Some(Ok(status)) if status != GrpcStatus::Ok => {
106                                    return Err(Error::from(ClientError::GrpcStatus(
107                                        status, headers,
108                                    )));
109                                }
110                                Some(Err(())) => {
111                                    return Err(Error::from(ClientError::Decode(
112                                        DecodeError::new("Cannot parse grpc status"),
113                                    )));
114                                }
115                                Some(Ok(_)) | None => {}
116                            }
117
118                            return Err(Error::from(ClientError::UnexpectedEof(
119                                pseudo.status,
120                                headers,
121                            )));
122                        }
123                        hdrs = headers;
124                        status = pseudo.status;
125                        continue;
126                    }
127                    h2::MessageKind::Data(data, _cap) => {
128                        payload.push(data);
129                        continue;
130                    }
131                    h2::MessageKind::Eof(data) => {
132                        match data {
133                            h2::StreamEof::Data(data) => {
134                                payload.push(data);
135                            }
136                            h2::StreamEof::Trailers(hdrs) => {
137                                // check grpc status
138                                match check_grpc_status(&hdrs) {
139                                    Some(Ok(GrpcStatus::Ok)) | None => Ok(()),
140                                    Some(Ok(GrpcStatus::DeadlineExceeded)) => {
141                                        return Err(Error::from(ClientError::DeadlineExceeded(
142                                            hdrs,
143                                        )));
144                                    }
145                                    Some(Ok(st)) => {
146                                        return Err(Error::from(ClientError::GrpcStatus(
147                                            st, hdrs,
148                                        )));
149                                    }
150                                    Some(Err(())) => Err(Error::from(ClientError::Decode(
151                                        DecodeError::new("Cannot parse grpc status"),
152                                    ))),
153                                }?;
154                                trailers = hdrs;
155                            }
156                            h2::StreamEof::Error(err) => {
157                                return Err(err.map(ClientError::Stream));
158                            }
159                        }
160                    }
161                    h2::MessageKind::Disconnect(err) => {
162                        return Err(err.map(ClientError::Operation));
163                    }
164                }
165
166                let mut data = payload.get();
167                match status {
168                    Some(st) => {
169                        if !st.is_success() {
170                            return Err(Error::from(ClientError::Response(Some(st), hdrs, data)));
171                        }
172                    }
173                    None => return Err(Error::from(ClientError::Response(None, hdrs, data))),
174                }
175                let _compressed = data.get_u8();
176                let len = data.get_u32();
177                let Some(mut block) = data.split_to_checked(len as usize) else {
178                    return Err(Error::from(ClientError::UnexpectedEof(None, hdrs)));
179                };
180
181                return match <T::Output as Message>::read(&mut block) {
182                    Ok(output) => Ok(Response {
183                        output,
184                        trailers,
185                        req_size,
186                        headers: hdrs,
187                        res_size: data.len(),
188                    }),
189                    Err(e) => Err(Error::from(ClientError::Decode(e))),
190                };
191            }
192        }
193        .await
194        .map_err(|e| e.set_service(self.service()))
195    }
196}
197
198fn check_grpc_status(hdrs: &HeaderMap) -> Option<Result<GrpcStatus, ()>> {
199    // check grpc status
200    if let Some(val) = hdrs.get(consts::GRPC_STATUS) {
201        if let Ok(status) = val
202            .to_str()
203            .map_err(|_| ())
204            .and_then(|v| u8::from_str(v).map_err(|_| ()))
205            .and_then(GrpcStatus::try_from)
206        {
207            Some(Ok(status))
208        } else {
209            Some(Err(()))
210        }
211    } else {
212        None
213    }
214}
ntex_grpc/client/transport.rs

ntex_grpc/client/
transport.rs
