Compare commits

...

2 commits

Author SHA1 Message Date
c435f8038d refactor into dedicated function 2023-04-04 22:22:18 -03:00
d33ab7540c add code that mirrors request to deepdanbooru instance 2023-04-04 22:16:13 -03:00
2 changed files with 137 additions and 3 deletions

View file

@ -6,5 +6,12 @@ edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies] [dependencies]
axum = { version = "0.6.12", features = ["tokio"] } anyhow = "1.0.70"
axum = { version = "0.6.12", features = ["tokio", "multipart"] }
axum-macros = "0.3.7"
log = "0.4.17"
pretty_env_logger = "0.4.0"
reqwest = { version = "0.11.16", features = ["json", "multipart"] }
serde = { version = "1.0.159", features = ["derive"] }
serde_json = "1.0.95"
tokio = { version = "1.27.0", features = ["full"] } tokio = { version = "1.27.0", features = ["full"] }

View file

@ -1,9 +1,17 @@
use axum::{routing::get, Router}; use axum::body::Bytes;
use axum::extract::{Multipart, Query};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::post;
use axum::{Json, Router};
use axum_macros::debug_handler;
use serde::{Deserialize, Serialize};
#[tokio::main] #[tokio::main]
async fn main() { async fn main() {
pretty_env_logger::init();
// build our application with a single route // build our application with a single route
let app = Router::new().route("/", get(|| async { "Hello, World!" })); let app = Router::new().route("/", post(upload_file));
// run it with hyper on localhost:3000 // run it with hyper on localhost:3000
axum::Server::bind(&"0.0.0.0:6679".parse().unwrap()) axum::Server::bind(&"0.0.0.0:6679".parse().unwrap())
@ -11,3 +19,122 @@ async fn main() {
.await .await
.unwrap(); .unwrap();
} }
#[derive(Deserialize)]
struct Options {
threshold: String,
}
#[derive(Serialize, Deserialize)]
#[serde(untagged)]
enum WrappedResponse {
Tags(Vec<String>),
Error(String),
}
// Make our own error that wraps `anyhow::Error`.
struct AppError(anyhow::Error);
// Tell axum how to convert `AppError` into a response.
impl IntoResponse for AppError {
fn into_response(self) -> Response {
println!("amogus");
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("Something went wrong: {}", self.0),
)
.into_response()
}
}
// This enables using `?` on functions that return `Result<_, anyhow::Error>` to turn them into
// `Result<_, AppError>`. That way you don't need to do that manually.
impl<E> From<E> for AppError
where
E: Into<anyhow::Error>,
{
fn from(err: E) -> Self {
Self(err.into())
}
}
fn try_thing() -> Result<(), anyhow::Error> {
anyhow::bail!("it failed!")
}
async fn test_handler() -> Result<(), AppError> {
try_thing()?;
Ok(())
}
async fn send_image_to_dd(
file_contents: Vec<u8>,
file_name: String,
file_mime_type: &str,
options: &Options,
) -> anyhow::Result<WrappedResponse> {
let part = reqwest::multipart::Part::bytes(file_contents)
.file_name(file_name)
.mime_str(file_mime_type)
.unwrap();
let form = reqwest::multipart::Form::new().part("file", part);
log::debug!("calling dd");
let resp = reqwest::Client::new()
.post("http://localhost:4443")
.multipart(form)
.header("authorization", "Bearer 123")
.query(&[("threshold", options.threshold.clone())])
.send()
.await?;
let body = resp.text().await?;
log::info!("body: {}", &body);
let json_response: WrappedResponse = serde_json::from_str(&body)?;
log::debug!("called!");
Ok(json_response)
}
#[debug_handler]
async fn upload_file(
options: Query<Options>,
mut multipart: Multipart,
) -> Result<(StatusCode, Json<WrappedResponse>), AppError> {
let mut maybe_file_contents: Option<axum::body::Bytes> = None;
let mut maybe_file_type: Option<String> = None;
let mut maybe_file_name: Option<String> = None;
while let Some(field) = multipart.next_field().await.unwrap() {
let name = field.name().unwrap().to_string();
let content_type = field.content_type().unwrap().to_string();
let filename = field.file_name().unwrap().to_string();
let data = field.bytes().await.unwrap();
log::info!("Length of `{}` is {} bytes", name, data.len());
if name == "file" {
maybe_file_contents = Some(data);
maybe_file_type = Some(content_type);
maybe_file_name = Some(filename);
}
}
if let Some(file_contents) = maybe_file_contents {
let json_response = send_image_to_dd(
file_contents.to_vec(),
maybe_file_name.unwrap(),
&maybe_file_type.unwrap(),
&options,
)
.await?;
Ok((StatusCode::OK, Json(json_response)))
} else {
Ok((
StatusCode::INTERNAL_SERVER_ERROR,
Json(WrappedResponse::Error(
"no file found in request".to_string(),
)),
))
}
}