ditch hardcoded pairs

harsh-ps-2003 · harsh-ps-2003 · commit 6ee535769401 · 2025-01-24T23:00:58.000+05:30
diff --git a/.gitignore b/.gitignore
@@ -4,3 +4,5 @@ target/
 *.key
 *.crt
 *.log
+
+.codegpt
diff --git a/build.rs b/build.rs
@@ -1,6 +1,4 @@
 fn main() -> Result<(), Box<dyn std::error::Error>> {
-    tonic_build::configure()
-        .type_attribute(".", "#[derive(Hash, Eq, PartialOrd, Ord)]")
-        .compile(&["proto/rumi.proto"], &["proto"])?;
+    tonic_build::compile_protos("proto/rumi.proto")?;
     Ok(())
 }
diff --git a/proto/rumi.proto b/proto/rumi.proto
@@ -1,29 +1,53 @@
 syntax = "proto3";
 package rumi;
 
+// The discovery service definition
 service Discovery {
+    // Get the set of registered identifiers
     rpc GetPublicSet (GetPublicSetRequest) returns (GetPublicSetResponse);
+    
+    // Find a user by their identifier
     rpc Find (FindRequest) returns (FindResponse);
+    
+    // Register a new identifier-UUID pair
+    rpc Register(RegisterRequest) returns (RegisterResponse);
 }
 
+// Request for getting the public set of identifiers
 message GetPublicSetRequest {}
 
+// Response containing the public set of identifiers
 message GetPublicSetResponse {
     repeated uint64 identifiers = 1;
 }
 
+// Request for finding a user
 message FindRequest {
     bytes hash_prefix = 1;
     bytes blinded_identifier = 2;
     string zksm_proof = 3;
 }
 
+// A bucket entry in the response
+message BucketEntry {
+    bytes blinded_identifier = 1;
+    bytes blinded_user_id = 2;
+}
+
+// Response for finding a user
 message FindResponse {
     bytes double_blinded_identifier = 1;
     repeated BucketEntry entries = 2;
 }
 
-message BucketEntry {
-    bytes blinded_identifier = 1;
-    bytes blinded_user_id = 2;
+// Request for registering a new identifier-UUID pair
+message RegisterRequest {
+    uint64 identifier = 1;
+    bytes uuid = 2;
+}
+
+// Response for registration
+message RegisterResponse {
+    bool success = 1;
+    string message = 2;
 }
diff --git a/src/bin/client.rs b/src/bin/client.rs
@@ -1,15 +1,19 @@
 use clap::Parser;
 use console::style;
 use rumi::Client;
-use rumi_proto::discovery_client::DiscoveryClient;
-use rumi_proto::{FindRequest, FindResponse, GetPublicSetRequest, GetPublicSetResponse};
+use rumi_proto::{
+    discovery_client::DiscoveryClient,
+    FindRequest, FindResponse, GetPublicSetRequest, GetPublicSetResponse,
+    RegisterRequest, RegisterResponse,
+};
 use std::collections::HashMap;
 use std::error::Error;
 use std::vec;
 use tonic::Response;
 use tracing::{debug, info, trace, warn, Level};
 use tracing_attributes::instrument;
 use tracing_subscriber::{fmt, prelude::*};
+use uuid::Uuid;
 
 pub mod rumi_proto {
     tonic::include_proto!("rumi");
@@ -28,6 +32,12 @@ enum Commands {
         #[arg(help = "The identifier to look up")]
         identifier: u64,
     },
+    Register {
+        #[arg(help = "The identifier to register")]
+        identifier: u64,
+        #[arg(help = "The UUID to register (optional, will generate if not provided)")]
+        uuid: Option<String>,
+    },
 }
 
 #[instrument(skip(client, rumi_client), fields(identifier = %identifier), ret)]
@@ -117,27 +127,72 @@ async fn lookup_identifier(
     Ok(())
 }
 
+#[instrument(skip(client), fields(identifier = %identifier), ret)]
+async fn register_identifier(
+    client: &mut DiscoveryClient<tonic::transport::Channel>,
+    identifier: u64,
+    uuid_str: Option<String>,
+) -> Result<(), Box<dyn Error>> {
+    let uuid = if let Some(uuid_str) = uuid_str {
+        Uuid::parse_str(&uuid_str)?
+    } else {
+        Uuid::new_v4()
+    };
+
+    info!(
+        "Registering identifier {} with UUID {}",
+        style(identifier).cyan(),
+        style(uuid).yellow()
+    );
+
+    let request = tonic::Request::new(RegisterRequest {
+        identifier,
+        uuid: uuid.as_bytes().to_vec(),
+    });
+
+    match client.register(request).await {
+        Ok(response) => {
+            let response = response.into_inner();
+            if response.success {
+                info!(
+                    "{} {}",
+                    style("✓").green().bold(),
+                    style(&response.message).green()
+                );
+            } else {
+                info!(
+                    "{} {}",
+                    style("✗").red().bold(),
+                    style(&response.message).red()
+                );
+            }
+        }
+        Err(status) => {
+            warn!(
+                "{} Registration failed: {}",
+                style("✗").red().bold(),
+                style(status).red()
+            );
+        }
+    }
+
+    Ok(())
+}
+
 #[tokio::main(flavor = "multi_thread")]
 async fn main() -> Result<(), Box<dyn Error>> {
-    // Set up single combined subscriber
-    let console_layer = console_subscriber::ConsoleLayer::builder()
-        .with_default_env()
-        .spawn();
-
-    tracing_subscriber::registry()
-        .with(console_layer)
-        .with(
-            fmt::layer()
-                .with_target(false)
-                .with_thread_ids(false)
-                .with_line_number(false)
-                .with_level(true),
-        )
+    // Set up logging based on RUST_LOG env var, defaulting to info level
+    let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
+        .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
+
+    tracing_subscriber::fmt()
+        .with_target(false)
+        .with_thread_ids(false)
+        .with_line_number(false)
+        .with_level(true)
+        .with_env_filter(env_filter)
         .init();
 
-    info!("RUMI Client starting up");
-    info!("Tokio Console available on http://127.0.0.1:6669");
-
     let cli = Cli::parse();
 
     match cli.command {
@@ -146,6 +201,10 @@ async fn main() -> Result<(), Box<dyn Error>> {
             let rumi_client = Client::new(rand::thread_rng());
             lookup_identifier(&mut client, &rumi_client, identifier).await?;
         }
+        Commands::Register { identifier, uuid } => {
+            let mut client = DiscoveryClient::connect("http://[::1]:50051").await?;
+            register_identifier(&mut client, identifier, uuid).await?;
+        }
     }
 
     Ok(())
diff --git a/src/bin/server.rs b/src/bin/server.rs
@@ -1,4 +1,8 @@
-use crate::rumi_proto::{GetPublicSetRequest, GetPublicSetResponse};
+use crate::rumi_proto::{
+    discovery_server::{Discovery, DiscoveryServer},
+    FindRequest, FindResponse, GetPublicSetRequest, GetPublicSetResponse,
+    RegisterRequest, RegisterResponse,
+};
 use console::style;
 use lazy_static::lazy_static;
 use p256::EncodedPoint;
@@ -22,11 +26,6 @@ pub mod rumi_proto {
     tonic::include_proto!("rumi");
 }
 
-use rumi_proto::{
-    discovery_server::{Discovery, DiscoveryServer},
-    FindRequest, FindResponse,
-};
-
 // Define metrics
 lazy_static! {
     static ref REGISTRY: Registry = Registry::new();
@@ -63,16 +62,11 @@ pub struct DiscoveryService {
 impl DiscoveryService {
     pub fn new() -> Self {
         let mut rng = rand::thread_rng();
-        let mut users = HashMap::new();
-
-        for i in 0..100 {
-            users.insert(1_000_000_000 + i, Uuid::new_v4());
-        }
+        let users = HashMap::new();  // Start with empty user set
 
         let server = Server::new(&mut rng, &users);
         debug!(
-            "Server initialized with identifiers: {:?}",
-            server.get_public_set()
+            "Server initialized with empty user set",
         );
 
         Self {
@@ -168,10 +162,62 @@ impl Discovery for DiscoveryService {
         timer.observe_duration();
         result
     }
+
+    #[instrument(skip(self, request), name = "register", ret)]
+    async fn register(
+        &self,
+        request: Request<RegisterRequest>,
+    ) -> Result<Response<RegisterResponse>, Status> {
+        let timer = REQUEST_DURATION
+            .with_label_values(&["register"])
+            .start_timer();
+        REQUEST_COUNTER.with_label_values(&["register"]).inc();
+
+        let result = {
+            let request_inner = request.into_inner();
+            let identifier = request_inner.identifier;
+            let uuid_bytes = request_inner.uuid;
+
+            let uuid = Uuid::from_slice(&uuid_bytes)
+                .map_err(|_| Status::invalid_argument("Invalid UUID format"))?;
+
+            let mut rng = rand::thread_rng();
+            let mut server = self
+                .server
+                .lock()
+                .map_err(|_| Status::internal("Server lock poisoned"))?;
+
+            match server.register(identifier, &uuid, &mut rng) {
+                Ok(()) => Ok(Response::new(RegisterResponse {
+                    success: true,
+                    message: format!("Successfully registered identifier {}", identifier),
+                })),
+                Err(e) => Ok(Response::new(RegisterResponse {
+                    success: false,
+                    message: e.to_string(),
+                })),
+            }
+        };
+
+        timer.observe_duration();
+        result
+    }
 }
 
 #[tokio::main(flavor = "multi_thread")]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // Set up logging based on RUST_LOG env var, defaulting to info level
+    let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
+        .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
+
+    tracing_subscriber::fmt()
+        .with_target(false)
+        .with_thread_ids(false)
+        .with_line_number(false)
+        .with_level(true)
+        .with_env_filter(env_filter)
+        .init();
+
     let addr = "[::1]:50051".parse()?;
     let service = DiscoveryService::new();
 
@@ -184,6 +230,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .unwrap();
     REGISTRY.register(Box::new(MEMORY_GAUGE.clone())).unwrap();
 
+    info!("RUMI Server starting up on {}", style(addr).cyan());
+
     // Start metrics pushing in background
     tokio::spawn(async {
         let client = reqwest::Client::new();
@@ -217,41 +265,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                     warn!("Failed to push metrics: {}", e);
                 }
             }
-            tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
+
+            tokio::time::sleep(tokio::time::Duration::from_secs(15)).await;
         }
     });
 
-    // Set up console and tracing
-    let console_layer = console_subscriber::ConsoleLayer::builder()
-        .with_default_env()
-        .spawn();
-
-    tracing_subscriber::registry()
-        .with(console_layer)
-        .with(
-            fmt::layer()
-                .with_target(false)
-                .with_thread_ids(true)
-                .with_file(true)
-                .with_line_number(true)
-                .with_level(true),
-        )
-        .init();
-
-    info!("{}", style("RUMI Discovery Server").green().bold());
-    info!("Listening on {}", style(addr).cyan());
-    info!(
-        "Initialized with {} identifiers",
-        style(service.server.lock().unwrap().get_public_set().len()).yellow()
-    );
-    info!("Tokio Console available on http://127.0.0.1:6669");
-    info!("Metrics being pushed to Prometheus");
-
-    debug!(
-        "Public set: {:?}",
-        service.server.lock().unwrap().get_public_set()
-    );
-
     TonicServer::builder()
         .add_service(DiscoveryServer::new(service))
         .serve(addr)
diff --git a/src/lib.rs b/src/lib.rs
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs

-Original file line number
+Diff line change
 *.key
 *.crt
 *.log
++
 +.codegpt
Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,4 @@`
`1`	`1`	`fn main() -> Result<(), Box<dyn std::error::Error>> {`
`2`		`- tonic_build::configure()`
`3`		`- .type_attribute(".", "#[derive(Hash, Eq, PartialOrd, Ord)]")`
`4`		`- .compile(&["proto/rumi.proto"], &["proto"])?;`
	`2`	`+ tonic_build::compile_protos("proto/rumi.proto")?;`
`5`	`3`	`Ok(())`
`6`	`4`	`}`