diff options
-rw-r--r-- | src/db.rs | 2 | ||||
-rw-r--r-- | src/service.rs | 26 |
2 files changed, 19 insertions, 9 deletions
@@ -11,7 +11,7 @@ pub fn open_connection<P>(p: P) -> Result<Connection> | |||
11 | where | 11 | where |
12 | P: AsRef<Path> + fmt::Display, | 12 | P: AsRef<Path> + fmt::Display, |
13 | { | 13 | { |
14 | info!("Opened connection to databse"); | 14 | info!("Opened connection to database"); |
15 | Ok(Connection::open_with_flags( | 15 | Ok(Connection::open_with_flags( |
16 | &p, | 16 | &p, |
17 | OpenFlags::SQLITE_OPEN_CREATE | OpenFlags::SQLITE_OPEN_READ_WRITE, | 17 | OpenFlags::SQLITE_OPEN_CREATE | OpenFlags::SQLITE_OPEN_READ_WRITE, |
diff --git a/src/service.rs b/src/service.rs index 4477ee8..35471ea 100644 --- a/src/service.rs +++ b/src/service.rs | |||
@@ -1,6 +1,6 @@ | |||
1 | // extern | 1 | // extern |
2 | use anyhow::{Context, Result}; | 2 | use anyhow::{Context, Result}; |
3 | use hyper::header::CONTENT_TYPE; | 3 | use hyper::header::{HeaderValue, CONTENT_TYPE}; |
4 | use hyper::{Body, Method, Request, Response, StatusCode}; | 4 | use hyper::{Body, Method, Request, Response, StatusCode}; |
5 | use log::{error, info, trace}; | 5 | use log::{error, info, trace}; |
6 | use multer::Multipart; | 6 | use multer::Multipart; |
@@ -11,12 +11,18 @@ use url::form_urlencoded; | |||
11 | // std | 11 | // std |
12 | use std::collections::HashMap; | 12 | use std::collections::HashMap; |
13 | 13 | ||
14 | fn respond_with_shortlink<S: AsRef<str>>(shortlink: S) -> Response<Body> { | 14 | fn get_host(req: &Request<Body>) -> Option<HeaderValue> { |
15 | let host = req.headers().get("host").map(|h| h.clone())?; | ||
16 | return Some(host); | ||
17 | } | ||
18 | |||
19 | fn respond_with_shortlink<S: AsRef<[u8]>>(shortlink: S, host: &[u8]) -> Response<Body> { | ||
20 | let url = [host, b"/", shortlink.as_ref()].concat(); | ||
15 | info!("Successfully generated shortlink"); | 21 | info!("Successfully generated shortlink"); |
16 | Response::builder() | 22 | Response::builder() |
17 | .status(StatusCode::OK) | 23 | .status(StatusCode::OK) |
18 | .header("content-type", "text/html") | 24 | .header("content-type", "text/html") |
19 | .body(Body::from(shortlink.as_ref().to_string())) | 25 | .body(Body::from(url)) |
20 | .unwrap() | 26 | .unwrap() |
21 | } | 27 | } |
22 | 28 | ||
@@ -51,11 +57,13 @@ fn get_link<S: AsRef<str>>(url: S, conn: &Connection) -> Result<Option<String>> | |||
51 | } | 57 | } |
52 | 58 | ||
53 | async fn process_multipart( | 59 | async fn process_multipart( |
54 | body: Body, | 60 | req: Request<Body>, |
55 | boundary: String, | 61 | boundary: String, |
56 | conn: &mut Connection, | 62 | conn: &mut Connection, |
57 | ) -> Result<Response<Body>> { | 63 | ) -> Result<Response<Body>> { |
58 | let mut m = Multipart::new(body, boundary); | 64 | let _h = get_host(&req); |
65 | let host = _h.as_ref().map(|h| h.as_bytes()).unwrap_or(b""); | ||
66 | let mut m = Multipart::new(req.into_body(), boundary); | ||
59 | if let Some(field) = m.next_field().await? { | 67 | if let Some(field) = m.next_field().await? { |
60 | if field.name() == Some("shorten") { | 68 | if field.name() == Some("shorten") { |
61 | trace!("Recieved valid multipart request"); | 69 | trace!("Recieved valid multipart request"); |
@@ -65,7 +73,7 @@ async fn process_multipart( | |||
65 | .with_context(|| format!("Expected field name"))?; | 73 | .with_context(|| format!("Expected field name"))?; |
66 | 74 | ||
67 | let shortlink = shorten(content, conn)?; | 75 | let shortlink = shorten(content, conn)?; |
68 | return Ok(respond_with_shortlink(shortlink)); | 76 | return Ok(respond_with_shortlink(shortlink, host)); |
69 | } | 77 | } |
70 | } | 78 | } |
71 | trace!("Unprocessable multipart request!"); | 79 | trace!("Unprocessable multipart request!"); |
@@ -73,6 +81,8 @@ async fn process_multipart( | |||
73 | } | 81 | } |
74 | 82 | ||
75 | async fn process_form(req: Request<Body>, conn: &mut Connection) -> Result<Response<Body>> { | 83 | async fn process_form(req: Request<Body>, conn: &mut Connection) -> Result<Response<Body>> { |
84 | let _h = get_host(&req); | ||
85 | let host = _h.as_ref().map(|h| h.as_bytes()).unwrap_or(b""); | ||
76 | let b = hyper::body::to_bytes(req) | 86 | let b = hyper::body::to_bytes(req) |
77 | .await | 87 | .await |
78 | .with_context(|| format!("Failed to stream request body!"))?; | 88 | .with_context(|| format!("Failed to stream request body!"))?; |
@@ -84,7 +94,7 @@ async fn process_form(req: Request<Body>, conn: &mut Connection) -> Result<Respo | |||
84 | if let Some(n) = params.get("shorten") { | 94 | if let Some(n) = params.get("shorten") { |
85 | trace!("POST: {}", &n); | 95 | trace!("POST: {}", &n); |
86 | let s = shorten(n, conn)?; | 96 | let s = shorten(n, conn)?; |
87 | return Ok(respond_with_shortlink(s)); | 97 | return Ok(respond_with_shortlink(s, host)); |
88 | } else { | 98 | } else { |
89 | error!("Invalid form"); | 99 | error!("Invalid form"); |
90 | return Ok(respond_with_status(StatusCode::UNPROCESSABLE_ENTITY)); | 100 | return Ok(respond_with_status(StatusCode::UNPROCESSABLE_ENTITY)); |
@@ -105,7 +115,7 @@ pub async fn shortner_service(req: Request<Body>, mut conn: Connection) -> Resul | |||
105 | } | 115 | } |
106 | 116 | ||
107 | trace!("Attempting to parse multipart request"); | 117 | trace!("Attempting to parse multipart request"); |
108 | return process_multipart(req.into_body(), boundary.unwrap(), &mut conn).await; | 118 | return process_multipart(req, boundary.unwrap(), &mut conn).await; |
109 | } | 119 | } |
110 | &Method::GET => { | 120 | &Method::GET => { |
111 | trace!("GET: {}", req.uri()); | 121 | trace!("GET: {}", req.uri()); |