Generate polling token on instance update

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(())
+    }
 }

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),
         })

src/grpc/enrollment.rs

@@ -609,7 +609,8 @@ impl EnrollmentServer {
     ) -> Result<DeviceConfigResponse, Status> {
         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
     }
 }
 

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<DeviceConfigResponse, Status> {
     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,
     })
 }