diff --git a/.sqlx/query-ace6a5da2485d065f1cd6c8d66ec5ee98d9d5e238fbce7d91c899163d0311234.json b/.sqlx/query-ace6a5da2485d065f1cd6c8d66ec5ee98d9d5e238fbce7d91c899163d0311234.json new file mode 100644 index 0000000000..b5eae9a2d5 --- /dev/null +++ b/.sqlx/query-ace6a5da2485d065f1cd6c8d66ec5ee98d9d5e238fbce7d91c899163d0311234.json @@ -0,0 +1,14 @@ +{ + "db_name": "PostgreSQL", + "query": "DELETE FROM pollingtoken WHERE device_id = $1", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Int8" + ] + }, + "nullable": [] + }, + "hash": "ace6a5da2485d065f1cd6c8d66ec5ee98d9d5e238fbce7d91c899163d0311234" +} diff --git a/src/db/models/polling_token.rs b/src/db/models/polling_token.rs index 99f611909d..6f8f888e3d 100644 --- a/src/db/models/polling_token.rs +++ b/src/db/models/polling_token.rs @@ -1,6 +1,6 @@ use chrono::{NaiveDateTime, Utc}; use model_derive::Model; -use sqlx::{query_as, Error as SqlxError}; +use sqlx::{query_as, Error as SqlxError, PgExecutor}; use super::DbPool; use crate::random::gen_alphanumeric; @@ -34,4 +34,14 @@ impl PollingToken { .fetch_optional(pool) .await } + + pub async fn delete_for_device_id<'e, E>(executor: E, device_id: i64) -> Result<(), SqlxError> + where + E: PgExecutor<'e>, + { + sqlx::query!("DELETE FROM pollingtoken WHERE device_id = $1", device_id) + .execute(executor) + .await?; + Ok(()) + } } diff --git a/src/enterprise/grpc/polling.rs b/src/enterprise/grpc/polling.rs index 79964e5b52..9780a90457 100644 --- a/src/enterprise/grpc/polling.rs +++ b/src/enterprise/grpc/polling.rs @@ -82,7 +82,7 @@ impl PollingServer { // Build & return polling info let device_config = - build_device_config_response(&self.pool, &device.wireguard_pubkey).await?; + build_device_config_response(&self.pool, &device.wireguard_pubkey, false).await?; Ok(InstanceInfoResponse { device_config: Some(device_config), }) diff --git a/src/grpc/enrollment.rs b/src/grpc/enrollment.rs index 656ca07d2b..48a3eae208 100644 --- a/src/grpc/enrollment.rs +++ b/src/grpc/enrollment.rs @@ -609,7 +609,8 @@ impl EnrollmentServer { ) -> Result { debug!("Getting network info for device: {:?}", request.pubkey); let _token = self.validate_session(&request.token).await?; - build_device_config_response(&self.pool, &request.pubkey).await + + build_device_config_response(&self.pool, &request.pubkey, true).await } } diff --git a/src/grpc/utils.rs b/src/grpc/utils.rs index 87cb177cab..6fade90a88 100644 --- a/src/grpc/utils.rs +++ b/src/grpc/utils.rs @@ -7,7 +7,10 @@ use super::{ }; use crate::{ db::{ - models::{device::WireguardNetworkDevice, wireguard::WireguardNetwork}, + models::{ + device::WireguardNetworkDevice, polling_token::PollingToken, + wireguard::WireguardNetwork, + }, DbPool, Device, Settings, User, }, enterprise::db::models::enterprise_settings::EnterpriseSettings, @@ -16,6 +19,8 @@ use crate::{ pub(crate) async fn build_device_config_response( pool: &DbPool, pubkey: &str, + // Whether to make a new polling token for the device + new_token: bool, ) -> Result { Device::validate_pubkey(pubkey).map_err(|_| { error!("Invalid pubkey {pubkey}"); @@ -91,6 +96,59 @@ pub(crate) async fn build_device_config_response( } } + let token = if new_token { + debug!( + "Making a new polling token for device {}", + device.wireguard_pubkey + ); + let mut transaction = pool.begin().await.map_err(|err| { + error!("Failed to start transaction while making a new polling token: {err}"); + Status::internal(format!("unexpected error: {err}")) + })?; + + // 1. Delete existing polling token for the device, if it exists + // 2. Create a new polling token for the device + PollingToken::delete_for_device_id( + &mut *transaction, + device.id.ok_or_else(|| { + error!( + "Device {} has no id, can't delete polling token", + device.wireguard_pubkey + ); + Status::internal("unexpected error") + })?, + ) + .await + .map_err(|err| { + error!("Failed to delete polling token: {err}"); + Status::internal(format!("unexpected error: {err}")) + })?; + let mut new_token = PollingToken::new(device.id.ok_or_else(|| { + error!( + "Device {} has no id, can't create a polling token", + device.wireguard_pubkey + ); + Status::internal("unexpected error") + })?); + new_token.save(&mut *transaction).await.map_err(|err| { + error!("Failed to save new polling token: {err}"); + Status::internal(format!("unexpected error: {err}")) + })?; + + transaction.commit().await.map_err(|err| { + error!("Failed to commit transaction while making a new polling token: {err}"); + Status::internal(format!("unexpected error: {err}")) + })?; + info!( + "New polling token created for device {}", + device.wireguard_pubkey + ); + + Some(new_token.token) + } else { + None + }; + info!( "User {}({:?}) device {}({:?}) config fetched", user.username, user.id, device.name, device.id, @@ -100,6 +158,6 @@ pub(crate) async fn build_device_config_response( device: Some(device.into()), configs, instance: Some(InstanceInfo::new(settings, &user.username, enterprise_settings).into()), - token: None, + token, }) }