tenebrous-sheets/src/models/proto.rs

81 lines
2.5 KiB
Rust

use crate::errors::Error;
use rocket::data::{Data, FromData, Outcome, ToByteUnit};
use rocket::request::Request;
use std::default::Default;
use std::ops::Deref;
pub mod cofd;
const CRATE_NAME: &'static str = env!("CARGO_BIN_NAME");
/// Convert an incoming protobuf content-type to the equivalent type
/// name produced by std::any::type_name(). Currently does NOT work
/// with nested types due to how prost generates the module names.
fn convert_to_rust_name(proto_type: &str) -> String {
format!("{}::{}", CRATE_NAME, proto_type.replace(".", "::"))
}
/// A struct wrapping a protobuf that allows it to be used as binary
/// data submitted via POST using fetch API. Can automatically be
/// dereferenced into its wrapped type.
#[derive(Debug)]
pub(crate) struct Proto<T>(T)
where
T: prost::Message + Default;
/// Converts the body of a POST request containing encoded protobuf
/// data into the wrapped type.
#[rocket::async_trait]
impl<T> FromData for Proto<T>
where
T: prost::Message + Default,
{
type Error = crate::errors::Error;
async fn from_data(req: &Request<'_>, data: Data) -> Outcome<Self, Error> {
use rocket::http::Status;
let content_type = req.content_type();
let is_protobuf = content_type
.map(|ct| ct.top() == "application" && ct.sub() == "x-protobuf")
.unwrap_or(false);
let message_type: Option<String> = content_type.and_then(|ct| {
ct.params()
.find(|&(name, _)| name == "messageType")
.map(|(_, message_type)| convert_to_rust_name(message_type))
});
if !is_protobuf {
return Outcome::Failure((Status::new(422, "invalid protobuf"), Error::InvalidInput));
}
if message_type.as_ref().map(String::as_str) != Some(std::any::type_name::<T>()) {
println!("message type is {:?}", message_type);
return Outcome::Forward(data);
}
let bytes: Vec<u8> = match data.open(2.mebibytes()).stream_to_vec().await {
Ok(read_bytes) => read_bytes,
Err(e) => return Outcome::Failure((Status::new(422, "invalid protobuf"), e.into())),
};
match T::decode(bytes.as_ref()) {
Ok(decoded) => Outcome::Success(Proto(decoded)),
Err(e) => Outcome::Failure((Status::new(422, "invalid protobuf"), e.into())),
}
}
}
/// Enable automatically calling methods on a decoded Proto instance.
impl<T> Deref for Proto<T>
where
T: prost::Message + Default,
{
type Target = T;
fn deref(&self) -> &T {
&self.0
}
}