ntr/mnml
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

pog

Browse Source
This commit is contained in:
ntr 2019-07-11 20:00:39 +10:00
parent e06f97252b
commit 7b50bfaffb
5 changed files with 87 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
server/src/account.rs
View File

@ -12,6 +12,7 @@ use names::{name as generate_name};
use construct::{Construct, construct_recover, construct_spawn};
use instance::{Instance, instance_delete};
use mtx::{Mtx, FREE_MTX};
use pg::Db;
use failure::Error;
use failure::{err_msg, format_err};
@ -48,7 +49,7 @@ pub fn select(tx: &mut Transaction, id: Uuid) -> Result<Account, Error> {
Ok(Account { id, name: row.get(1), credits, subscribed })
}
pub fn from_token(tx: &mut Transaction, token: String) -> Result<Account, Error> {
pub fn from_token(db: &Db, token: String) -> Result<Account, Error> {
let query = "
SELECT id, name, subscribed, credits
FROM accounts
@ -56,7 +57,7 @@ pub fn from_token(tx: &mut Transaction, token: String) -> Result<Account, Error>
AND token_expiry > now();
";
let result = tx
let result = db
.query(query, &[&token])?;
let row = result.iter().next()

1
server/src/main.rs
View File

@ -25,6 +25,7 @@ extern crate persistent;
extern crate router;
extern crate cookie;
extern crate tungstenite;
extern crate crossbeam_channel;
mod account;
mod construct;

2
server/src/net.rs
View File

@ -137,7 +137,7 @@ fn login_res(token: String) -> Response {
let v = Cookie::build(TOKEN_HEADER, token)
.http_only(true)
.same_site(SameSite::Strict)
.max_age(Duration::seconds(-1)) // 1 week aligns with db set
.max_age(Duration::weeks(1)) // 1 week aligns with db set
.finish();
let mut res = Response::with(status::Ok);

2
server/src/rpc.rs
View File

@ -66,7 +66,7 @@ enum RpcRequest {
VboxReclaim { instance_id: Uuid, index: usize },
}
pub fn receive(data: Vec<u8>, db: &Db, _client: &mut WebSocket<TcpStream>, begin: Instant, account: Option<&Account>) -> Result<RpcResult, Error> {
pub fn receive(data: Vec<u8>, db: &Db, _client: &mut WebSocket<TcpStream>, begin: Instant, account: &Option<Account>) -> Result<RpcResult, Error> {
// cast the msg to this type to receive method name
match from_slice::<RpcRequest>(&data) {
Ok(v) => {

109
server/src/ws.rs
View File

@ -1,10 +1,16 @@
use std::time::{Instant};
use std::net::{TcpListener};
use std::thread::{spawn};
use std::str;
use cookie::Cookie;
use tungstenite::server::accept_hdr;
use tungstenite::Message::Binary;
use tungstenite::handshake::server::Request;
use tungstenite::handshake::server::{Request, ErrorResponse};
use tungstenite::http::StatusCode;
use crossbeam_channel::{unbounded, Sender};
use failure::Error;
use failure::err_msg;
@ -12,8 +18,9 @@ use failure::err_msg;
use serde_cbor::{to_vec};
use net::TOKEN_HEADER;
use rpc::{receive};
use rpc;
use pg::PgPool;
use account;
#[derive(Debug,Clone,Serialize)]
struct RpcError {
@ -25,42 +32,88 @@ pub fn start(pool: PgPool) {
for stream in ws_server.incoming() {
let ws_pool = pool.clone();
spawn(move || {
let (acc_s, acc_r) = unbounded();
let cb = |req: &Request| {
let token = req.headers.find_first(TOKEN_HEADER);
println!("{:?}", token);
let err = || ErrorResponse {
error_code: StatusCode::FORBIDDEN,
headers: None,
body: Some("Unauthorized".into()),
};
if let Some(cl) = req.headers.find_first("Cookie") {
let cookie_list = str::from_utf8(cl).or(Err(err()))?;
for s in cookie_list.split(";").map(|s| s.trim()) {
let cookie = Cookie::parse(s).or(Err(err()))?;
// got auth token
if cookie.name() == TOKEN_HEADER {
info!("{:?}", cookie.value().to_string());
acc_s.send(Some(cookie.value().to_string())).or(Err(err()))?;
}
};
};
acc_s.send(None).unwrap();
Ok(None)
};
let mut websocket = accept_hdr(stream.unwrap(), cb).unwrap();
let account = match acc_r.recv().unwrap() {
Some(t) => {
let db = ws_pool.get()
.expect("unable to get db connection");
match account::from_token(&db, t) {
Ok(a) => {
let state = to_vec(&rpc::RpcResult::AccountState(a.clone())).unwrap();
websocket.write_message(Binary(state)).unwrap();
Some(a)
},
Err(e) => {
warn!("{:?}", e);
return;
},
}
},
None => None,
};
loop {
match websocket.read_message() {
Ok(msg) => {
let begin = Instant::now();
let db_connection = ws_pool.get()
.expect("unable to get db connection");
match msg {
Binary(data) => {
let begin = Instant::now();
let db_connection = ws_pool.get()
.expect("unable to get db connection");
let data = msg.into_data();
match receive(data, &db_connection, &mut websocket, begin, None) {
Ok(reply) => {
let response = to_vec(&reply)
.expect("failed to serialize response");
match rpc::receive(data, &db_connection, &mut websocket, begin, &account) {
Ok(reply) => {
let response = to_vec(&reply)
.expect("failed to serialize response");
if let Err(e) = websocket.write_message(Binary(response)) {
// connection closed
debug!("{:?}", e);
return;
};
if let Err(e) = websocket.write_message(Binary(response)) {
// connection closed
debug!("{:?}", e);
return;
};
},
Err(e) => {
warn!("{:?}", e);
let response = to_vec(&RpcError { err: e.to_string() })
.expect("failed to serialize error response");
if let Err(e) = websocket.write_message(Binary(response)) {
// connection closed
debug!("{:?}", e);
return;
};
}
}
},
Err(e) => {
warn!("{:?}", e);
let response = to_vec(&RpcError { err: e.to_string() })
.expect("failed to serialize error response");
if let Err(e) = websocket.write_message(Binary(response)) {
// connection closed
debug!("{:?}", e);
return;
};
}
_ => (),
}
},
// connection is closed