53 lines
1.5 KiB
Rust
53 lines
1.5 KiB
Rust
use prost::Message;
|
|
use rocket::Request;
|
|
use rocket::data::{Data, FromData, Outcome, ToByteUnit};
|
|
use rocket::http::{ContentType, Status};
|
|
use std::ops::{Deref, DerefMut};
|
|
|
|
pub struct Proto<T>(pub T);
|
|
|
|
impl<T> Proto<T> {
|
|
pub fn into_inner(self) -> T {
|
|
self.0
|
|
}
|
|
}
|
|
|
|
impl<T> Deref for Proto<T> {
|
|
type Target = T;
|
|
|
|
fn deref(&self) -> &Self::Target {
|
|
&self.0
|
|
}
|
|
}
|
|
|
|
impl<T> DerefMut for Proto<T> {
|
|
fn deref_mut(&mut self) -> &mut Self::Target {
|
|
&mut self.0
|
|
}
|
|
}
|
|
|
|
#[rocket::async_trait]
|
|
impl<'r, T: Message + Default> FromData<'r> for Proto<T> {
|
|
type Error = String;
|
|
|
|
async fn from_data(req: &'r Request<'_>, data: Data<'r>) -> Outcome<'r, Self> {
|
|
if req.content_type() != Some(&ContentType::new("application", "protobuf"))
|
|
&& req.content_type() != Some(&ContentType::new("application", "octet-stream"))
|
|
{
|
|
return Outcome::Forward((data, Status::NotFound));
|
|
}
|
|
|
|
let limit = req.limits().get("protobuf").unwrap_or(1.mebibytes());
|
|
let bytes = match data.open(limit).into_bytes().await {
|
|
Ok(bytes) if bytes.is_complete() => bytes.into_inner(),
|
|
Ok(_) => return Outcome::Error((Status::PayloadTooLarge, "Payload too large".into())),
|
|
Err(e) => return Outcome::Error((Status::InternalServerError, e.to_string())),
|
|
};
|
|
|
|
match T::decode(&bytes[..]) {
|
|
Ok(msg) => Outcome::Success(Proto(msg)),
|
|
Err(e) => Outcome::Error((Status::UnprocessableEntity, e.to_string())),
|
|
}
|
|
}
|
|
}
|