grpcio_proto/
util.rs

1// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
2
3use crate::google::rpc::Status;
4use grpcio::{
5    ChannelCredentials, ChannelCredentialsBuilder, ServerCredentials, ServerCredentialsBuilder,
6};
7use std::convert::TryFrom;
8
9#[cfg(all(
10    any(feature = "protobuf-codec", feature = "protobufv3-codec"),
11    not(feature = "prost-codec")
12))]
13use crate::proto::protobuf::testing::messages::{Payload, ResponseParameters};
14#[cfg(feature = "prost-codec")]
15use crate::testing::{Payload, ResponseParameters};
16
17#[cfg(feature = "protobufv3-codec")]
18use protobufv3 as protobuf;
19
20/// Create a payload with the specified size.
21pub fn new_payload(size: usize) -> Payload {
22    Payload {
23        body: vec![0; size],
24        ..Default::default()
25    }
26}
27
28pub fn new_parameters(size: i32) -> ResponseParameters {
29    ResponseParameters {
30        size,
31        ..Default::default()
32    }
33}
34
35pub fn create_test_server_credentials() -> ServerCredentials {
36    let private_key = include_str!("../data/server1.key");
37    let cert = include_str!("../data/server1.pem");
38    ServerCredentialsBuilder::new()
39        .add_cert(cert.into(), private_key.into())
40        .build()
41}
42
43pub fn create_test_channel_credentials() -> ChannelCredentials {
44    let ca = include_str!("../data/ca.pem");
45    ChannelCredentialsBuilder::new()
46        .root_cert(ca.into())
47        .build()
48}
49
50impl TryFrom<grpcio::RpcStatus> for Status {
51    type Error = grpcio::Error;
52
53    fn try_from(value: grpcio::RpcStatus) -> grpcio::Result<Self> {
54        let mut s = Status::default();
55        #[cfg(any(feature = "protobuf-codec", feature = "protobufv3-codec"))]
56        protobuf::Message::merge_from_bytes(&mut s, value.details())?;
57        #[cfg(feature = "prost-codec")]
58        prost::Message::merge(&mut s, value.details())?;
59        if s.code == value.code().into() {
60            if s.message == value.message() {
61                Ok(s)
62            } else {
63                Err(grpcio::Error::Codec(
64                    format!(
65                        "message doesn't match {:?} != {:?}",
66                        s.message,
67                        value.message()
68                    )
69                    .into(),
70                ))
71            }
72        } else {
73            Err(grpcio::Error::Codec(
74                format!("code doesn't match {} != {}", s.code, value.code()).into(),
75            ))
76        }
77    }
78}
79
80impl TryFrom<Status> for grpcio::RpcStatus {
81    type Error = grpcio::Error;
82
83    fn try_from(value: Status) -> grpcio::Result<Self> {
84        #[cfg(any(feature = "protobuf-codec", feature = "protobufv3-codec"))]
85        let details = protobuf::Message::write_to_bytes(&value)?;
86        #[cfg(feature = "prost-codec")]
87        let details = {
88            let mut v = vec![];
89            prost::Message::encode(&value, &mut v).unwrap();
90            v
91        };
92        Ok(grpcio::RpcStatus::with_details(
93            value.code,
94            value.message,
95            details,
96        ))
97    }
98}