wiremock_grpc/wiremock/
tonic_ext.rs

1use prost::{bytes::BufMut, Message};
2
3use tonic::{codec::Codec, Code};
4
5pub(crate) struct GenericSvc(pub(crate) Vec<u8>);
6impl tonic::server::UnaryService<Vec<u8>> for GenericSvc {
7    type Response = Vec<u8>;
8    type Future = tonic::codegen::BoxFuture<tonic::Response<Self::Response>, tonic::Status>;
9    fn call(&mut self, _: tonic::Request<Vec<u8>>) -> Self::Future {
10        let body = self.0.clone();
11        let fut = async move { Ok(tonic::Response::new(body)) };
12
13        Box::pin(fut)
14    }
15}
16
17#[derive(Default)]
18pub(crate) struct GenericCodec;
19
20impl Codec for GenericCodec {
21    type Encode = Vec<u8>;
22    type Decode = Vec<u8>;
23
24    type Encoder = GenericProstEncoder;
25    type Decoder = GenericProstDecoder;
26
27    fn encoder(&mut self) -> Self::Encoder {
28        GenericProstEncoder {}
29    }
30
31    fn decoder(&mut self) -> Self::Decoder {
32        GenericProstDecoder {}
33    }
34}
35
36/// A [`Encoder`] that knows how to encode `T`.
37#[derive(Debug, Clone, Default)]
38pub struct GenericProstEncoder;
39
40impl tonic::codec::Encoder for GenericProstEncoder {
41    type Item = Vec<u8>;
42    type Error = tonic::Status;
43
44    fn encode(
45        &mut self,
46        item: Self::Item,
47        buf: &mut tonic::codec::EncodeBuf<'_>,
48    ) -> Result<(), Self::Error> {
49        // copy the respone body `item` to the buffer
50        for i in item {
51            buf.put_u8(i);
52        }
53
54        Ok(())
55    }
56}
57
58/// A [`Decoder`] that knows how to decode `U`.
59#[derive(Debug, Clone, Default)]
60pub struct GenericProstDecoder;
61
62impl tonic::codec::Decoder for GenericProstDecoder {
63    type Item = Vec<u8>;
64    type Error = tonic::Status;
65
66    fn decode(
67        &mut self,
68        buf: &mut tonic::codec::DecodeBuf<'_>,
69    ) -> Result<Option<Self::Item>, Self::Error> {
70        let item = Message::decode(buf)
71            .map(Option::Some)
72            .map_err(from_decode_error)?;
73
74        Ok(item)
75    }
76}
77
78fn from_decode_error(error: prost::DecodeError) -> tonic::Status {
79    // Map Protobuf parse errors to an INTERNAL status code, as per
80    // https://github.com/grpc/grpc/blob/master/doc/statuscodes.md
81    tonic::Status::new(Code::Internal, error.to_string())
82}