Gateway database modifications for different modes

common/gateway-storage/migrations/20240910120000_generic_client_id.sql

@@ -0,0 +1,98 @@
+/*
+ * Copyright 2024 - Nym Technologies SA <[email protected]>
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+CREATE TABLE clients (
+   id              INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
+   client_type     TEXT NOT NULL CHECK(client_type IN ('entry_mixnet', 'exit_mixnet', 'entry_wireguard', 'exit_wireguard'))
+);
+
+INSERT INTO clients (id, client_type)
+SELECT id, 'entry_mixnet'
+FROM shared_keys;
+
+CREATE TABLE shared_keys_tmp (
+   client_id                                INTEGER NOT NULL PRIMARY KEY REFERENCES clients(id),
+   client_address_bs58                      TEXT NOT NULL UNIQUE,
+   derived_aes128_ctr_blake3_hmac_keys_bs58 TEXT NOT NULL
+);
+
+INSERT INTO shared_keys_tmp (client_id, client_address_bs58, derived_aes128_ctr_blake3_hmac_keys_bs58)
+SELECT id as client_id, client_address_bs58, derived_aes128_ctr_blake3_hmac_keys_bs58 FROM shared_keys;
+
+CREATE TABLE available_bandwidth_tmp (
+   client_id INTEGER NOT NULL PRIMARY KEY REFERENCES clients(id),
+   available INTEGER NOT NULL,
+   expiration TIMESTAMP WITHOUT TIME ZONE
+);
+
+INSERT INTO available_bandwidth_tmp
+SELECT * FROM available_bandwidth;
+
+DROP TABLE available_bandwidth;
+ALTER TABLE available_bandwidth_tmp RENAME TO available_bandwidth;
+
+CREATE TABLE received_ticket_tmp (
+   id                  INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
+   client_id           INTEGER NOT NULL REFERENCES clients(id),
+   received_at         TIMESTAMP WITHOUT TIME ZONE NOT NULL,
+   rejected            BOOLEAN
+);
+
+INSERT INTO received_ticket_tmp
+SELECT * FROM received_ticket;
+
+DROP INDEX received_ticket_index;
+CREATE INDEX received_ticket_index ON received_ticket_tmp (received_at);
+
+ -- received tickets that are in the process of verifying
+CREATE TABLE ticket_data_tmp (
+   ticket_id           INTEGER NOT NULL PRIMARY KEY REFERENCES received_ticket_tmp(id),
+
+   -- serial_number, alongside the entire row, will get purged after redemption is complete
+   serial_number       BLOB NOT NULL UNIQUE,
+
+   --    data will get purged after 80% of signers verifies it
+   data                BLOB
+);
+
+INSERT INTO ticket_data_tmp
+SELECT * FROM ticket_data;
+
+DROP TABLE ticket_data;
+ALTER TABLE ticket_data_tmp RENAME TO ticket_data;
+
+-- result of a verification from a single signer (API)
+CREATE TABLE ticket_verification_tmp (
+   ticket_id           INTEGER NOT NULL REFERENCES received_ticket_tmp(id),
+   signer_id           INTEGER NOT NULL,
+   verified_at         TIMESTAMP WITHOUT TIME ZONE NOT NULL,
+   accepted            BOOLEAN NOT NULL,
+
+   PRIMARY KEY (ticket_id, signer_id)
+);
+
+DROP INDEX ticket_verification_index;
+CREATE INDEX ticket_verification_index ON ticket_verification_tmp (ticket_id);
+
+DROP TABLE ticket_verification;
+ALTER TABLE ticket_verification_tmp RENAME TO ticket_verification;
+
+-- verified tickets that are yet to be redeemed
+CREATE TABLE verified_tickets_tmp (
+   ticket_id           INTEGER NOT NULL PRIMARY KEY REFERENCES received_ticket_tmp(id),
+   proposal_id         INTEGER REFERENCES redemption_proposals(proposal_id)
+);
+
+DROP INDEX verified_tickets_index;
+CREATE INDEX verified_tickets_index ON verified_tickets_tmp (proposal_id);
+
+DROP TABLE verified_tickets;
+ALTER TABLE verified_tickets_tmp RENAME TO verified_tickets;
+
+DROP TABLE received_ticket;
+ALTER TABLE received_ticket_tmp RENAME TO received_ticket;
+
+DROP TABLE shared_keys;
+ALTER TABLE shared_keys_tmp RENAME TO shared_keys;

common/gateway-storage/src/clients.rs

@@ -0,0 +1,89 @@
+// Copyright 2024 - Nym Technologies SA <[email protected]>
+// SPDX-License-Identifier: GPL-3.0-only
+
+use std::str::FromStr;
+
+use crate::models::Client;
+
+#[derive(Debug, PartialEq, sqlx::Type)]
+#[sqlx(type_name = "TEXT")] // SQLite TEXT type
+pub enum ClientType {
+    EntryMixnet,
+    ExitMixnet,
+    EntryWireguard,
+    ExitWireguard,
+}
+
+impl FromStr for ClientType {
+    type Err = &'static str;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "entry_mixnet" => Ok(ClientType::EntryMixnet),
+            "exit_mixnet" => Ok(ClientType::ExitMixnet),
+            "entry_wireguard" => Ok(ClientType::EntryWireguard),
+            "exit_wireguard" => Ok(ClientType::ExitWireguard),
+            _ => Err("Invalid client type"),
+        }
+    }
+}
+
+impl std::fmt::Display for ClientType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let s = match self {
+            ClientType::EntryMixnet => "entry_mixnet",
+            ClientType::ExitMixnet => "exit_mixnet",
+            ClientType::EntryWireguard => "entry_wireguard",
+            ClientType::ExitWireguard => "exit_wireguard",
+        };
+        write!(f, "{}", s)
+    }
+}
+
+#[derive(Clone)]
+pub(crate) struct ClientManager {
+    connection_pool: sqlx::SqlitePool,
+}
+
+impl ClientManager {
+    /// Creates new instance of the `ClientManager` with the provided sqlite connection pool.
+    ///
+    /// # Arguments
+    ///
+    /// * `connection_pool`: database connection pool to use.
+    pub(crate) fn new(connection_pool: sqlx::SqlitePool) -> Self {
+        ClientManager { connection_pool }
+    }
+
+    /// Inserts new client to the storage, specifying its type.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_type`: Type of the client that gets inserted
+    pub(crate) async fn insert_client(&self, client_type: ClientType) -> Result<i64, sqlx::Error> {
+        let client_id = sqlx::query!("INSERT INTO clients(client_type) VALUES (?)", client_type)
+            .execute(&self.connection_pool)
+            .await?
+            .last_insert_rowid();
+        Ok(client_id)
+    }
+
+    /// Tries to retrieve a particular client.
+    ///
+    /// # Arguments
+    ///
+    /// * `id`: The client id
+    pub(crate) async fn get_client(&self, id: i64) -> Result<Option<Client>, sqlx::Error> {
+        sqlx::query_as!(
+            Client,
+            r#"
+            SELECT id, client_type as "client_type: ClientType"
+            FROM clients
+            WHERE id = ?
+            "#,
+            id
+        )
+        .fetch_optional(&self.connection_pool)
+        .await
+    }
+}

common/gateway-storage/src/lib.rs

@@ -3,11 +3,12 @@
 
 use async_trait::async_trait;
 use bandwidth::BandwidthManager;
+use clients::{ClientManager, ClientType};
 use error::StorageError;
 use inboxes::InboxManager;
 use models::{
-    PersistedBandwidth, PersistedSharedKeys, RedemptionProposal, StoredMessage, VerifiedTicket,
-    WireguardPeer,
+    Client, PersistedBandwidth, PersistedSharedKeys, RedemptionProposal, StoredMessage,
+    VerifiedTicket, WireguardPeer,
 };
 use nym_credentials_interface::ClientTicket;
 use nym_gateway_requests::registration::handshake::SharedKeys;
@@ -20,6 +21,7 @@ use time::OffsetDateTime;
 use tracing::{debug, error};
 
 pub mod bandwidth;
+mod clients;
 pub mod error;
 mod inboxes;
 pub mod models;
@@ -29,7 +31,7 @@ mod wireguard_peers;
 
 #[async_trait]
 pub trait Storage: Send + Sync {
-    async fn get_client_id(
+    async fn get_mixnet_client_id(
         &self,
         client_address: DestinationAddressBytes,
     ) -> Result<i64, StorageError>;
@@ -39,7 +41,7 @@ pub trait Storage: Send + Sync {
     ///
     /// # Arguments
     ///
-    /// * `client_address`: address of the client
+    /// * `client_address`: base58-encoded address of the client
     /// * `shared_keys`: shared encryption (AES128CTR) and mac (hmac-blake3) derived shared keys to store.
     async fn insert_shared_keys(
         &self,
@@ -70,6 +72,14 @@ pub trait Storage: Send + Sync {
         client_address: DestinationAddressBytes,
     ) -> Result<(), StorageError>;
 
+    /// Tries to retrieve a particular client.
+    ///
+    /// # Arguments
+    ///
+    /// * `client_id`: id of the client
+    #[allow(dead_code)]
+    async fn get_client(&self, client_id: i64) -> Result<Option<Client>, StorageError>;
+
     /// Inserts new message to the storage for an offline client for future retrieval.
     ///
     /// # Arguments
@@ -246,6 +256,7 @@ pub trait Storage: Send + Sync {
 // note that clone here is fine as upon cloning the same underlying pool will be used
 #[derive(Clone)]
 pub struct PersistentStorage {
+    client_manager: ClientManager,
     shared_key_manager: SharedKeysManager,
     inbox_manager: InboxManager,
     bandwidth_manager: BandwidthManager,
@@ -294,6 +305,7 @@ impl PersistentStorage {
 
         // the cloning here are cheap as connection pool is stored behind an Arc
         Ok(PersistentStorage {
+            client_manager: clients::ClientManager::new(connection_pool.clone()),
             wireguard_peer_manager: wireguard_peers::WgPeerManager::new(connection_pool.clone()),
             shared_key_manager: SharedKeysManager::new(connection_pool.clone()),
             inbox_manager: InboxManager::new(connection_pool.clone(), message_retrieval_limit),
@@ -305,7 +317,7 @@ impl PersistentStorage {
 
 #[async_trait]
 impl Storage for PersistentStorage {
-    async fn get_client_id(
+    async fn get_mixnet_client_id(
         &self,
         client_address: DestinationAddressBytes,
     ) -> Result<i64, StorageError> {
@@ -321,8 +333,12 @@ impl Storage for PersistentStorage {
         shared_keys: &SharedKeys,
     ) -> Result<i64, StorageError> {
         let client_id = self
-            .shared_key_manager
+            .client_manager
+            .insert_client(ClientType::EntryMixnet)
+            .await?;
+        self.shared_key_manager
             .insert_shared_keys(
+                client_id,
                 client_address.as_base58_string(),
                 shared_keys.to_base58_string(),
             )
@@ -352,6 +368,11 @@ impl Storage for PersistentStorage {
         Ok(())
     }
 
+    async fn get_client(&self, client_id: i64) -> Result<Option<Client>, StorageError> {
+        let client = self.client_manager.get_client(client_id).await?;
+        Ok(client)
+    }
+
     async fn store_message(
         &self,
         client_address: DestinationAddressBytes,

common/gateway-storage/src/models.rs

@@ -6,9 +6,14 @@ use nym_credentials_interface::{AvailableBandwidth, ClientTicket, CredentialSpen
 use sqlx::FromRow;
 use time::OffsetDateTime;
 
+pub struct Client {
+    pub id: i64,
+    pub client_type: crate::clients::ClientType,
+}
+
 pub struct PersistedSharedKeys {
     #[allow(dead_code)]
-    pub id: i64,
+    pub client_id: i64,
 
     #[allow(dead_code)]
     pub client_address_bs58: String,

common/gateway-storage/src/shared_keys.rs

@@ -20,12 +20,12 @@ impl SharedKeysManager {
 
     pub(crate) async fn client_id(&self, client_address_bs58: &str) -> Result<i64, sqlx::Error> {
         let client_id = sqlx::query!(
-            "SELECT id FROM shared_keys WHERE client_address_bs58 = ?",
+            "SELECT client_id FROM shared_keys WHERE client_address_bs58 = ?",
             client_address_bs58
         )
         .fetch_one(&self.connection_pool)
         .await?
-        .id;
+        .client_id;
         Ok(client_id)
     }
 
@@ -34,26 +34,30 @@ impl SharedKeysManager {
     ///
     /// # Arguments
     ///
-    /// * `shared_keys`: shared encryption (AES128CTR) and mac (hmac-blake3) derived shared keys to store.
+    /// * `client_id`: The client id for which the shared keys are stored
+    /// * `client_address_bs58`: base58-encoded address of the client
+    /// * `derived_aes128_ctr_blake3_hmac_keys_bs58`: shared encryption (AES128CTR) and mac (hmac-blake3) derived shared keys to store.
     pub(crate) async fn insert_shared_keys(
         &self,
+        client_id: i64,
         client_address_bs58: String,
         derived_aes128_ctr_blake3_hmac_keys_bs58: String,
-    ) -> Result<i64, sqlx::Error> {
+    ) -> Result<(), sqlx::Error> {
         // https://stackoverflow.com/a/20310838
         // we don't want to be using `INSERT OR REPLACE INTO` due to the foreign key on `available_bandwidth` if the entry already exists
         sqlx::query!(
             r#"
-                INSERT OR IGNORE INTO shared_keys(client_address_bs58, derived_aes128_ctr_blake3_hmac_keys_bs58) VALUES (?, ?);
+                INSERT OR IGNORE INTO shared_keys(client_id, client_address_bs58, derived_aes128_ctr_blake3_hmac_keys_bs58) VALUES (?, ?, ?);
                 UPDATE shared_keys SET derived_aes128_ctr_blake3_hmac_keys_bs58 = ? WHERE client_address_bs58 = ?
             "#,
+            client_id,
             client_address_bs58,
             derived_aes128_ctr_blake3_hmac_keys_bs58,
             derived_aes128_ctr_blake3_hmac_keys_bs58,
             client_address_bs58,
         ).execute(&self.connection_pool).await?;
 
-        self.client_id(&client_address_bs58).await
+        Ok(())
     }
 
     /// Tries to retrieve shared keys stored for the particular client.

gateway/src/node/client_handling/websocket/connection_handler/fresh.rs

@@ -570,7 +570,11 @@ where
             return Ok(InitialAuthResult::new_failed(Some(negotiated_protocol)));
         };
 
-        let client_id = self.shared_state.storage.get_client_id(address).await?;
+        let client_id = self
+            .shared_state
+            .storage
+            .get_mixnet_client_id(address)
+            .await?;
 
         let available_bandwidth: AvailableBandwidth = self
             .shared_state