started rewriting datastore and gRPC

rewrite/Cargo.toml

@@ -5,6 +5,16 @@ edition = "2021"
 
 [dependencies]
 actix-web = "4.9.0"
+async-stream = "0.3.5"
+dashmap = "6.1.0"
+prost = "0.13.2"
+prost-types = "0.13.2"
 rand = "0.8.5"
 serde = { version = "1.0.210", features = ["derive"] }
+solana-sdk = "2.0.10"
 tokio = { version = "1.40.0", features = ["macros", "rt-multi-thread", "fs"] }
+tokio-stream = { version = "0.1.16", features = ["sync"] }
+tonic = "0.12.1"
+
+[build-dependencies]
+tonic-build = "0.12.1"

rewrite/build.rs

@@ -0,0 +1,9 @@
+fn main() {
+    tonic_build::configure()
+        .build_server(true)
+        .compile(
+            &["proto/challenge.proto"],
+            &["proto"],
+        )
+        .unwrap_or_else(|e| panic!("Failed to compile protos {:?}", e));
+}

rewrite/proto/challenge.proto

@@ -0,0 +1,25 @@
+syntax = "proto3";
+package challenge;
+
+import "google/protobuf/timestamp.proto";
+
+message NodeUpdate {
+  string ip = 1;
+  google.protobuf.Timestamp started_at = 2;
+  google.protobuf.Timestamp keepalive = 3;
+  uint64 total_mints = 4;
+  uint64 ratls_conns = 5;
+  uint64 ratls_attacks = 6;
+  bool public = 7;
+}
+
+message Keypair {
+  string keypair = 1;
+}
+
+message Empty {}
+
+service Update {
+  rpc GetKeypair (Empty) returns (Keypair);
+  rpc GetUpdates (stream NodeUpdate) returns (stream NodeUpdate);
+}

rewrite/proto/google/protobuf/timestamp.proto

@@ -0,0 +1,147 @@
+// Protocol Buffers - Google's data interchange format
+// Copyright 2008 Google Inc.  All rights reserved.
+// https://developers.google.com/protocol-buffers/
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are
+// met:
+//
+//     * Redistributions of source code must retain the above copyright
+// notice, this list of conditions and the following disclaimer.
+//     * Redistributions in binary form must reproduce the above
+// copyright notice, this list of conditions and the following disclaimer
+// in the documentation and/or other materials provided with the
+// distribution.
+//     * Neither the name of Google Inc. nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+syntax = "proto3";
+
+package google.protobuf;
+
+option csharp_namespace = "Google.Protobuf.WellKnownTypes";
+option cc_enable_arenas = true;
+option go_package = "google.golang.org/protobuf/types/known/timestamppb";
+option java_package = "com.google.protobuf";
+option java_outer_classname = "TimestampProto";
+option java_multiple_files = true;
+option objc_class_prefix = "GPB";
+
+// A Timestamp represents a point in time independent of any time zone or local
+// calendar, encoded as a count of seconds and fractions of seconds at
+// nanosecond resolution. The count is relative to an epoch at UTC midnight on
+// January 1, 1970, in the proleptic Gregorian calendar which extends the
+// Gregorian calendar backwards to year one.
+//
+// All minutes are 60 seconds long. Leap seconds are "smeared" so that no leap
+// second table is needed for interpretation, using a [24-hour linear
+// smear](https://developers.google.com/time/smear).
+//
+// The range is from 0001-01-01T00:00:00Z to 9999-12-31T23:59:59.999999999Z. By
+// restricting to that range, we ensure that we can convert to and from [RFC
+// 3339](https://www.ietf.org/rfc/rfc3339.txt) date strings.
+//
+// # Examples
+//
+// Example 1: Compute Timestamp from POSIX `time()`.
+//
+//     Timestamp timestamp;
+//     timestamp.set_seconds(time(NULL));
+//     timestamp.set_nanos(0);
+//
+// Example 2: Compute Timestamp from POSIX `gettimeofday()`.
+//
+//     struct timeval tv;
+//     gettimeofday(&tv, NULL);
+//
+//     Timestamp timestamp;
+//     timestamp.set_seconds(tv.tv_sec);
+//     timestamp.set_nanos(tv.tv_usec * 1000);
+//
+// Example 3: Compute Timestamp from Win32 `GetSystemTimeAsFileTime()`.
+//
+//     FILETIME ft;
+//     GetSystemTimeAsFileTime(&ft);
+//     UINT64 ticks = (((UINT64)ft.dwHighDateTime) << 32) | ft.dwLowDateTime;
+//
+//     // A Windows tick is 100 nanoseconds. Windows epoch 1601-01-01T00:00:00Z
+//     // is 11644473600 seconds before Unix epoch 1970-01-01T00:00:00Z.
+//     Timestamp timestamp;
+//     timestamp.set_seconds((INT64) ((ticks / 10000000) - 11644473600LL));
+//     timestamp.set_nanos((INT32) ((ticks % 10000000) * 100));
+//
+// Example 4: Compute Timestamp from Java `System.currentTimeMillis()`.
+//
+//     long millis = System.currentTimeMillis();
+//
+//     Timestamp timestamp = Timestamp.newBuilder().setSeconds(millis / 1000)
+//         .setNanos((int) ((millis % 1000) * 1000000)).build();
+//
+//
+// Example 5: Compute Timestamp from Java `Instant.now()`.
+//
+//     Instant now = Instant.now();
+//
+//     Timestamp timestamp =
+//         Timestamp.newBuilder().setSeconds(now.getEpochSecond())
+//             .setNanos(now.getNano()).build();
+//
+//
+// Example 6: Compute Timestamp from current time in Python.
+//
+//     timestamp = Timestamp()
+//     timestamp.GetCurrentTime()
+//
+// # JSON Mapping
+//
+// In JSON format, the Timestamp type is encoded as a string in the
+// [RFC 3339](https://www.ietf.org/rfc/rfc3339.txt) format. That is, the
+// format is "{year}-{month}-{day}T{hour}:{min}:{sec}[.{frac_sec}]Z"
+// where {year} is always expressed using four digits while {month}, {day},
+// {hour}, {min}, and {sec} are zero-padded to two digits each. The fractional
+// seconds, which can go up to 9 digits (i.e. up to 1 nanosecond resolution),
+// are optional. The "Z" suffix indicates the timezone ("UTC"); the timezone
+// is required. A proto3 JSON serializer should always use UTC (as indicated by
+// "Z") when printing the Timestamp type and a proto3 JSON parser should be
+// able to accept both UTC and other timezones (as indicated by an offset).
+//
+// For example, "2017-01-15T01:30:15.01Z" encodes 15.01 seconds past
+// 01:30 UTC on January 15, 2017.
+//
+// In JavaScript, one can convert a Date object to this format using the
+// standard
+// [toISOString()](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Date/toISOString)
+// method. In Python, a standard `datetime.datetime` object can be converted
+// to this format using
+// [`strftime`](https://docs.python.org/2/library/time.html#time.strftime) with
+// the time format spec '%Y-%m-%dT%H:%M:%S.%fZ'. Likewise, in Java, one can use
+// the Joda Time's [`ISODateTimeFormat.dateTime()`](
+// http://www.joda.org/joda-time/apidocs/org/joda/time/format/ISODateTimeFormat.html#dateTime%2D%2D
+// ) to obtain a formatter capable of generating timestamps in this format.
+//
+//
+message Timestamp {
+  // Represents seconds of UTC time since Unix epoch
+  // 1970-01-01T00:00:00Z. Must be from 0001-01-01T00:00:00Z to
+  // 9999-12-31T23:59:59Z inclusive.
+  int64 seconds = 1;
+
+  // Non-negative fractions of a second at nanosecond resolution. Negative
+  // second values with fractions must still have non-negative nanos values
+  // that count forward in time. Must be from 0 to 999,999,999
+  // inclusive.
+  int32 nanos = 2;
+}

rewrite/src/datastore.rs

@@ -0,0 +1,105 @@
+#![allow(dead_code)]
+use crate::grpc::challenge::NodeUpdate;
+use dashmap::DashMap;
+use dashmap::DashSet;
+use solana_sdk::signature::keypair::Keypair;
+use std::time::SystemTime;
+
+type IP = String;
+
+#[derive(Clone, PartialEq, Debug)]
+pub struct NodeInfo {
+    pub started_at: SystemTime,
+    pub keepalive: SystemTime,
+    pub total_mints: u64,
+    pub ratls_conns: u64,
+    pub ratls_attacks: u64,
+    pub public: bool,
+}
+
+/// Keypair must already be known when creating a Store
+/// This means the first node of the network creates the key
+/// Second node will grab the key from the first node
+pub struct Store {
+    key: Keypair,
+    nodes: DashMap<IP, NodeInfo>,
+    conns: DashSet<IP>,
+    // TODO: write persistence
+    // persistence: FileManager,
+}
+
+impl Store {
+    pub fn init(key: Keypair) -> Self {
+        Self {
+            key,
+            nodes: DashMap::new(),
+            conns: DashSet::new(),
+        }
+    }
+
+    pub fn get_keypair_base58(&self) -> String {
+        self.key.to_base58_string()
+    }
+
+    pub async fn add_conn(&self, ip: &str) {
+        self.conns.insert(ip.to_string());
+    }
+
+    pub async fn delete_conn(&self, ip: &str) {
+        self.conns.remove(ip);
+    }
+
+    pub async fn get_localhost(&self) -> NodeUpdate {
+        // TODO trigger reset_localhost_keys on error instead of expects
+        let node = self.nodes.get("localhost").expect("no localhost node");
+        NodeUpdate {
+            ip: "localhost".to_string(),
+            started_at: Some(prost_types::Timestamp::from(node.started_at)),
+            keepalive: Some(prost_types::Timestamp::from(SystemTime::now())),
+            total_mints: node.total_mints,
+            ratls_conns: node.ratls_conns,
+            ratls_attacks: node.ratls_attacks,
+            public: false,
+        }
+    }
+
+    /// This returns true if NodeInfo got modified.
+    ///
+    /// On a side note, there are two types of people in this world:
+    ///     1. Those that can extrapolate... WAT?
+    pub async fn process_node_update(&self, _node: NodeUpdate) -> bool {
+        // TODO: write this function
+        true
+    }
+
+    pub async fn get_full_node_list(&self) -> Vec<NodeUpdate> {
+        self.nodes
+            .iter()
+            .map(|node| NodeUpdate {
+                ip: node.key().to_string(),
+                started_at: Some(prost_types::Timestamp::from(node.value().started_at)),
+                keepalive: Some(prost_types::Timestamp::from(node.value().keepalive)),
+                total_mints: node.value().total_mints,
+                ratls_conns: node.value().ratls_conns,
+                ratls_attacks: node.value().ratls_attacks,
+                public: node.value().public,
+            })
+            .collect()
+    }
+
+    // returns a random node that does not have an active connection
+    pub async fn get_random_node(&self) -> Option<String> {
+        use rand::{rngs::OsRng, RngCore};
+        let len = self.nodes.len();
+        if len == 0 {
+            return None;
+        }
+        let skip = OsRng.next_u64().try_into().unwrap_or(0) % len;
+        self.nodes
+            .iter()
+            .map(|n| n.key().clone())
+            .cycle()
+            .skip(skip)
+            .find(|k| !self.conns.contains(k))
+    }
+}

rewrite/src/grpc/client.rs

@@ -0,0 +1,90 @@
+#![allow(dead_code)]
+use super::challenge::NodeUpdate;
+use crate::datastore::Store;
+use crate::grpc::challenge::update_client::UpdateClient;
+use crate::grpc::challenge::Empty;
+use solana_sdk::signature::keypair::Keypair;
+use std::sync::Arc;
+use tokio::sync::broadcast::Sender;
+use tokio::time::{sleep, Duration};
+use tokio_stream::wrappers::BroadcastStream;
+use tokio_stream::StreamExt;
+
+#[derive(Clone)]
+pub struct ConnManager {
+    ds: Arc<Store>,
+    tx: Sender<NodeUpdate>,
+}
+
+impl ConnManager {
+    pub fn init(ds: Arc<Store>, tx: Sender<NodeUpdate>) -> Self {
+        Self { ds, tx }
+    }
+
+    pub async fn start_with_node(self, node_ip: String) {
+        self.connect_wrapper(node_ip).await;
+    }
+
+    pub async fn start(self) {
+        loop {
+            if let Some(node) = self.ds.get_random_node().await {
+                if node != "localhost" {
+                    self.connect_wrapper(node.clone()).await;
+                }
+            }
+            sleep(Duration::from_secs(3)).await;
+        }
+    }
+
+    async fn connect_wrapper(&self, node_ip: String) {
+        let ds = self.ds.clone();
+        ds.add_conn(&node_ip).await;
+        if let Err(e) = self.connect(node_ip.clone()).await {
+            println!("Client connection for {node_ip} failed: {e:?}");
+        }
+        ds.delete_conn(&node_ip).await;
+    }
+
+    async fn connect(&self, node_ip: String) -> Result<(), Box<dyn std::error::Error>> {
+        println!("Connecting to {node_ip}...");
+        let mut client = UpdateClient::connect(format!("http://{node_ip}:31373")).await?;
+
+        let rx = self.tx.subscribe();
+        let rx_stream = BroadcastStream::new(rx).filter_map(|n| n.ok());
+        let response = client.get_updates(rx_stream).await?;
+        let mut resp_stream = response.into_inner();
+
+        let _ = self.tx.send(self.ds.get_localhost().await);
+
+        while let Some(mut update) = resp_stream.message().await? {
+            // "localhost" IPs need to be changed to the real IP of the counterpart
+            if update.ip == "localhost" {
+                update.ip = node_ip.clone();
+                // since we are connecting TO this server, we have a guarantee that this
+                // server is not behind NAT, so we can set it public
+                update.public = true;
+            }
+
+            // update the entire network in case the information is new
+            if self.ds.process_node_update(update.clone()).await {
+                if let Err(_) = self.tx.send(update.clone()) {
+                    println!("tokio broadcast receivers had an issue consuming the channel");
+                }
+            };
+        }
+
+        Ok(())
+    }
+}
+
+pub async fn key_grabber(node_ip: String) -> Result<Keypair, Box<dyn std::error::Error>> {
+    let mut client = UpdateClient::connect(format!("http://{node_ip}:31373")).await?;
+    let response = client.get_keypair(tonic::Request::new(Empty {})).await?;
+    let response = &response.into_inner().keypair;
+    let keypair =
+        match std::panic::catch_unwind(|| Keypair::from_base58_string(&response)) {
+            Ok(k) => k,
+            Err(_) => return Err("Could not parse key".into()),
+        };
+    Ok(keypair)
+}

rewrite/src/grpc/mod.rs

@@ -0,0 +1,6 @@
+pub mod server;
+pub mod client;
+
+pub mod challenge {
+    tonic::include_proto!("challenge");
+}

rewrite/src/grpc/server.rs

@@ -0,0 +1,101 @@
+#![allow(dead_code)]
+
+use super::challenge::update_server::UpdateServer;
+use super::challenge::{Empty, Keypair, NodeUpdate};
+use crate::datastore::Store;
+use crate::grpc::challenge::update_server::Update;
+use std::pin::Pin;
+use std::sync::Arc;
+use tokio::sync::broadcast::Sender;
+use tokio_stream::{Stream, StreamExt};
+use tonic::{transport::Server, Request, Response, Status, Streaming};
+
+pub struct MyServer {
+    ds: Arc<Store>,
+    tx: Sender<NodeUpdate>,
+}
+
+impl MyServer {
+    pub fn init(ds: Arc<Store>, tx: Sender<NodeUpdate>) -> Self {
+        Self { ds, tx }
+    }
+
+    pub async fn start(self) {
+        let addr = "0.0.0.0:31373".parse().unwrap();
+        if let Err(e) = Server::builder()
+            .add_service(UpdateServer::new(self))
+            .serve(addr)
+            .await
+        {
+            println!("gRPC server failed: {e:?}");
+        };
+    }
+}
+
+#[tonic::async_trait]
+impl Update for MyServer {
+    type GetUpdatesStream = Pin<Box<dyn Stream<Item = Result<NodeUpdate, Status>> + Send>>;
+
+    async fn get_updates(
+        &self,
+        req: Request<Streaming<NodeUpdate>>,
+    ) -> Result<Response<Self::GetUpdatesStream>, Status> {
+        let remote_ip = req.remote_addr().unwrap().ip().to_string();
+        let tx = self.tx.clone();
+        let mut rx = self.tx.subscribe();
+        let mut inbound = req.into_inner();
+        let ds = self.ds.clone();
+
+        let stream = async_stream::stream! {
+            let full_update_list = ds.get_full_node_list().await;
+            for update in full_update_list {
+                yield Ok(update);
+            }
+
+            loop {
+                tokio::select! {
+                    Some(msg) = inbound.next() => {
+                        match msg {
+                            Ok(mut update) => {
+                                if update.ip == "localhost" {
+                                    update.ip = remote_ip.clone();
+                                    // note that we don't set this node online,
+                                    // as it can be behind NAT
+                                }
+                                if update.ip != "127.0.0.1" && ds.process_node_update(update.clone()).await {
+                                    if let Err(_) = tx.send(update.clone()) {
+                                        println!("tokio broadcast receivers had an issue consuming the channel");
+                                    }
+                                };
+                            }
+                            Err(e) => {
+                                yield Err(Status::internal(format!("Error receiving client stream: {}", e)));
+                                break;
+                            }
+                        }
+                    }
+                    Ok(update) = rx.recv() => {
+                        yield Ok(update);
+                        // disconnect client if too many connections are active
+                        if tx.receiver_count() > 9 {
+                            yield Err(Status::internal("Already have too many clients. Connect to another server."));
+                            return;
+                        }
+                    }
+
+                }
+            }
+        };
+
+        Ok(Response::new(Box::pin(stream) as Self::GetUpdatesStream))
+    }
+
+    async fn get_keypair(&self, request: Request<Empty>) -> Result<Response<Keypair>, Status> {
+        println!("Got a request from {:?}", request.remote_addr());
+
+        let reply = Keypair {
+            keypair: self.ds.get_keypair_base58(),
+        };
+        Ok(Response::new(reply))
+    }
+}

rewrite/src/http_server.rs

@@ -1,10 +1,17 @@
+#![allow(dead_code)]
 use actix_web::{
-    error::ResponseError, get, http::StatusCode, post, web, App, HttpResponse, HttpServer,
+    // http::StatusCode, error::ResponseError, Result,
-    Responder, Result,
+    get,
+    post,
+    web,
+    App,
+    HttpResponse,
+    HttpServer,
+    Responder,
 };
 use rand::Rng;
 use serde::{Deserialize, Serialize};
-use std::{fmt, sync::Arc};
+// use std::{fmt, sync::Arc};
 
 const HOMEPAGE: &str = r#"Welcome, beloved hacker!
 
@@ -56,7 +63,7 @@ struct MintReq {
 }
 
 #[post("/mint")]
-async fn mint(req: web::Json<MintReq>) -> impl Responder {
+async fn mint(_: web::Json<MintReq>) -> impl Responder {
     HttpResponse::Ok().json({})
 }
 

rewrite/src/main.rs

@@ -1,3 +1,5 @@
+mod grpc;
+mod datastore;
 mod http_server;
 
 use tokio::task::JoinSet;