diff --git a/src/db.rs b/src/db.rs index d547392..92552e5 100644 --- a/src/db.rs +++ b/src/db.rs @@ -18,7 +18,8 @@ fn get_db_environment() -> String { #[derive(RustcEncodable, Debug)] pub struct Record { pub public_ip: String, - pub message: String, + pub client: String, + pub message: String, pub timestamp: i64 // i64 because of the database type. } @@ -29,7 +30,7 @@ fn escape(string: &str) -> String { pub enum FindFilter { PublicIp(String), - PublicIpAndMessage(String, String) + PublicIpAndClient(String, String) } pub struct Db { @@ -45,6 +46,7 @@ impl Db { let connection = Connection::open(get_db_environment()).unwrap(); connection.execute("CREATE TABLE IF NOT EXISTS boxes ( public_ip TEXT NOT NULL, + client TEXT NOT NULL, message TEXT, timestamp INTEGER )", &[]).unwrap(); @@ -78,11 +80,11 @@ impl Db { ); try!(stmt.query(&[&escape(&public_ip)])) }, - FindFilter::PublicIpAndMessage(public_ip, message) => { + FindFilter::PublicIpAndClient(public_ip, client_id) => { stmt = try!( - self.connection.prepare("SELECT * FROM boxes WHERE (public_ip=$1 and message=$2)") + self.connection.prepare("SELECT * FROM boxes WHERE (public_ip=$1 and client=$2)") ); - try!(stmt.query(&[&escape(&public_ip), &escape(&message)])) + try!(stmt.query(&[&escape(&public_ip), &escape(&client_id)])) } }; @@ -91,8 +93,9 @@ impl Db { let row = try!(result_row); records.push(Record { public_ip: row.get(0), - message: row.get(1), - timestamp: row.get(2) + client: row.get(1), + message: row.get(2), + timestamp: row.get(3) }); } Ok(records) @@ -100,18 +103,18 @@ impl Db { pub fn update(&self, record: Record) -> rusqlite::Result { self.connection.execute("UPDATE boxes - SET public_ip=$1, message=$2, timestamp=$3 - WHERE (public_ip=$4 AND message=$5)", - &[&record.public_ip, &record.message, - &record.public_ip, &record.message]) + SET public_ip=$1, message=$2, client=$3, timestamp=$4 + WHERE (public_ip=$5 AND client=$6)", + &[&record.public_ip, &record.message, &record.client, + &record.timestamp, &record.public_ip, &record.client]) } pub fn add(&self, record: Record) -> rusqlite::Result { self.connection.execute("INSERT INTO boxes - (public_ip, message, timestamp) - VALUES ($1, $2, $3)", - &[&record.public_ip, &record.message, - &record.timestamp]) + (public_ip, client, message, timestamp) + VALUES ($1, $2, $3, $4)", + &[&record.public_ip, &record.client, &record.message, + &record.timestamp]) } pub fn delete_older_than(&self, timestamp: i64) -> rusqlite::Result { @@ -124,7 +127,7 @@ fn test_db() { let db = Db::new(); // Look for a record, but the db is empty. - match db.find(FindFilter::PublicIpAndMessage("127.0.0.1".to_owned(), ".knilxof.org".to_owned())) { + match db.find(FindFilter::PublicIpAndClient("127.0.0.1".to_owned(), "".to_owned())) { Ok(vec) => { assert!(vec.is_empty()); }, Err(err) => { println!("Unexpected error: {}", err); assert!(false); } } @@ -132,7 +135,8 @@ fn test_db() { let mut r = Record { public_ip: "127.0.0.1".to_owned(), - message: ".knilxof.org".to_owned(), + message: "".to_owned(), + client: "".to_owned(), timestamp: now }; @@ -142,7 +146,7 @@ fn test_db() { Err(err) => { println!("Unexpected error: {}", err); assert!(false); } } // Check that we find it. - match db.find(FindFilter::PublicIpAndMessage("127.0.0.1".to_owned(), ".knilxof.org".to_owned())) { + match db.find(FindFilter::PublicIpAndClient("127.0.0.1".to_owned(), "".to_owned())) { Ok(records) => { assert_eq!(records.len(), 1); assert_eq!(records[0].timestamp, now); @@ -153,7 +157,8 @@ fn test_db() { // Add another record with the same public IP, but a different local one. r = Record { public_ip: "127.0.0.1".to_owned(), - message: ".knilxof.org".to_owned(), + message: "".to_owned(), + client: "".to_owned(), timestamp: now }; match db.add(r) { @@ -165,8 +170,11 @@ fn test_db() { match db.find(FindFilter::PublicIp("127.0.0.1".to_owned())) { Ok(records) => { assert_eq!(records.len(), 2); - assert_eq!(records[0].message, ".knilxof.org"); - assert_eq!(records[1].message, ".knilxof.org"); + assert_eq!(records[0].message, ""); + assert_eq!(records[1].message, ""); + + assert_eq!(records[0].client, ""); + assert_eq!(records[1].client, ""); }, Err(err) => { println!("Unexpected error: {}", err); assert!(false); } } diff --git a/src/routes.rs b/src/routes.rs index 4d7f0cd..c92a3a9 100644 --- a/src/routes.rs +++ b/src/routes.rs @@ -30,6 +30,7 @@ fn register(req: &mut Request) -> IronResult { // Get the local IP and optional tunnel url from the body, #[derive(RustcDecodable, Debug)] struct RegisterBody { + client: String, message: String, } @@ -44,6 +45,7 @@ fn register(req: &mut Request) -> IronResult { }; let message = body.message; + let client_id = body.client; // And the public IP from the socket. let public_ip = format!("{}", req.remote_addr.ip()); @@ -51,21 +53,22 @@ fn register(req: &mut Request) -> IronResult { // Get the current number of seconds since epoch. let now = Db::seconds_from_epoch(); - info!("POST /register public_ip={} message={} time is {}", - public_ip, message, now); + info!("POST /register public_ip={} client={} message={} time is {}", + public_ip, client_id, message, now); // Save this registration in the database. // If we already have the same (local, tunnel, public) match, update it, // if not create a new match. let db = Db::new(); match db.find( - FindFilter::PublicIpAndMessage(public_ip.clone(), message.clone()) + FindFilter::PublicIpAndClient(public_ip.clone(), client_id.clone()) ) { Ok(rvect) => { // If the vector is empty, create a new record, if not update // the existing one with the new timestamp. let record = Record { public_ip: public_ip, + client: client_id, message: message, timestamp: now, }; @@ -83,10 +86,12 @@ fn register(req: &mut Request) -> IronResult { Err(_) => { let record = Record { public_ip: public_ip, + client: client_id, message: message, timestamp: now, }; - if let Err(_) = db.add(record) { + if let Err(e) = db.add(record) { + error!("Error {}", e); return EndpointError::with(status::InternalServerError, 501) } }