feat: refactor session management by removing unnecessary keys and adding pending sessions table

AdamLBS · AdamLBS · commit bcf1f79917d7 · 2025-10-22T19:29:43.000+02:00
diff --git a/sql_models/seed.sql b/sql_models/seed.sql
@@ -1,99 +1,163 @@
+-- Enable UUID generator
 CREATE EXTENSION IF NOT EXISTS "pgcrypto";
 
+-- =========================
+-- Users
+-- =========================
 CREATE TABLE users (
   id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
   username TEXT UNIQUE NOT NULL,
   created_at TIMESTAMP DEFAULT NOW()
 );
 
+-- =========================
+-- Devices
+-- =========================
 CREATE TABLE devices (
   id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
   user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
-  identity_pubkey TEXT NOT NULL, 
+  identity_pubkey TEXT NOT NULL,
   prekey_pubkey TEXT NOT NULL,
   signed_prekey_pub TEXT NOT NULL,
   signed_prekey_sig TEXT NOT NULL,
-  one_time_prekeys JSONB NOT NULL,
+  one_time_prekeys JSONB NOT NULL,   -- array of one-time public prekeys
   device_label TEXT,
   push_token TEXT,
   last_seen TIMESTAMP DEFAULT NOW(),
   created_at TIMESTAMP DEFAULT NOW(),
-  UNIQUE(user_id, identity_pubkey)
+  UNIQUE (user_id, identity_pubkey)
 );
 
+-- =========================
+-- Chats
+-- =========================
 CREATE TABLE chats (
-    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
-    user_a_id UUID REFERENCES users(id) ON DELETE CASCADE,
-    user_b_id UUID REFERENCES users(id) ON DELETE CASCADE,
-    created_at TIMESTAMP DEFAULT NOW(),
-    UNIQUE(user_a_id, user_b_id)
+  id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+  type TEXT CHECK (type IN ('direct','group')) DEFAULT 'direct',
+
+  -- Direct chat participants
+  user_a UUID REFERENCES users(id) ON DELETE CASCADE,
+  user_b UUID REFERENCES users(id) ON DELETE CASCADE,
+
+  -- Group chat info
+  name TEXT,
+  owner_id UUID REFERENCES users(id) ON DELETE SET NULL,
+
+  -- Will be linked after messages table is created
+  last_message_id UUID,
+
+  created_at TIMESTAMP DEFAULT NOW(),
+  updated_at TIMESTAMP DEFAULT NOW(),
+
+  -- For direct chats: both users must be set, and enforce strict ordering
+  CONSTRAINT chats_direct_shape
+    CHECK (
+      type <> 'direct'
+      OR (user_a IS NOT NULL AND user_b IS NOT NULL AND user_a < user_b)
+    )
 );
 
+-- Prevent duplicates (A,B) vs (B,A)
+CREATE UNIQUE INDEX uniq_direct_chat
+  ON chats (LEAST(user_a, user_b), GREATEST(user_a, user_b))
+  WHERE type = 'direct';
 
-CREATE TABLE messages (
-  id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+-- =========================
+-- Chat Members (for group chats)
+-- =========================
+CREATE TABLE chat_members (
   chat_id UUID REFERENCES chats(id) ON DELETE CASCADE,
-  from_device_id UUID REFERENCES devices(id) ON DELETE SET NULL,
-  plaintext_hash TEXT,                -- SHA256 du plaintext avant chiffrement
-  created_at TIMESTAMP DEFAULT NOW()
+  user_id UUID REFERENCES users(id) ON DELETE CASCADE,
+  role TEXT DEFAULT 'member',
+  joined_at TIMESTAMP DEFAULT NOW(),
+  PRIMARY KEY (chat_id, user_id)
 );
 
-CREATE TABLE message_deliveries (
+-- =========================
+-- Messages
+-- =========================
+CREATE TABLE messages (
   id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
-  message_id UUID REFERENCES messages(id) ON DELETE CASCADE,
+
+  -- Logical message id (shared by all device fanouts)
+  logical_msg_id TEXT NOT NULL,
+
+  chat_id UUID REFERENCES chats(id) ON DELETE CASCADE,
+
+  from_user_id UUID REFERENCES users(id) ON DELETE CASCADE,
+  from_device_id UUID REFERENCES devices(id) ON DELETE CASCADE,
+
+  to_user_id UUID REFERENCES users(id) ON DELETE CASCADE,
   to_device_id UUID REFERENCES devices(id) ON DELETE CASCADE,
-  ciphertext TEXT NOT NULL,
-  delivered BOOLEAN DEFAULT FALSE,
+
+  header JSONB NOT NULL,       -- Double Ratchet header (DH pubkey, counters, etc.)
+  ciphertext TEXT NOT NULL,    -- base64(nonce || cipher || mac)
+
+  delivered_at TIMESTAMP,
+  read_at TIMESTAMP,
   created_at TIMESTAMP DEFAULT NOW()
 );
 
+-- Add FK once messages table exists
+ALTER TABLE chats
+  ADD CONSTRAINT fk_chats_last_message
+  FOREIGN KEY (last_message_id)
+  REFERENCES messages(id)
+  ON DELETE SET NULL;
+
+-- Indexes
+CREATE INDEX idx_messages_todevice ON messages(to_device_id, created_at);
+CREATE INDEX idx_messages_chatid   ON messages(chat_id, created_at);
+CREATE INDEX idx_messages_logical  ON messages(logical_msg_id);
 
+-- =========================
+-- Sessions (metadata only, no secret keys)
+-- =========================
 CREATE TABLE sessions (
   id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
   chat_id UUID REFERENCES chats(id) ON DELETE CASCADE,
   sender_device_id UUID REFERENCES devices(id) ON DELETE CASCADE,
   receiver_device_id UUID REFERENCES devices(id) ON DELETE CASCADE,
-  root_key BYTEA NOT NULL,
-  send_chain_key BYTEA,
-  recv_chain_key BYTEA,
-  send_counter INTEGER DEFAULT 0,
-  recv_counter INTEGER DEFAULT 0,
-  ratchet_pub BYTEA,
-  last_remote_pub BYTEA,
   created_at TIMESTAMP DEFAULT NOW(),
   updated_at TIMESTAMP DEFAULT NOW(),
-  UNIQUE(sender_device_id, receiver_device_id)
+  UNIQUE (sender_device_id, receiver_device_id)
 );
 
+-- =========================
+-- Used Tokens (anti-replay or enrollment)
+-- =========================
 CREATE TABLE used_tokens (
   token TEXT PRIMARY KEY,
   used_at TIMESTAMP DEFAULT NOW()
 );
 
-
+-- =========================
+-- Pending Sessions (X3DH handshake initialization)
+-- =========================
 CREATE TABLE pending_sessions (
   id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
-  sender_device_id UUID NOT NULL REFERENCES devices(id),
-  recipient_device_id UUID NOT NULL REFERENCES devices(id),
-  ephemeral_pubkey TEXT NOT NULL,
-  sender_prekey_pub TEXT NOT NULL,
-  otpk_used TEXT NOT NULL,
-  ciphertext TEXT NOT NULL,
+  sender_device_id UUID NOT NULL REFERENCES devices(id) ON DELETE CASCADE,
+  recipient_device_id UUID NOT NULL REFERENCES devices(id) ON DELETE CASCADE,
+  ephemeral_pubkey TEXT NOT NULL,     -- Ephemeral key (EK_A pub)
+  sender_prekey_pub TEXT NOT NULL,    -- IK/SPK_A used by sender
+  otpk_used TEXT NOT NULL,    -- Whether a one-time prekey was consumed
+  ciphertext TEXT NOT NULL,           -- Encrypted payload (init message)
   created_at TIMESTAMP DEFAULT NOW(),
-  state TEXT DEFAULT 'initiated' CHECK (state IN ('initiated', 'responded', 'completed'))
+  state TEXT DEFAULT 'initiated' CHECK (state IN ('initiated','responded','completed'))
 );
 
-ALTER TABLE chats
-ADD CONSTRAINT chats_unique_pair CHECK (user_a_id < user_b_id);
-
+CREATE INDEX idx_pending_sessions_recipient ON pending_sessions(recipient_device_id, created_at);
 
+-- =========================
+-- Convenience View: Public device list
+-- =========================
 CREATE VIEW user_devices_view AS
-SELECT 
+SELECT
   u.username,
   d.id AS device_id,
   d.identity_pubkey,
   d.signed_prekey_pub,
   d.signed_prekey_sig,
-  (d.one_time_prekeys -> 0) AS one_time_prekey_pub
+  (d.one_time_prekeys ->> 0) AS one_time_prekey_pub
 FROM users u
-JOIN devices d ON u.id = d.user_id;
+JOIN devices d ON u.id = d.user_id;
diff --git a/src/controllers/session_controller.rs b/src/controllers/session_controller.rs
@@ -30,9 +30,6 @@ pub struct ConfirmSessionBody {
     pub pending_session_id: Uuid,
     pub sender_device_id: Uuid,
     pub receiver_device_id: Uuid,
-    pub root_key: String,
-    pub ratchet_pub: String,
-    pub last_remote_pub: String,
 }
 
 
@@ -96,7 +93,6 @@ pub async fn confirm_session(
     AuthenticatedDevice(device): AuthenticatedDevice,
     Json(payload): Json<ConfirmSessionBody>,
 ) -> Result<impl IntoResponse, (StatusCode, &'static str)> {
-    // 1️⃣ Vérifie que la pending session existe et appartient au device courant
     let pending_session = session_repository::get_pending_session_by_id(
         &state.pool,
         &payload.pending_session_id,
@@ -110,19 +106,6 @@ pub async fn confirm_session(
         None => return Err((StatusCode::NOT_FOUND, "Pending session not found or not owned by device")),
     };
 
-    // 2️⃣ Decode les clés Base64 -> bytes
-    let root_key_bytes = general_purpose::STANDARD
-        .decode(&payload.root_key)
-        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid root_key"))?;
-
-    let ratchet_pub_bytes = general_purpose::STANDARD
-        .decode(&payload.ratchet_pub)
-        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid ratchet_pub"))?;
-
-    let last_remote_pub_bytes = general_purpose::STANDARD
-        .decode(&payload.last_remote_pub)
-        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid last_remote_pub"))?;
-
     // 3️⃣ Récupère ou crée le chat_id
     let chat_id = session_repository::get_or_create_chat_id(
         &state.pool,
@@ -141,10 +124,7 @@ pub async fn confirm_session(
         &state.pool,
         &chat_id,
         &payload.sender_device_id,
-        &payload.receiver_device_id,
-        &root_key_bytes,
-        &ratchet_pub_bytes,
-        &last_remote_pub_bytes,
+        &payload.receiver_device_id
     )
     .await
     .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Failed to insert session"))?;
diff --git a/src/models/message.rs b/src/models/message.rs
diff --git a/src/models/mod.rs b/src/models/mod.rs
@@ -3,3 +3,4 @@ pub mod enrollment_token;
 pub mod keys;
 pub mod session;
 pub mod user;
+pub mod message;
diff --git a/src/models/session.rs b/src/models/session.rs
@@ -20,13 +20,6 @@ pub struct Session {
     pub chat_id: Option<Uuid>,
     pub sender_device_id: Uuid,
     pub receiver_device_id: Uuid,
-    pub root_key: Vec<u8>,
-    pub send_chain_key: Option<Vec<u8>>,
-    pub recv_chain_key: Option<Vec<u8>>,
-    pub send_counter: i32,
-    pub recv_counter: i32,
-    pub ratchet_pub: Option<Vec<u8>>,
-    pub last_remote_pub: Option<Vec<u8>>,
     pub created_at: Option<NaiveDateTime>,
     pub updated_at: Option<NaiveDateTime>,
 }
diff --git a/src/repository/message_repository.rs b/src/repository/message_repository.rs
diff --git a/src/repository/mod.rs b/src/repository/mod.rs
@@ -3,3 +3,4 @@ pub mod enrollment_token_repository;
 pub mod keys_repository;
 pub mod session_repository;
 pub mod user_repository;
+pub mod message_repository;
diff --git a/src/repository/session_repository.rs b/src/repository/session_repository.rs
@@ -106,11 +106,11 @@ pub async fn get_or_create_chat_id(
     if let Some(chat_id) = sqlx::query_scalar!(
         r#"
         SELECT id FROM chats
-        WHERE (user_a_id = LEAST(
+        WHERE (user_a= LEAST(
                     (SELECT user_id FROM devices WHERE id=$1),
                     (SELECT user_id FROM devices WHERE id=$2)
                )
-           AND user_b_id = GREATEST(
+           AND user_b = GREATEST(
                     (SELECT user_id FROM devices WHERE id=$1),
                     (SELECT user_id FROM devices WHERE id=$2)
                ))
@@ -126,7 +126,7 @@ pub async fn get_or_create_chat_id(
     // sinon crée-le proprement en respectant la contrainte
     let new_chat_id = sqlx::query_scalar!(
         r#"
-        INSERT INTO chats (user_a_id, user_b_id)
+        INSERT INTO chats (user_a, user_b)
         VALUES (
             LEAST(
               (SELECT user_id FROM devices WHERE id=$1),
@@ -153,30 +153,20 @@ pub async fn insert_or_update_session(
     chat_id: &Uuid,
     sender_device_id: &Uuid,
     receiver_device_id: &Uuid,
-    root_key: &[u8],
-    ratchet_pub: &[u8],
-    last_remote_pub: &[u8],
 ) -> Result<(), sqlx::Error> {
     sqlx::query!(
         r#"
         INSERT INTO sessions (
-          chat_id, sender_device_id, receiver_device_id,
-          root_key, ratchet_pub, last_remote_pub
+          chat_id, sender_device_id, receiver_device_id
         )
-        VALUES ($1, $2, $3, $4, $5, $6)
+        VALUES ($1, $2, $3)
         ON CONFLICT (sender_device_id, receiver_device_id)
         DO UPDATE SET
-          root_key = EXCLUDED.root_key,
-          ratchet_pub = EXCLUDED.ratchet_pub,
-          last_remote_pub = EXCLUDED.last_remote_pub,
           updated_at = NOW()
         "#,
         chat_id,
         sender_device_id,
         receiver_device_id,
-        root_key,
-        ratchet_pub,
-        last_remote_pub
     )
     .execute(pool)
     .await?;
