Bit.Auth-r/src/main.rs

mod funcs;
mod types;
mod html;
mod url;

use {
	std::{
		sync::Arc,
		net::SocketAddr,
		collections::HashMap,
		process::exit,
	},
	http_body_util::{
		Full,
	},
	hyper::{
		StatusCode,
		Request,
		Response,
		HeaderMap,
		header::HeaderValue,
		body::{
			Bytes,
			Incoming,
		},
		server::conn::http1 as Server,
		service::service_fn,
	},
	hyper_util::{
		rt::TokioIo,
	},
	tokio::{
		net::TcpListener,
	},
	serde_json::{
		Value as Json,
		json,
	},
	skytable::{
		query,
		Config,
		pool::{
			self,
			ConnectionMgrTcp
		},
		error::Error as SkyError,
	},
	bb8::{
		Pool,
	},
	urlencoding::{
		encode as url_encode,
	},
	crate::{
		types::{
			users::Users,
			sites::Sites,
		},
		html::*,
		funcs::*,
	},
};

type Res<T, E> = std::result::Result<T, E>;
type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
type DBPool = Arc<Pool<ConnectionMgrTcp>>;
type FullBytes = Result<Response<Full<Bytes>>>;


const PORT: u16 = 8083;

const DB_POOL: u32 = 32;
const DB_ADDR: &str = "127.0.0.1";
const DB_PORT: u16 = 2003;
const DB_USER: &str = "root";
const DB_PASS: &str = "5a14775a-a490-4212345678";

const TOKEN_LIFETIME: u32 = 300;
const REFRESH_LIFETIME: u32 = 2_678_400;

const APIV0_LIFETIME: u32 = 120;


#[tokio::main]
async fn main() -> Result<()> {
	println!("Starting.");
	
	print!("Binding port...");
	let addr = SocketAddr::from(([0, 0, 0, 0], PORT));
	let listener = TcpListener::bind(addr).await;
	if listener.is_err() {
		println!(" Error.");
		exit(1);
	}
	let listener = listener?;
	println!(" OK.");

	print!("Connecting to DB...");
	let pool = pool::get_async(DB_POOL, Config::new(DB_ADDR, DB_PORT, DB_USER, DB_PASS)).await;
	let pool = Arc::new(pool.unwrap());

	if pool.get().await.is_err() {
		println!(" Error.");
		exit(1);
	}
	println!(" OK.");

	init_tables(pool.clone()).await?;

	println!("Server started on port: {}", PORT);

	loop {
		let (stream, _) = listener.accept().await?;

		let ip = format!("{:?}", stream.peer_addr().unwrap().ip());

		let io = TokioIo::new(stream);

		let this_pool = pool.clone();
		tokio::task::spawn(async move {
			if let Err(err) = Server::Builder::new().serve_connection(
				io, service_fn(move |req| handle_connection(
					req,
					this_pool.clone(),
					ip.clone()
				))
			).await {
				println!("Error serving connection: {:?}", err);
			}
		});
	}
}

async fn handle_connection(req: Request<Incoming>, pool: DBPool, ip: String) -> FullBytes {
	let t = time_mcs();

	if ip == "1.1.1.1" {
	}

	let res = Response::new(Full::new(Bytes::new()));
	let (mut parts, _) = res.into_parts();
	let body: String;
	let restype: HeaderValue;

	let mut headers = parts.headers.clone();

	let cookies = get_cookies(req.headers().clone());
	let mut token = match cookies.get("token") {
		x if x.is_none() => "".to_owned(),
		x => x.unwrap().to_owned()
	};
	let mut logged = false;

	match <str as AsRef<str>>::as_ref(req.uri().path()) {
		x if x.starts_with("/api/") => {},
		_ => 'jwt_check: {
			if token == "" { break 'jwt_check; }

			let is_live = jwt_verify(pool.clone(), &token)
				.await?
				.claims
				.as_object()
				.unwrap()
				.len() != 0;

			match is_live {
				true => { logged = true; },
				_ => 'ref_check: {
					let reftoken = cookies.get("refresh");
					let reftoken = match reftoken.is_none() {
						false => reftoken.unwrap(),
						_ => ""
					};

					if reftoken == "" { break 'ref_check; }

					let ref_data: HashMap<String, Json> = jwt_verify(pool.clone(), reftoken)
						.await?
						.claims
						.as_object()
						.unwrap()
						.clone()
						.into_iter()
						.collect();

					if ref_data.clone().len() != 0 {
						let mut con = pool.get().await.unwrap();

						let tokenid = ref_data.get("uuid").unwrap().as_str().unwrap();
						let newref = format!("{}", uuid_v4().as_hyphenated());
						let time = time();

						let q = con.query_parse::<(String,)>(&query!(
							r#"
								SELECT uid
								FROM bitauth.tokens
								WHERE uuid = ?
							"#,
							tokenid
						)).await;
						if q.is_err() { break 'ref_check; }
						let (uuid,) = q.unwrap();

						let (login,) = con.query_parse::<(String,)>(&query!(
							r#"
								SELECT login
								FROM bitauth.users
								WHERE uuid = ?
							"#,
							uuid.clone()
						)).await?;					
						let _ = con.query_parse::<()>(&query!(
							r#"
								UPDATE bitauth.tokens
								SET ref = ?, refend = ?
								WHERE uuid = ?
							"#,
							newref.clone(), time + REFRESH_LIFETIME, tokenid
						)).await;

						token = jwt_sign(pool.clone(), json!({
							"login": login.clone(),
							"uuid": uuid.clone(),
							"iat": time,
							"exp": time + TOKEN_LIFETIME
						})).await.unwrap();

						set_cookie(&mut headers, "token", &token);
						set_cookie(&mut headers, "refresh",
							&jwt_sign(pool.clone(), json!({
								"uuid": tokenid,
								"iat": time,
								"ref": newref.clone(),
								"exp": time + REFRESH_LIFETIME
							})).await.unwrap()
						);

						logged = true;
					}
				}
			}
		}
	}

	(body, parts.status, restype) = match req.uri().path().as_ref() {
		"/" => url::index::index(),
		"/cabinet" => url::login::login(req, pool.clone(), &mut headers).await?,
		"/login" => url::login::login(req, pool.clone(), &mut headers).await?,
		x if x == "/authorize" && logged => url::authorize::authorize(req, pool.clone(), token).await?,
		"/authorize" => uri_auth_required(req, &mut headers).await?,
		"/register" => url::register::register(req, pool.clone(), &mut headers).await?,
		"/recover" => url::recover::recover(),
		x if x.starts_with("/@") => url::user::user(req, pool.clone()).await?,
		x if x.starts_with("/api/") => url::api::endpoint(req, pool.clone()).await,
		_ => url::nf::nf()
	};

	headers.insert(hyper::header::CONTENT_TYPE, restype);
	parts.headers = headers;

	let body = body.replace("PAGE_TIME", &format!("{}", time_mcs() - t));
	Ok(Response::from_parts(parts, Full::new(Bytes::from(body))))
}


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))
}