From 5c80d1a11de2a2396b10b9a644b7fa5d01487066 Mon Sep 17 00:00:00 2001
From: Pranay Tulugu <pranay@althea.systems>
Date: Thu, 3 Aug 2023 23:05:49 -0700
Subject: [PATCH] Update exit checkin to send a list of registered wgkey

The exit checkin now sends a list of WgKey it retrieves from the database during exit checkin. Ops can use this
to display more accurately whether a router is registered or not.
---
 althea_types/src/interop.rs                  |  2 +
 rita_exit/src/operator_update/mod.rs         | 41 ++++++++++++++++++++
 rita_exit/src/operator_update/update_loop.rs |  1 +
 3 files changed, 44 insertions(+)

diff --git a/althea_types/src/interop.rs b/althea_types/src/interop.rs
index 5d9e2575a..00961d91e 100644
--- a/althea_types/src/interop.rs
+++ b/althea_types/src/interop.rs
@@ -770,6 +770,8 @@ pub struct OperatorExitCheckinMessage {
     pub pass: String,
     /// This is to keep track of the rita exit uptime for debugging purposes
     pub exit_uptime: Duration,
+    /// A list of registered wg keys that ops can use to display routers to be registered
+    pub registered_keys: Option<Vec<WgKey>>,
     /// Number of users online
     pub users_online: Option<u32>,
 }
diff --git a/rita_exit/src/operator_update/mod.rs b/rita_exit/src/operator_update/mod.rs
index 13c8a5a85..09ebbd373 100644
--- a/rita_exit/src/operator_update/mod.rs
+++ b/rita_exit/src/operator_update/mod.rs
@@ -4,10 +4,16 @@ pub mod update_loop;
 use althea_types::ExitClientIdentity;
 use althea_types::OperatorExitCheckinMessage;
 use althea_types::OperatorExitUpdateMessage;
+use althea_types::WgKey;
+use diesel::QueryDsl;
+use diesel::RunQueryDsl;
+use exit_db::schema::clients::dsl::clients as db_client;
+use exit_db::schema::clients::wg_pubkey;
 use rita_common::KI;
 use std::time::{Duration, Instant};
 
 use crate::database::signup_client;
+use crate::get_database_connection;
 use crate::rita_loop::EXIT_INTERFACE;
 
 pub struct UptimeStruct {
@@ -52,6 +58,7 @@ pub async fn operator_update(rita_started: Instant) {
                 id,
                 pass,
                 exit_uptime: rita_started.elapsed(),
+                registered_keys: get_registered_list(),
                 // Since this checkin works only from b20, we only need to look on wg_exit_v2
                 users_online: KI.get_wg_exit_clients_online(EXIT_INTERFACE).ok(),
             })
@@ -93,3 +100,37 @@ async fn register_op_clients(clients: Vec<ExitClientIdentity>) {
         };
     }
 }
+
+pub fn get_registered_list() -> Option<Vec<WgKey>> {
+    match get_database_connection() {
+        Ok(conn) => {
+            let registered_routers = db_client.select(wg_pubkey);
+            let registered_routers = match registered_routers.load::<String>(&conn) {
+                Ok(a) => a,
+                Err(e) => {
+                    error!("Unable to retrive wg keys {}", e);
+                    return None;
+                }
+            };
+            Some(
+                registered_routers
+                    .iter()
+                    .filter_map(|r| match r.parse() {
+                        Ok(a) => Some(a),
+                        Err(_) => {
+                            error!("Invalid wg key in database! {}", r);
+                            None
+                        }
+                    })
+                    .collect::<Vec<WgKey>>(),
+            )
+        }
+        Err(e) => {
+            error!(
+                "Unable to get a database connection to retrieve registered exits: {}",
+                e
+            );
+            None
+        }
+    }
+}
diff --git a/rita_exit/src/operator_update/update_loop.rs b/rita_exit/src/operator_update/update_loop.rs
index ee62aad90..223922cef 100644
--- a/rita_exit/src/operator_update/update_loop.rs
+++ b/rita_exit/src/operator_update/update_loop.rs
@@ -13,6 +13,7 @@ pub fn start_operator_update_loop() {
     thread::spawn(move || {
         // this will always be an error, so it's really just a loop statement
         // with some fancy destructuring
+
         while let Err(e) = {
             thread::spawn(move || loop {
                 let start = Instant::now();