From e28a8c2893c5fa49520fc805b3be49e0887cda37 Mon Sep 17 00:00:00 2001
From: BitHeaven <bit.heaven@bk.ru>
Date: Fri, 12 Apr 2024 21:49:53 +0500
Subject: [PATCH] Worked v0 API

---
 src/main.rs | 55 ++++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 42 insertions(+), 13 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 60a249e..4ad5083 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -73,6 +73,10 @@ use {
 	bb8::{
 		Pool,
 	},
+	urlencoding::{
+		encode as url_encode,
+		decode as url_decode,
+	},
 	crate::{
 		types::{
 			users::Users,
@@ -219,15 +223,16 @@ async fn handle_connection(req: Request<Incoming>, pool: DBPool, ip: String) ->
 						let newref = format!("{}", uuid_v4().as_hyphenated());
 						let time = time();
 
-						// TODO: FIX ERROR IF TOKEN INVALID
-						let (uuid,) = con.query_parse::<(String,)>(&query!(
+						let q = con.query_parse::<(String,)>(&query!(
 							r#"
 								SELECT uid
 								FROM bitauth.tokens
 								WHERE uuid = ?
 							"#,
 							tokenid
-						)).await?;
+						)).await;
+						if q.is_err() { break 'ref_check; }
+						let (uuid,) = q.unwrap();
 
 						let (login,) = con.query_parse::<(String,)>(&query!(
 							r#"
@@ -278,7 +283,7 @@ async fn handle_connection(req: Request<Incoming>, pool: DBPool, ip: String) ->
 		"/cabinet" => uri_login(req, pool.clone(), &mut headers).await?,
 		"/login" => uri_login(req, pool.clone(), &mut headers).await?,
 		x if x == "/authorize" && logged => uri_authorize(req, pool.clone(), token).await?,
-//		"/authorize" => uri_authorize(req, pool.clone()).await?,
+		"/authorize" => uri_auth_required(req, &mut headers).await?,
 		"/register" => uri_register(req, pool.clone(), &mut headers).await?,
 		"/recover" => uri_recover(),
 		x if x.starts_with("/@") => uri_user(req, pool.clone()).await?,
@@ -323,19 +328,34 @@ fn set_location(headers: &mut HeaderMap, location: &str) {
 }
 
 async fn uri_login(req: Request<Incoming>, pool: DBPool, headers: &mut HeaderMap) -> Result<(String, StatusCode, HeaderValue)> {
-	if *req.method() == Method::POST {
-		let body = get_body_from_request(req).await?;
-		let body = String::from_utf8(body).unwrap();
-		let body = double_split(body, "&", "=");
+	let mut body = build_html(LOGIN_HTML);
+	let mut status = StatusCode::OK;
+	let restype: HeaderValue = "text/html".parse().unwrap();
 
-		let (access, refresh) = login_user(pool.clone(), body).await?;
+	if *req.method() == Method::POST {
+		let r = double_split(req.uri().query().or(Some("")).unwrap().to_owned(), "&", "=");
+
+		let post = get_body_from_request(req).await?;
+		let post = String::from_utf8(post).unwrap();
+		let post = double_split(post, "&", "=");
+
+		let (access, refresh) = login_user(pool.clone(), post).await?;
 
 		set_cookie(headers, "token", &access);
 		set_cookie(headers, "refresh", &refresh);
+
+		let r = r.get("q");
+		match r.is_some() {
+			true => {
+				status = StatusCode::FOUND;
+				body = "".to_owned();
+				set_location(headers, format!("{}", url_decode(r.unwrap())?).as_str());
+			},
+			_ => {}
+		}
 	}
 
-	let restype: HeaderValue = "text/html".parse().unwrap();
-	Ok((build_html(LOGIN_HTML), StatusCode::OK, restype))
+	Ok((body, status, restype))
 }
 
 async fn uri_authorize(req: Request<Incoming>, pool: DBPool, token: String) -> Result<(String, StatusCode, HeaderValue)> {
@@ -357,6 +377,15 @@ async fn uri_authorize(req: Request<Incoming>, pool: DBPool, token: String) -> R
 	Ok((build_html(AUTHORIZE_HTML), StatusCode::OK, restype))
 }
 
+async fn uri_auth_required(req: Request<Incoming>, headers: &mut HeaderMap) -> Result<(String, StatusCode, HeaderValue)> {
+	let url = url_encode(req.uri().path_and_query().unwrap().as_str());
+
+	set_location(headers, format!("/login?q={}", url).as_str());
+
+	let restype: HeaderValue = "text/html".parse().unwrap();
+	Ok(("".to_owned(), StatusCode::FOUND, restype))
+}
+
 async fn authorize_user(pool: DBPool, token: String, session: String) {
 	let mut con = pool.get().await.unwrap();
 
@@ -376,7 +405,7 @@ async fn authorize_user(pool: DBPool, token: String, session: String) {
 			session: ?,
 			login: ?,
 			uuid: ?,
-			expire: ?,
+			expire: ?
 		}"#,
 		session,
 		login,
@@ -489,7 +518,7 @@ async fn init_tables(pool: DBPool) -> Res<(), SkyError> {
 			session: string,
 			login: string,
 			uuid: string,
-			expire: uint32,
+			expire: uint32
 		)
 	"#)).await;