diff options
author | Akshay <[email protected]> | 2020-08-23 15:32:53 +0100 |
---|---|---|
committer | Akshay <[email protected]> | 2020-08-23 15:32:53 +0100 |
commit | e16743598ccdaae9709034a7215c35ccfe2a49fa (patch) | |
tree | f9180a143c3a4d48185c52113aa23a44b2cff7a7 | |
parent | 99d59ca3be5ce319e43671f63449b36344bb3f84 (diff) |
refactor src into db and service
-rw-r--r-- | src/db.rs | 19 | ||||
-rw-r--r-- | src/main.rs | 145 | ||||
-rw-r--r-- | src/service.rs | 121 |
3 files changed, 145 insertions, 140 deletions
diff --git a/src/db.rs b/src/db.rs new file mode 100644 index 0000000..a5c0f85 --- /dev/null +++ b/src/db.rs | |||
@@ -0,0 +1,19 @@ | |||
1 | use anyhow::Result; | ||
2 | use rusqlite::{Connection, OpenFlags, NO_PARAMS}; | ||
3 | |||
4 | use std::path::Path; | ||
5 | |||
6 | pub fn init_db<P: AsRef<Path>>(p: P) -> Result<Connection> { | ||
7 | let conn = Connection::open_with_flags( | ||
8 | p, | ||
9 | OpenFlags::SQLITE_OPEN_CREATE | OpenFlags::SQLITE_OPEN_READ_WRITE, | ||
10 | )?; | ||
11 | conn.execute( | ||
12 | "CREATE TABLE IF NOT EXISTS urls ( | ||
13 | link TEXT PRIMARY KEY, | ||
14 | shortlink TEXT NOT NULL | ||
15 | )", | ||
16 | NO_PARAMS, | ||
17 | )?; | ||
18 | return Ok(conn); | ||
19 | } | ||
diff --git a/src/main.rs b/src/main.rs index 86af042..d31abfe 100644 --- a/src/main.rs +++ b/src/main.rs | |||
@@ -1,145 +1,10 @@ | |||
1 | use anyhow::{Context, Result}; | 1 | use anyhow::Result; |
2 | use hyper::header::CONTENT_TYPE; | ||
3 | use hyper::service::{make_service_fn, service_fn}; | 2 | use hyper::service::{make_service_fn, service_fn}; |
4 | use hyper::{Body, Method, Request, Response, Server, StatusCode}; | 3 | use hyper::Server; |
5 | use multer::Multipart; | ||
6 | use nanoid::nanoid; | ||
7 | use rusqlite::{params, Connection, OpenFlags, NO_PARAMS}; | ||
8 | use url::form_urlencoded; | ||
9 | 4 | ||
10 | use std::collections::HashMap; | 5 | mod db; |
11 | use std::path::Path; | 6 | mod service; |
12 | 7 | use service::shortner_service; | |
13 | fn respond_with_shortlink<S: AsRef<str>>(shortlink: S) -> Response<Body> { | ||
14 | Response::builder() | ||
15 | .status(StatusCode::OK) | ||
16 | .header("content-type", "text/html") | ||
17 | .body(Body::from(shortlink.as_ref().to_string())) | ||
18 | .unwrap() | ||
19 | } | ||
20 | |||
21 | fn respond_with_status(s: StatusCode) -> Response<Body> { | ||
22 | Response::builder().status(s).body(Body::empty()).unwrap() | ||
23 | } | ||
24 | |||
25 | fn shorten<S: AsRef<str>>(url: S, conn: &mut Connection) -> Result<String> { | ||
26 | let mut stmt = conn.prepare("select * from urls where link = ?1")?; | ||
27 | let mut rows = stmt.query(params![url.as_ref().to_string()])?; | ||
28 | if let Some(row) = rows.next()? { | ||
29 | return Ok(row.get(1)?); | ||
30 | } else { | ||
31 | let new_id = nanoid!(4); | ||
32 | conn.execute( | ||
33 | "insert into urls (link, shortlink) values (?1, ?2)", | ||
34 | params![url.as_ref().to_string(), new_id], | ||
35 | )?; | ||
36 | return Ok(new_id); | ||
37 | } | ||
38 | } | ||
39 | |||
40 | fn get_link<S: AsRef<str>>(url: S, conn: &mut Connection) -> Result<Option<String>> { | ||
41 | let url = url.as_ref(); | ||
42 | let mut stmt = conn.prepare("select * from urls where shortlink = ?1")?; | ||
43 | let mut rows = stmt.query(params![url.to_string()])?; | ||
44 | if let Some(row) = rows.next()? { | ||
45 | return Ok(row.get(0)?); | ||
46 | } else { | ||
47 | return Ok(None); | ||
48 | } | ||
49 | } | ||
50 | |||
51 | async fn process_multipart( | ||
52 | body: Body, | ||
53 | boundary: String, | ||
54 | conn: &mut Connection, | ||
55 | ) -> Result<Response<Body>> { | ||
56 | let mut m = Multipart::new(body, boundary); | ||
57 | if let Some(field) = m.next_field().await? { | ||
58 | if field.name() == Some("shorten") { | ||
59 | let content = field | ||
60 | .text() | ||
61 | .await | ||
62 | .with_context(|| format!("Expected field name"))?; | ||
63 | |||
64 | let shortlink = shorten(content, conn)?; | ||
65 | return Ok(respond_with_shortlink(shortlink)); | ||
66 | } | ||
67 | } | ||
68 | Ok(Response::builder() | ||
69 | .status(StatusCode::OK) | ||
70 | .body(Body::empty())?) | ||
71 | } | ||
72 | |||
73 | async fn shortner_service(req: Request<Body>) -> Result<Response<Body>> { | ||
74 | let mut conn = init_db("./urls.db_3").unwrap(); | ||
75 | |||
76 | match req.method() { | ||
77 | &Method::POST => { | ||
78 | let boundary = req | ||
79 | .headers() | ||
80 | .get(CONTENT_TYPE) | ||
81 | .and_then(|ct| ct.to_str().ok()) | ||
82 | .and_then(|ct| multer::parse_boundary(ct).ok()); | ||
83 | |||
84 | if boundary.is_none() { | ||
85 | let b = hyper::body::to_bytes(req) | ||
86 | .await | ||
87 | .with_context(|| format!("Failed to stream request body!"))?; | ||
88 | |||
89 | let params = form_urlencoded::parse(b.as_ref()) | ||
90 | .into_owned() | ||
91 | .collect::<HashMap<String, String>>(); | ||
92 | |||
93 | if let Some(n) = params.get("shorten") { | ||
94 | let s = shorten(n, &mut conn)?; | ||
95 | return Ok(respond_with_shortlink(s)); | ||
96 | } else { | ||
97 | return Ok(respond_with_status(StatusCode::UNPROCESSABLE_ENTITY)); | ||
98 | } | ||
99 | } | ||
100 | |||
101 | return process_multipart(req.into_body(), boundary.unwrap(), &mut conn).await; | ||
102 | } | ||
103 | &Method::GET => { | ||
104 | let shortlink = req.uri().path().to_string(); | ||
105 | let link = get_link(&shortlink[1..], &mut conn); | ||
106 | if let Some(l) = link.unwrap() { | ||
107 | Ok(Response::builder() | ||
108 | .header("Location", &l) | ||
109 | .header("content-type", "text/html") | ||
110 | .status(StatusCode::MOVED_PERMANENTLY) | ||
111 | .body(Body::from(format!( | ||
112 | "You will be redirected to: {}. If not, click the link.", | ||
113 | &l | ||
114 | )))?) | ||
115 | } else { | ||
116 | Ok(Response::builder() | ||
117 | .status(StatusCode::NOT_FOUND) | ||
118 | .body(Body::empty())?) | ||
119 | } | ||
120 | } | ||
121 | _ => { | ||
122 | return Ok(Response::builder() | ||
123 | .status(StatusCode::NOT_FOUND) | ||
124 | .body(Body::empty())?) | ||
125 | } | ||
126 | } | ||
127 | } | ||
128 | |||
129 | fn init_db<P: AsRef<Path>>(p: P) -> Result<Connection> { | ||
130 | let conn = Connection::open_with_flags( | ||
131 | p, | ||
132 | OpenFlags::SQLITE_OPEN_CREATE | OpenFlags::SQLITE_OPEN_READ_WRITE, | ||
133 | )?; | ||
134 | conn.execute( | ||
135 | "CREATE TABLE IF NOT EXISTS urls ( | ||
136 | link TEXT PRIMARY KEY, | ||
137 | shortlink TEXT NOT NULL | ||
138 | )", | ||
139 | NO_PARAMS, | ||
140 | )?; | ||
141 | return Ok(conn); | ||
142 | } | ||
143 | 8 | ||
144 | fn main() -> Result<()> { | 9 | fn main() -> Result<()> { |
145 | smol::run(async { | 10 | smol::run(async { |
diff --git a/src/service.rs b/src/service.rs new file mode 100644 index 0000000..55a42bf --- /dev/null +++ b/src/service.rs | |||
@@ -0,0 +1,121 @@ | |||
1 | use anyhow::{Context, Result}; | ||
2 | use hyper::header::CONTENT_TYPE; | ||
3 | use hyper::{Body, Method, Request, Response, StatusCode}; | ||
4 | use multer::Multipart; | ||
5 | use nanoid::nanoid; | ||
6 | use rusqlite::{params, Connection}; | ||
7 | use url::form_urlencoded; | ||
8 | |||
9 | use std::collections::HashMap; | ||
10 | |||
11 | use crate::db::init_db; | ||
12 | |||
13 | fn respond_with_shortlink<S: AsRef<str>>(shortlink: S) -> Response<Body> { | ||
14 | Response::builder() | ||
15 | .status(StatusCode::OK) | ||
16 | .header("content-type", "text/html") | ||
17 | .body(Body::from(shortlink.as_ref().to_string())) | ||
18 | .unwrap() | ||
19 | } | ||
20 | |||
21 | fn respond_with_status(s: StatusCode) -> Response<Body> { | ||
22 | Response::builder().status(s).body(Body::empty()).unwrap() | ||
23 | } | ||
24 | |||
25 | fn shorten<S: AsRef<str>>(url: S, conn: &mut Connection) -> Result<String> { | ||
26 | let mut stmt = conn.prepare("select * from urls where link = ?1")?; | ||
27 | let mut rows = stmt.query(params![url.as_ref().to_string()])?; | ||
28 | if let Some(row) = rows.next()? { | ||
29 | return Ok(row.get(1)?); | ||
30 | } else { | ||
31 | let new_id = nanoid!(4); | ||
32 | conn.execute( | ||
33 | "insert into urls (link, shortlink) values (?1, ?2)", | ||
34 | params![url.as_ref().to_string(), new_id], | ||
35 | )?; | ||
36 | return Ok(new_id); | ||
37 | } | ||
38 | } | ||
39 | |||
40 | fn get_link<S: AsRef<str>>(url: S, conn: &mut Connection) -> Result<Option<String>> { | ||
41 | let url = url.as_ref(); | ||
42 | let mut stmt = conn.prepare("select * from urls where shortlink = ?1")?; | ||
43 | let mut rows = stmt.query(params![url.to_string()])?; | ||
44 | if let Some(row) = rows.next()? { | ||
45 | return Ok(row.get(0)?); | ||
46 | } else { | ||
47 | return Ok(None); | ||
48 | } | ||
49 | } | ||
50 | |||
51 | async fn process_multipart( | ||
52 | body: Body, | ||
53 | boundary: String, | ||
54 | conn: &mut Connection, | ||
55 | ) -> Result<Response<Body>> { | ||
56 | let mut m = Multipart::new(body, boundary); | ||
57 | if let Some(field) = m.next_field().await? { | ||
58 | if field.name() == Some("shorten") { | ||
59 | let content = field | ||
60 | .text() | ||
61 | .await | ||
62 | .with_context(|| format!("Expected field name"))?; | ||
63 | |||
64 | let shortlink = shorten(content, conn)?; | ||
65 | return Ok(respond_with_shortlink(shortlink)); | ||
66 | } | ||
67 | } | ||
68 | Ok(Response::builder() | ||
69 | .status(StatusCode::OK) | ||
70 | .body(Body::empty())?) | ||
71 | } | ||
72 | |||
73 | pub async fn shortner_service(req: Request<Body>) -> Result<Response<Body>> { | ||
74 | let mut conn = init_db("./urls.db_3").unwrap(); | ||
75 | |||
76 | match req.method() { | ||
77 | &Method::POST => { | ||
78 | let boundary = req | ||
79 | .headers() | ||
80 | .get(CONTENT_TYPE) | ||
81 | .and_then(|ct| ct.to_str().ok()) | ||
82 | .and_then(|ct| multer::parse_boundary(ct).ok()); | ||
83 | |||
84 | if boundary.is_none() { | ||
85 | let b = hyper::body::to_bytes(req) | ||
86 | .await | ||
87 | .with_context(|| format!("Failed to stream request body!"))?; | ||
88 | |||
89 | let params = form_urlencoded::parse(b.as_ref()) | ||
90 | .into_owned() | ||
91 | .collect::<HashMap<String, String>>(); | ||
92 | |||
93 | if let Some(n) = params.get("shorten") { | ||
94 | let s = shorten(n, &mut conn)?; | ||
95 | return Ok(respond_with_shortlink(s)); | ||
96 | } else { | ||
97 | return Ok(respond_with_status(StatusCode::UNPROCESSABLE_ENTITY)); | ||
98 | } | ||
99 | } | ||
100 | |||
101 | return process_multipart(req.into_body(), boundary.unwrap(), &mut conn).await; | ||
102 | } | ||
103 | &Method::GET => { | ||
104 | let shortlink = req.uri().path().to_string(); | ||
105 | let link = get_link(&shortlink[1..], &mut conn); | ||
106 | if let Some(l) = link.unwrap() { | ||
107 | Ok(Response::builder() | ||
108 | .header("Location", &l) | ||
109 | .header("content-type", "text/html") | ||
110 | .status(StatusCode::MOVED_PERMANENTLY) | ||
111 | .body(Body::from(format!( | ||
112 | "You will be redirected to: {}. If not, click the link.", | ||
113 | &l | ||
114 | )))?) | ||
115 | } else { | ||
116 | Ok(respond_with_status(StatusCode::NOT_FOUND)) | ||
117 | } | ||
118 | } | ||
119 | _ => Ok(respond_with_status(StatusCode::NOT_FOUND)), | ||
120 | } | ||
121 | } | ||