use crate::errors::Error; use prost::bytes::BytesMut; use rocket::data::{Data, FromData, Outcome, ToByteUnit}; use rocket::http::{ContentType, Status}; use rocket::request::Request; use rocket::response::status; use rocket::response::{self, Responder, Response}; use std::default::Default; use std::io::Cursor; use std::ops::{Deref, DerefMut}; pub mod cofd; /// 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(pub 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 FromData for Proto where T: prost::Message + Default, { type Error = crate::errors::Error; async fn from_data(req: &Request<'_>, data: Data) -> Outcome { 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); if !is_protobuf { return Outcome::Failure((Status::new(422, "invalid protobuf"), Error::InvalidInput)); } let bytes: Vec = 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())), } } } impl<'r, T> Responder<'r, 'static> for Proto where T: prost::Message + Default, { fn respond_to(self, req: &Request) -> response::Result<'static> { let mut buf = BytesMut::with_capacity(std::mem::size_of_val(&self.0)); match self.0.encode(&mut buf) { Ok(_) => Response::build() .header(ContentType::new("application", "x-protobuf")) .sized_body(buf.len(), Cursor::new(buf)) .ok(), Err(e) => status::Custom(Status::InternalServerError, e.to_string()).respond_to(req), } } } /// Enable automatically calling methods on a decoded Proto instance. impl Deref for Proto where T: prost::Message + Default, { type Target = T; fn deref(&self) -> &T { &self.0 } } impl DerefMut for Proto where T: prost::Message + Default, { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.0 } }