nym/common/gateway-storage/src/shared_keys.rs

// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: GPL-3.0-only

use crate::models::PersistedSharedKeys;
use time::OffsetDateTime;

#[derive(Clone)]
pub(crate) struct SharedKeysManager {
    connection_pool: sqlx::SqlitePool,
}

impl SharedKeysManager {
    /// Creates new instance of the `SharedKeysManager` with the provided sqlite connection pool.
    ///
    /// # Arguments
    ///
    /// * `connection_pool`: database connection pool to use.
    pub(crate) fn new(connection_pool: sqlx::SqlitePool) -> Self {
        SharedKeysManager { connection_pool }
    }

    pub(crate) async fn client_id(&self, client_address_bs58: &str) -> Result<i64, sqlx::Error> {
        let client_id = sqlx::query!(
            "SELECT client_id FROM shared_keys WHERE client_address_bs58 = ?",
            client_address_bs58
        )
        .fetch_one(&self.connection_pool)
        .await?
        .client_id;
        Ok(client_id)
    }

    /// Inserts provided derived shared keys into the database.
    /// If keys previously existed for the provided client, they are overwritten with the new data.
    ///
    /// # Arguments
    ///
    /// * `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: Option<&String>,
        derived_aes256_gcm_siv_key: Option<&Vec<u8>>,
    ) -> 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_id, client_address_bs58, derived_aes128_ctr_blake3_hmac_keys_bs58, derived_aes256_gcm_siv_key) VALUES (?, ?, ?, ?);

                UPDATE shared_keys
                SET
                    derived_aes128_ctr_blake3_hmac_keys_bs58 = ?,
                    derived_aes256_gcm_siv_key = ?
                WHERE client_address_bs58 = ?
            "#,
            client_id,
            client_address_bs58,
            derived_aes128_ctr_blake3_hmac_keys_bs58,
            derived_aes256_gcm_siv_key,
            derived_aes128_ctr_blake3_hmac_keys_bs58,
            derived_aes256_gcm_siv_key,
            client_address_bs58,
        ).execute(&self.connection_pool).await?;

        Ok(())
    }

    pub(crate) async fn update_last_used_authentication_timestamp(
        &self,
        client_id: i64,
        last_used: OffsetDateTime,
    ) -> Result<(), sqlx::Error> {
        sqlx::query!(
            "UPDATE shared_keys SET last_used_authentication = ? WHERE client_id = ?;",
            last_used,
            client_id
        )
        .execute(&self.connection_pool)
        .await?;

        Ok(())
    }

    /// Tries to retrieve shared keys stored for the particular client.
    ///
    /// # Arguments
    ///
    /// * `client_address_bs58`: base58-encoded address of the client
    pub(crate) async fn get_shared_keys(
        &self,
        client_address_bs58: &str,
    ) -> Result<Option<PersistedSharedKeys>, sqlx::Error> {
        sqlx::query_as("SELECT * FROM shared_keys WHERE client_address_bs58 = ?")
            .bind(client_address_bs58)
            .fetch_optional(&self.connection_pool)
            .await
    }

    /// Removes from the database shared keys derived with the particular client.
    ///
    /// # Arguments
    ///
    /// * `client_address_bs58`: base58-encoded address of the client
    // currently there is no code flow that causes removal (not overwriting)
    // of the stored keys. However, retain the function for consistency and completion sake
    #[allow(dead_code)]
    pub(crate) async fn remove_shared_keys(
        &self,
        client_address_bs58: &str,
    ) -> Result<(), sqlx::Error> {
        sqlx::query!(
            "DELETE FROM shared_keys WHERE client_address_bs58 = ?",
            client_address_bs58
        )
        .execute(&self.connection_pool)
        .await?;
        Ok(())
    }
}