diff options
| author | Akshay <[email protected]> | 2020-08-24 07:07:03 +0100 |
|---|---|---|
| committer | Akshay <[email protected]> | 2020-08-24 07:07:03 +0100 |
| commit | 0a0952e6edb2f4a8c7883b661ef6d4bb4eb14405 (patch) | |
| tree | e7bddf107823494ab73de40fba34aaa8abecc9c3 | |
| parent | 28194780aa62b8f02432eb29d857adbd26d88569 (diff) | |
parse host header and return full url
| -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()); |