VPN-4401: Cache JWT skew

nym-vpn-core/crates/nym-vpn-api-client/src/client.rs

@@ -1,15 +1,19 @@
 // Copyright 2024 - Nym Technologies SA <[email protected]>
 // SPDX-License-Identifier: GPL-3.0-only
 
-use std::time::Duration;
+use std::{
+    sync::Arc,
+    time::{Duration, Instant},
+};
 
 use backon::Retryable;
 use nym_credential_proxy_requests::api::v1::ticketbook::models::PartialVerificationKeysResponse;
 use nym_http_api_client::{
     ApiClient, Client, HttpClientError, NO_PARAMS, Params, PathSegments, Url, UserAgent,
 };
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
-use time::OffsetDateTime;
+use time::{Duration as TimeDuration, OffsetDateTime};
+use tokio::sync::RwLock;
 
 use crate::{
     ResolverOverrides, api_urls_to_urls,
@@ -40,11 +44,48 @@ pub(crate) const DEVICE_AUTHORIZATION_HEADER: &str = "x-device-authorization";
 // GET requests can unfortunately take a long time over the mixnet
 pub(crate) const NYM_VPN_API_TIMEOUT: Duration = Duration::from_secs(60);
 
+const SKEW_CACHE_TTL: Duration = Duration::from_secs(4 * 60 * 60); // 4 hours
+
+#[derive(Debug)]
+struct SkewState {
+    skew: TimeDuration,
+    expires_at: Instant,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum SkewStatus {
+    Expired(),
+    Valid(TimeDuration),
+}
+
+impl SkewState {
+    fn new(skew: TimeDuration, now: Instant) -> Self {
+        Self {
+            skew,
+            expires_at: now + SKEW_CACHE_TTL,
+        }
+    }
+
+    fn update(&mut self, skew: TimeDuration, now: Instant) {
+        self.skew = skew;
+        self.expires_at = now + SKEW_CACHE_TTL;
+    }
+
+    fn status(&self, now: Instant) -> SkewStatus {
+        if self.expires_at > now {
+            SkewStatus::Valid(self.skew)
+        } else {
+            SkewStatus::Expired()
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct VpnApiClient {
     inner: Client,
     urls: Vec<Url>,
     user_agent: UserAgent,
+    skew_state: Arc<RwLock<Option<SkewState>>>,
 }
 
 impl VpnApiClient {
@@ -65,6 +106,7 @@ impl VpnApiClient {
             inner,
             urls,
             user_agent,
+            skew_state: Arc::new(RwLock::new(None)),
         })
     }
 
@@ -96,6 +138,7 @@ impl VpnApiClient {
             inner,
             urls,
             user_agent,
+            skew_state: Arc::new(RwLock::new(None)),
         })
     }
 
@@ -153,8 +196,55 @@ impl VpnApiClient {
         }
     }
 
-    async fn sync_with_remote_time(&self) -> Result<Option<VpnApiTime>> {
+    async fn refresh_skew(&self) -> Result<VpnApiTime> {
         let remote_time = self.get_remote_time().await?;
+        let skew = remote_time.local_time_ahead_skew();
+        let now = Instant::now();
+
+        {
+            let mut skew_state = self.skew_state.write().await;
+            match skew_state.as_mut() {
+                Some(state) => state.update(skew, now),
+                None => *skew_state = Some(SkewState::new(skew, now)),
+            }
+        }
+
+        tracing::debug!(skew = ?skew, "Refreshed VPN API time skew");
+
+        Ok(remote_time)
+    }
+
+    async fn current_remote_time(&self) -> Result<Option<VpnApiTime>> {
+        let now = Instant::now();
+        let status = {
+            let state = self.skew_state.read().await;
+            state.as_ref().map(|state| state.status(now))
+        };
+
+        let cached_remote_time = match status {
+            Some(SkewStatus::Valid(skew)) => {
+                tracing::debug!("Valid VPN API time skew");
+                let local_time = OffsetDateTime::now_utc();
+                let estimated_remote_time = local_time - skew;
+
+                VpnApiTime::from_estimated_remote_time(local_time, estimated_remote_time)
+            }
+            Some(SkewStatus::Expired()) | None => {
+                tracing::debug!("VPN API time skew expired or not present, refreshing");
+
+                self.refresh_skew().await?
+            }
+        };
+
+        Ok(if Self::use_remote_time(cached_remote_time) {
+            Some(cached_remote_time)
+        } else {
+            None
+        })
+    }
+
+    async fn sync_with_remote_time(&self) -> Result<Option<VpnApiTime>> {
+        let remote_time = self.refresh_skew().await?;
 
         if Self::use_remote_time(remote_time) {
             Ok(Some(remote_time))
@@ -198,7 +288,18 @@ impl VpnApiClient {
     where
         T: DeserializeOwned,
     {
-        match self.get_query::<T>(path, account, device, None).await {
+        let jwt = match self.current_remote_time().await {
+            Ok(remote_time) => remote_time,
+            Err(err) => {
+                tracing::debug!(
+                    error = %err,
+                    "Failed to determine cached remote time"
+                );
+                None
+            }
+        };
+
+        match self.get_query::<T>(path, account, device, jwt).await {
             Ok(response) => Ok(response),
             Err(err) => {
                 if let HttpClientError::EndpointFailure { error, .. } = &err
@@ -364,8 +465,19 @@ impl VpnApiClient {
         T: DeserializeOwned,
         B: Serialize,
     {
+        let jwt = match self.current_remote_time().await {
+            Ok(remote_time) => remote_time,
+            Err(err) => {
+                tracing::debug!(
+                    error = %err,
+                    "Failed to determine cached remote time"
+                );
+                None
+            }
+        };
+
         match self
-            .post_query::<T, B>(path, json_body, account, device, None)
+            .post_query::<T, B>(path, json_body, account, device, jwt)
             .await
         {
             Ok(response) => Ok(response),