Merge pull request #34 from oxen-io/auth-token-handling

Fix Multi Device Auth Token Handling

Cargo.lock

@@ -1555,6 +1555,16 @@ dependencies = [
  "smallvec",
 ]
 
+[[package]]
+name = "rusqlite_migration"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc814a80978c4b6aae9748dacf5570384dcb5255c17593b904d9b3b1482141a8"
+dependencies = [
+ "log",
+ "rusqlite",
+]
+
 [[package]]
 name = "rustc-demangle"
 version = "0.1.19"
@@ -1756,6 +1766,7 @@ dependencies = [
  "rand_core 0.5.1",
  "reqwest",
  "rusqlite",
+ "rusqlite_migration",
  "serde",
  "serde_json",
  "sha2",

Cargo.toml

@@ -29,6 +29,7 @@ rand = "0.8"
 rand_core = "0.5"
 reqwest = { version = "0.11", features = ["json"] }
 rusqlite = { version = "0.24", features = ["bundled"] }
+rusqlite_migration = "0.4"
 r2d2_sqlite = "0.17"
 r2d2 = "0.8"
 serde = { version = "1.0", features = ["derive"] }

src/handlers.rs

@@ -407,10 +407,11 @@ pub fn claim_auth_token(
     let token = &pending_tokens[index].1;
     // Store the claimed token
     let stmt = format!(
-        "INSERT OR REPLACE INTO {} (public_key, token) VALUES (?1, ?2)",
+        "INSERT INTO {} (public_key, timestamp, token) VALUES (?1, ?2, ?3)",
         storage::TOKENS_TABLE
     );
-    match conn.execute(&stmt, params![public_key, hex::encode(token)]) {
+    let now = chrono::Utc::now().timestamp();
+    match conn.execute(&stmt, params![public_key, now, hex::encode(token)]) {
         Ok(_) => (),
         Err(e) => {
             error!("Couldn't insert token due to error: {}.", e);
@@ -909,7 +910,7 @@ pub fn get_member_count(
     // Get a database connection
     let conn = pool.get().map_err(|_| Error::DatabaseFailedInternally)?;
     // Query the database
-    let raw_query = format!("SELECT COUNT(public_key) FROM {}", storage::TOKENS_TABLE);
+    let raw_query = format!("SELECT COUNT(DISTINCT public_key) FROM {}", storage::TOKENS_TABLE);
     let mut query = conn.prepare(&raw_query).map_err(|_| Error::DatabaseFailedInternally)?;
     let rows = match query.query_map(params![], |row| row.get(0)) {
         Ok(rows) => rows,

src/storage.rs

@@ -6,6 +6,7 @@ use std::sync::Mutex;
 use log::{error, info};
 use r2d2_sqlite::SqliteConnectionManager;
 use rusqlite::params;
+use rusqlite_migration::{Migrations, M};
 
 use super::errors::Error;
 
@@ -141,9 +142,9 @@ fn create_room_tables_if_needed(conn: &DatabaseConnection) {
     // The token is stored as hex here (rather than as bytes) because it's more convenient for lookup
     let tokens_table_cmd = format!(
         "CREATE TABLE IF NOT EXISTS {} (
-        public_key TEXT PRIMARY KEY,
+        public_key TEXT,
         timestamp INTEGER,
-        token TEXT
+        token TEXT PRIMARY KEY
     )",
         TOKENS_TABLE
     );
@@ -157,6 +158,7 @@ fn create_room_tables_if_needed(conn: &DatabaseConnection) {
         FILES_TABLE
     );
     conn.execute(&files_table_cmd, params![]).expect("Couldn't create files table.");
+    // User activity table
     let user_activity_table_cmd = format!(
         "CREATE TABLE IF NOT EXISTS {} (
         public_key TEXT PRIMARY KEY,
@@ -318,15 +320,27 @@ pub async fn prune_files(file_expiration: i64) {
 // Migration
 
 pub fn perform_migration() {
-    // ensure all rooms schemas are up to date
     let rooms = match get_all_room_ids() {
         Ok(ids) => ids,
         Err(_e) => {
             return error!("Couldn't get all room IDs.");
         }
     };
+    let create_tokens_table_cmd = format!(
+        "CREATE TABLE IF NOT EXISTS {} (
+        public_key TEXT,
+        timestamp INTEGER,
+        token TEXT PRIMARY KEY
+    )",
+        TOKENS_TABLE
+    );
+    let migrations =
+        Migrations::new(vec![M::up("DROP TABLE tokens"), M::up(&create_tokens_table_cmd)]);
     for room in rooms {
         create_database_if_needed(&room);
+        let pool = pool_by_room_id(&room);
+        let mut conn = pool.get().unwrap();
+        migrations.to_latest(&mut conn).unwrap();
     }
 }