From 1224f18ff9e7b675bfecb9482389f18ea58e46fa Mon Sep 17 00:00:00 2001
From: Macpie <pierre@nova-labs.com>
Date: Wed, 16 Oct 2024 12:31:39 -0700
Subject: [PATCH] Update location cache to handle cbrs and wifi

---
 .../src/heartbeats/last_location.rs           | 83 ++++++++++---------
 mobile_verifier/src/heartbeats/mod.rs         |  5 +-
 2 files changed, 45 insertions(+), 43 deletions(-)

diff --git a/mobile_verifier/src/heartbeats/last_location.rs b/mobile_verifier/src/heartbeats/last_location.rs
index cb6e45c76..794926251 100644
--- a/mobile_verifier/src/heartbeats/last_location.rs
+++ b/mobile_verifier/src/heartbeats/last_location.rs
@@ -1,9 +1,14 @@
-use std::sync::Arc;
-
 use chrono::{DateTime, Duration, Utc};
 use helium_crypto::PublicKeyBinary;
 use retainer::Cache;
 use sqlx::PgPool;
+use std::sync::Arc;
+
+#[derive(Debug, Clone, Eq, Ord, PartialEq, PartialOrd)]
+pub enum Key {
+    CbrsId(String),
+    WifiPubKey(PublicKeyBinary),
+}
 
 #[derive(sqlx::FromRow, Copy, Clone)]
 pub struct LastLocation {
@@ -38,7 +43,7 @@ impl LastLocation {
 #[derive(Clone)]
 pub struct LocationCache {
     pool: PgPool,
-    locations: Arc<Cache<PublicKeyBinary, Option<LastLocation>>>,
+    locations: Arc<Cache<Key, Option<LastLocation>>>,
 }
 
 impl LocationCache {
@@ -56,9 +61,39 @@ impl LocationCache {
         }
     }
 
-    async fn fetch_from_db_and_set(
+    pub async fn fetch_last_location(&self, key: Key) -> anyhow::Result<Option<LastLocation>> {
+        Ok(
+            if let Some(last_location) = self.locations.get(&key).await {
+                *last_location
+            } else {
+                match key {
+                    Key::WifiPubKey(pub_key_bin) => self.fetch_wifi_and_set(pub_key_bin).await?,
+                    Key::CbrsId(_) => None,
+                }
+            },
+        )
+    }
+
+    pub async fn set_last_location(
+        &self,
+        key: Key,
+        last_location: LastLocation,
+    ) -> anyhow::Result<()> {
+        let duration_to_expiration = last_location.duration_to_expiration();
+        self.locations
+            .insert(key, Some(last_location), duration_to_expiration.to_std()?)
+            .await;
+        Ok(())
+    }
+
+    /// Only used for testing.
+    pub async fn delete_last_location(&self, key: Key) {
+        self.locations.remove(&key).await;
+    }
+
+    async fn fetch_wifi_and_set(
         &self,
-        hotspot: &PublicKeyBinary,
+        pub_key_bin: PublicKeyBinary,
     ) -> anyhow::Result<Option<LastLocation>> {
         let last_location: Option<LastLocation> = sqlx::query_as(
             r#"
@@ -72,12 +107,12 @@ impl LocationCache {
             "#,
         )
         .bind(Utc::now() - Duration::hours(12))
-        .bind(hotspot)
+        .bind(pub_key_bin.clone())
         .fetch_optional(&self.pool)
         .await?;
         self.locations
             .insert(
-                hotspot.clone(),
+                Key::WifiPubKey(pub_key_bin),
                 last_location,
                 last_location
                     .map(|x| x.duration_to_expiration())
@@ -87,38 +122,4 @@ impl LocationCache {
             .await;
         Ok(last_location)
     }
-
-    pub async fn fetch_last_location(
-        &self,
-        hotspot: &PublicKeyBinary,
-    ) -> anyhow::Result<Option<LastLocation>> {
-        Ok(
-            if let Some(last_location) = self.locations.get(hotspot).await {
-                *last_location
-            } else {
-                self.fetch_from_db_and_set(hotspot).await?
-            },
-        )
-    }
-
-    pub async fn set_last_location(
-        &self,
-        hotspot: &PublicKeyBinary,
-        last_location: LastLocation,
-    ) -> anyhow::Result<()> {
-        let duration_to_expiration = last_location.duration_to_expiration();
-        self.locations
-            .insert(
-                hotspot.clone(),
-                Some(last_location),
-                duration_to_expiration.to_std()?,
-            )
-            .await;
-        Ok(())
-    }
-
-    /// Only used for testing.
-    pub async fn delete_last_location(&self, hotspot: &PublicKeyBinary) {
-        self.locations.remove(hotspot).await;
-    }
 }
diff --git a/mobile_verifier/src/heartbeats/mod.rs b/mobile_verifier/src/heartbeats/mod.rs
index 932207f3d..afc107dbd 100644
--- a/mobile_verifier/src/heartbeats/mod.rs
+++ b/mobile_verifier/src/heartbeats/mod.rs
@@ -19,6 +19,7 @@ use futures::stream::{Stream, StreamExt};
 use h3o::{CellIndex, LatLng};
 use helium_crypto::PublicKeyBinary;
 use helium_proto::services::poc_mobile::{self as proto, LocationSource};
+use last_location::Key;
 use retainer::Cache;
 use rust_decimal::{prelude::ToPrimitive, Decimal};
 use rust_decimal_macros::dec;
@@ -521,7 +522,7 @@ impl ValidatedHeartbeat {
                 let is_valid = match heartbeat.location_validation_timestamp {
                     None => {
                         if let Some(last_location) = last_location_cache
-                            .fetch_last_location(&heartbeat.hotspot_key)
+                            .fetch_last_location(Key::WifiPubKey(heartbeat.hotspot_key.clone()))
                             .await?
                         {
                             heartbeat.lat = last_location.lat;
@@ -538,7 +539,7 @@ impl ValidatedHeartbeat {
                     Some(location_validation_timestamp) => {
                         last_location_cache
                             .set_last_location(
-                                &heartbeat.hotspot_key,
+                                Key::WifiPubKey(heartbeat.hotspot_key.clone()),
                                 LastLocation::new(
                                     location_validation_timestamp,
                                     heartbeat.timestamp,