564 lines
13 KiB
Rust
564 lines
13 KiB
Rust
mod types;
|
|
mod html;
|
|
mod api;
|
|
|
|
use {
|
|
std::{
|
|
sync::Arc,
|
|
net::SocketAddr,
|
|
collections::HashMap,
|
|
process::exit,
|
|
time::{
|
|
SystemTime,
|
|
UNIX_EPOCH,
|
|
},
|
|
},
|
|
chrono::{
|
|
DateTime,
|
|
},
|
|
http_body_util::{
|
|
Full,
|
|
BodyExt,
|
|
},
|
|
hyper::{
|
|
StatusCode,
|
|
Request,
|
|
Response,
|
|
Method,
|
|
HeaderMap,
|
|
header::HeaderValue,
|
|
body::{
|
|
Bytes,
|
|
Incoming,
|
|
},
|
|
server::conn::http1 as Server,
|
|
service::service_fn,
|
|
},
|
|
hyper_util::{
|
|
rt::TokioIo,
|
|
},
|
|
tokio::{
|
|
net::TcpListener,
|
|
},
|
|
rsa::{
|
|
pkcs1::{
|
|
EncodeRsaPublicKey,
|
|
EncodeRsaPrivateKey,
|
|
},
|
|
RsaPrivateKey,
|
|
RsaPublicKey,
|
|
},
|
|
urlencoding::{
|
|
decode as url_decode,
|
|
},
|
|
jsonwebtoken as jwt,
|
|
jwt::{
|
|
Header,
|
|
Algorithm,
|
|
TokenData,
|
|
Validation,
|
|
EncodingKey,
|
|
DecodingKey,
|
|
encode as jwt_encode,
|
|
decode as jwt_decode,
|
|
},
|
|
serde_json::{
|
|
Value as Json,
|
|
json,
|
|
},
|
|
uuid::{
|
|
Uuid,
|
|
},
|
|
skytable::{
|
|
query,
|
|
Config,
|
|
pool::{
|
|
self,
|
|
ConnectionMgrTcp
|
|
},
|
|
error::Error as SkyError,
|
|
},
|
|
bb8::{
|
|
Pool,
|
|
},
|
|
crate::{
|
|
types::{
|
|
users::Users,
|
|
sites::Sites,
|
|
},
|
|
html::*,
|
|
},
|
|
};
|
|
|
|
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 = "192.168.1.49";
|
|
const DB_PORT: u16 = 2003;
|
|
const DB_USER: &str = "root";
|
|
const DB_PASS: &str = "rootpass12345678";
|
|
|
|
const TOKEN_LIFETIME: u32 = 300;
|
|
const REFRESH_LIFETIME: u32 = 2_678_400;
|
|
|
|
|
|
#[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_ns();
|
|
|
|
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 token = cookies.get("token");
|
|
let token = match token.is_none() {
|
|
false => token.unwrap(),
|
|
_ => ""
|
|
};
|
|
|
|
match <str as AsRef<str>>::as_ref(req.uri().path()) {
|
|
x if x.starts_with("/api/") => {}
|
|
_ => {
|
|
// println!("{}", token);
|
|
|
|
if token != "" && jwt_verify(pool.clone(), token)
|
|
.await?.claims.as_object().unwrap().len() == 0
|
|
{
|
|
println!("Invalid suka");
|
|
parts.status = StatusCode::FOUND;
|
|
set_cookie(&mut headers, "token", "");
|
|
set_location(&mut headers, "/");
|
|
parts.headers = headers;
|
|
return Ok(Response::from_parts(parts, Full::new(Bytes::new())));
|
|
}
|
|
}
|
|
}
|
|
|
|
(body, parts.status, restype) = match req.uri().path().as_ref() {
|
|
"/" => uri_index(),
|
|
"/cabinet" => uri_login(req, pool.clone(), &mut headers).await?,
|
|
"/login" => uri_login(req, pool.clone(), &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?,
|
|
x if x.starts_with("/api/") => api::endpoint(req, pool.clone()).await,
|
|
_ => uri_404()
|
|
};
|
|
|
|
headers.insert(hyper::header::CONTENT_TYPE, restype);
|
|
parts.headers = headers;
|
|
|
|
let body = body.replace("RENDER_TIME", &format!("{}", time_ns() - t));
|
|
Ok(Response::from_parts(parts, Full::new(Bytes::from(body))))
|
|
}
|
|
|
|
fn build_html(body: &str) -> String {
|
|
format!("{}{}{}{}", HEADER_HTML, CSS3, body, FOOTER_HTML)
|
|
}
|
|
|
|
fn set_cookie(headers: &mut HeaderMap, key: &str, value: &str) {
|
|
let time = DateTime::from_timestamp((time() + REFRESH_LIFETIME) as i64, 0)
|
|
.expect("REASON")
|
|
.to_rfc2822();
|
|
let time = time.replace("+0000", "GMT");
|
|
headers.append(
|
|
hyper::header::SET_COOKIE,
|
|
HeaderValue::from_str(format!("{}={}; HttpOnly; Expires={}", key, value, time).as_str())
|
|
.unwrap()
|
|
);
|
|
}
|
|
|
|
fn get_cookies(headers: HeaderMap) -> HashMap<String, String> {
|
|
let header = headers.get(hyper::header::COOKIE);
|
|
let cookies = match header.is_none() {
|
|
false => header.unwrap().to_str().unwrap(),
|
|
_ => ""
|
|
};
|
|
|
|
double_split(cookies.to_owned(), ";", "=")
|
|
}
|
|
|
|
fn set_location(headers: &mut HeaderMap, location: &str) {
|
|
headers.append(
|
|
hyper::header::LOCATION,
|
|
HeaderValue::from_str(format!("{}", location).as_str()).unwrap()
|
|
);
|
|
}
|
|
|
|
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 (access, refresh) = login_user(pool.clone(), body).await?;
|
|
|
|
set_cookie(headers, "token", &access);
|
|
set_cookie(headers, "refresh", &refresh);
|
|
}
|
|
|
|
let restype: HeaderValue = "text/html".parse().unwrap();
|
|
Ok((build_html(LOGIN_HTML), StatusCode::OK, restype))
|
|
}
|
|
|
|
fn uri_index() -> (String, StatusCode, HeaderValue) {
|
|
let restype: HeaderValue = "text/html".parse().unwrap();
|
|
(build_html(INDEX_HTML), StatusCode::OK, restype)
|
|
}
|
|
|
|
async fn uri_register(req: Request<Incoming>, pool: DBPool, headers: &mut HeaderMap) -> Result<(String, StatusCode, HeaderValue)> {
|
|
let mut body = "".to_owned();
|
|
let mut status = StatusCode::OK;
|
|
let restype: HeaderValue = "text/html".parse().unwrap();
|
|
|
|
if *req.method() == Method::POST {
|
|
let request = get_body_from_request(req).await?;
|
|
let request = String::from_utf8(request).unwrap();
|
|
let request = double_split(request, "&", "=");
|
|
|
|
match create_user(pool.clone(), request).await? {
|
|
true => {
|
|
println!("Created");
|
|
set_location(headers, "/login");
|
|
},
|
|
_ => {
|
|
println!("Failed");
|
|
set_location(headers, "/register");
|
|
}
|
|
}
|
|
|
|
status = StatusCode::FOUND;
|
|
}
|
|
else {
|
|
body = build_html(REG_HTML);
|
|
}
|
|
|
|
Ok((body, status, restype))
|
|
}
|
|
|
|
fn uri_recover() -> (String, StatusCode, HeaderValue) {
|
|
let restype: HeaderValue = "text/html".parse().unwrap();
|
|
(build_html(RECOVER_HTML), StatusCode::OK, restype)
|
|
}
|
|
|
|
async fn uri_user(req: Request<Incoming>, pool: DBPool) -> Res<(String, StatusCode, HeaderValue), SkyError> {
|
|
let uri: &str = req.uri().path().as_ref();
|
|
let login = &uri[2..uri.len()];
|
|
let body: String;
|
|
|
|
let user = get_user(pool, login.to_string()).await;
|
|
if user.is_ok() {
|
|
body = format!("{}", user?.uuid);
|
|
}
|
|
else {
|
|
body = "Not fond :(".to_owned();
|
|
}
|
|
|
|
let restype: HeaderValue = "text/html".parse().unwrap();
|
|
Ok((build_html(&body), StatusCode::OK, restype))
|
|
}
|
|
|
|
fn uri_404() -> (String, StatusCode, HeaderValue) {
|
|
let restype: HeaderValue = "text/html".parse().unwrap();
|
|
(build_html(NF_HTML), StatusCode::NOT_FOUND, restype)
|
|
}
|
|
|
|
async fn init_tables(pool: DBPool) -> Res<(), SkyError> {
|
|
let mut con = pool.get().await.unwrap();
|
|
|
|
let _ = con.query_parse::<()>(&query!("CREATE SPACE IF NOT EXISTS bitauth")).await;
|
|
let _ = con.query_parse::<()>(&query!(r#"
|
|
CREATE MODEL IF NOT EXISTS bitauth.users(
|
|
login: string,
|
|
uuid: string,
|
|
password: string,
|
|
email: string,
|
|
tokens: list {type: string}
|
|
)
|
|
"#)).await;
|
|
let _ = con.query_parse::<()>(&query!(r#"
|
|
CREATE MODEL IF NOT EXISTS bitauth.sites(
|
|
uuid: string,
|
|
uid: string,
|
|
domain: string,
|
|
pkey: binary,
|
|
skey: binary
|
|
)
|
|
"#)).await;
|
|
let _ = con.query_parse::<()>(&query!(r#"
|
|
CREATE MODEL IF NOT EXISTS bitauth.tokens(
|
|
uuid: string,
|
|
uid: string,
|
|
sid: string,
|
|
ref: string,
|
|
refend: uint32
|
|
)
|
|
"#)).await;
|
|
|
|
let q = con.query_parse::<Sites>(&query!("SELECT * FROM bitauth.sites WHERE uuid = 0")).await;
|
|
if q.is_err() {
|
|
let (skey, pkey) = rsa_gen();
|
|
let _ = con.query_parse::<()>(&query!(
|
|
"INSERT INTO bitauth.sites { uuid: ?, uid: ?, domain: ?, pkey: ?, skey: ? }",
|
|
"0", "0", "", pkey, skey
|
|
)).await;
|
|
};
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn get_user(pool: DBPool, login: String) -> Res<Users, SkyError> {
|
|
let mut con = pool.get().await.unwrap();
|
|
|
|
let q = con
|
|
.query_parse::<Users>(&query!("SELECT * FROM bitauth.users WHERE login = ?", login))
|
|
.await;
|
|
|
|
Ok(q?)
|
|
}
|
|
|
|
async fn login_user(pool: DBPool, data: HashMap<String, String>) -> Res<(String, String), SkyError> {
|
|
let mut err = false;
|
|
let mut ret: (String, String) = Default::default();
|
|
let mut con = pool.get().await.unwrap();
|
|
|
|
let login = &data.get("login").unwrap().trim().to_lowercase();
|
|
let pass = data.get("password").unwrap().trim();
|
|
|
|
let q = con
|
|
.query_parse::<Users>(&query!("SELECT * FROM bitauth.users WHERE login = ?", login))
|
|
.await;
|
|
if q.is_err() { err = true };
|
|
|
|
if !err {
|
|
let q = q.unwrap();
|
|
if bcrypt::verify(pass, q.password.as_str()).unwrap() {
|
|
let token = format!("{}", uuid_v4().as_hyphenated());
|
|
let reftoken = format!("{}", uuid_v4().as_hyphenated());
|
|
let time = time();
|
|
let uuid: String = q.uuid;
|
|
let login: String = q.login;
|
|
|
|
let _ = con.query_parse::<()>(&query!(
|
|
"INSERT INTO bitauth.tokens { uuid: ?, uid: ?, sid: ?, ref: ?, refend: ? }",
|
|
token.clone(), uuid.clone(), "0", reftoken.clone(), time + REFRESH_LIFETIME
|
|
)).await;
|
|
let _ = con.query_parse::<()>(&query!(
|
|
"UPDATE bitauth.users SET tokens += ? WHERE login = ?",
|
|
token.clone(), login.clone()
|
|
)).await;
|
|
|
|
ret = (
|
|
jwt_sign(pool.clone(), json!({
|
|
"login": login.clone(),
|
|
"uuid": uuid.clone(),
|
|
"iat": time,
|
|
"exp": time + TOKEN_LIFETIME
|
|
})).await.unwrap(),
|
|
jwt_sign(pool.clone(), json!({
|
|
"uuid": token.clone(),
|
|
"iat": time,
|
|
"ref": reftoken.clone(),
|
|
"exp": time + REFRESH_LIFETIME
|
|
})).await.unwrap()
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(ret)
|
|
}
|
|
|
|
async fn create_user(pool: DBPool, data: HashMap<String, String>) -> Res<bool, SkyError> {
|
|
let mut ret = true;
|
|
let mut con = pool.get().await.unwrap();
|
|
|
|
let login = &data.get("login").unwrap().trim().to_lowercase();
|
|
let email = &data.get("email").unwrap().trim().to_lowercase();
|
|
let pass = data.get("password").unwrap().trim();
|
|
let pass2 = data.get("password2").unwrap().trim();
|
|
|
|
if pass != pass2 { ret = false };
|
|
if pass.len() < 8 { ret = false };
|
|
|
|
let q = con
|
|
.query_parse::<(String,)>(&query!("SELECT uuid FROM bitauth.users WHERE login = ?", login))
|
|
.await;
|
|
if q.is_ok() { ret = false };
|
|
|
|
if ret {
|
|
let uuid = format!("{}", uuid_v4().as_hyphenated());
|
|
let pass = bcrypt::hash(pass, 12).unwrap();
|
|
|
|
let q = con.query_parse::<()>(&query!(
|
|
r#"INSERT INTO bitauth.users {
|
|
login: ?,
|
|
uuid: ?,
|
|
password: ?,
|
|
email: ?,
|
|
tokens: []
|
|
}"#,
|
|
login,
|
|
uuid,
|
|
pass,
|
|
email,
|
|
)).await;
|
|
|
|
if q.is_err() { ret = false }
|
|
}
|
|
|
|
Ok(ret)
|
|
}
|
|
|
|
async fn get_body_from_request(mut req: Request<Incoming>) -> Result<Vec<u8>> {
|
|
let mut body: Vec<u8> = vec![];
|
|
while let Some(next) = req.frame().await {
|
|
let frame = next?;
|
|
if let Some(chunk) = frame.data_ref() {
|
|
body.extend_from_slice(chunk);
|
|
}
|
|
}
|
|
|
|
Ok(body)
|
|
}
|
|
|
|
fn double_split(body: String, first: &str, second: &str) -> HashMap<String, String> {
|
|
body.split(first)
|
|
.filter_map(|c| {
|
|
c.split_once(second)
|
|
.map(|(l, r)| (
|
|
l.trim().to_owned(),
|
|
format!("{}", url_decode(r).expect("UTF-8")).trim().to_owned()
|
|
))
|
|
})
|
|
.collect::<HashMap<String, String>>()
|
|
}
|
|
|
|
fn rsa_gen() -> (Vec<u8>, Vec<u8>) {
|
|
let mut rng = rand::thread_rng();
|
|
let bits = 2048;
|
|
let skey = RsaPrivateKey::new(&mut rng, bits).expect("RSA err");
|
|
let pkey = RsaPublicKey::from(&skey);
|
|
|
|
(
|
|
skey.to_pkcs1_der().unwrap().as_bytes().to_vec(),
|
|
pkey.to_pkcs1_der().unwrap().as_bytes().to_vec()
|
|
)
|
|
}
|
|
|
|
async fn jwt_sign(pool: DBPool, data: Json) -> Result<String> {
|
|
let mut con = pool.get().await.unwrap();
|
|
|
|
let q = con
|
|
.query_parse::<Sites>(&query!("SELECT * FROM bitauth.sites WHERE uuid = ?", "0"))
|
|
.await;
|
|
if q.is_err() { return Ok("".to_owned()) };
|
|
let skey = q.unwrap().skey;
|
|
|
|
let skey = EncodingKey::from_rsa_der(&skey);
|
|
let header = Header::new(Algorithm::RS512);
|
|
let token = jwt_encode(&header, &data, &skey)?;
|
|
|
|
Ok(token)
|
|
}
|
|
|
|
async fn jwt_verify(pool: DBPool, token: &str) -> Result<TokenData<Json>> {
|
|
let mut con = pool.get().await.unwrap();
|
|
let mut ret = TokenData { header: Header::new(Algorithm::RS512), claims: json![{}] };
|
|
|
|
let q = con
|
|
.query_parse::<Sites>(&query!("SELECT * FROM bitauth.sites WHERE uuid = ?", "0"))
|
|
.await;
|
|
if q.is_err() { return Ok(ret) };
|
|
let pkey = q.unwrap().pkey;
|
|
|
|
let pkey = DecodingKey::from_rsa_der(&pkey);
|
|
let token = jwt_decode::<Json>(&token, &pkey, &Validation::new(Algorithm::RS512));
|
|
|
|
match token.is_ok() {
|
|
true => ret = token.clone().unwrap(),
|
|
_ => {}
|
|
}
|
|
|
|
Ok(ret)
|
|
}
|
|
|
|
fn uuid_v4() -> Uuid {
|
|
Uuid::new_v4()
|
|
}
|
|
|
|
fn time() -> u32 {
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_secs() as u32
|
|
}
|
|
|
|
fn time_ns() -> u128 {
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_micros()
|
|
}
|