use bytes::{BufMut, Bytes, BytesMut};
use http::status::InvalidStatusCode;
use prost::Message;
pub trait CodeExt {
fn from_u16(code: u16) -> Result<tonic::Code, InvalidStatusCode>;
fn from_http(status_code: http::StatusCode) -> tonic::Code;
}
impl CodeExt for tonic::Code {
fn from_u16(code: u16) -> Result<tonic::Code, InvalidStatusCode> {
let status_code = http::StatusCode::from_u16(code)?;
Ok(tonic::Code::from_http(status_code))
}
fn from_http(status_code: http::StatusCode) -> tonic::Code {
match status_code {
http::StatusCode::INTERNAL_SERVER_ERROR => tonic::Code::Internal,
http::StatusCode::UNPROCESSABLE_ENTITY | http::StatusCode::BAD_REQUEST => {
tonic::Code::InvalidArgument
}
http::StatusCode::UNAUTHORIZED => tonic::Code::Unauthenticated,
http::StatusCode::FORBIDDEN => tonic::Code::PermissionDenied,
http::StatusCode::NOT_FOUND => tonic::Code::NotFound,
http::StatusCode::NOT_IMPLEMENTED => tonic::Code::Unimplemented,
http::StatusCode::TOO_MANY_REQUESTS
| http::StatusCode::BAD_GATEWAY
| http::StatusCode::SERVICE_UNAVAILABLE
| http::StatusCode::GATEWAY_TIMEOUT => tonic::Code::Unavailable,
http::StatusCode::OK => tonic::Code::Ok,
_ => tonic::Code::Unknown,
}
}
}
pub trait MessageExt {
fn to_bytes(&self) -> Bytes;
}
impl<T: Message> MessageExt for T {
fn to_bytes(&self) -> Bytes {
let mut buf = BytesMut::with_capacity(256);
buf.reserve(5);
unsafe {
buf.advance_mut(5);
}
self.encode(&mut buf).unwrap();
{
let len = buf.len() - 5;
let mut buf = &mut buf[..5];
buf.put_u8(0); buf.put_u32(len as u32);
}
buf.freeze()
}
}