one commit contains all impl

.gitignore

@@ -0,0 +1 @@
+/target

Cargo.toml

@@ -0,0 +1,7 @@
+[workspace]
+members = [
+  "libvanity",
+  "vanity",
+  "vanityd",
+  "vanity-worker",
+]

libvanity/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "libvanity"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+tonic = "0.8"
+prost = "0.11"
+
+[build-dependencies]
+tonic-build = "0.8"

libvanity/build.rs

@@ -0,0 +1,11 @@
+fn main() -> std::io::Result<()> {
+    let descriptor_path =
+        std::path::PathBuf::from(std::env::var("OUT_DIR").unwrap()).join("descriptor.bin");
+    tonic_build::configure()
+        .build_client(true)
+        .build_server(true)
+        .file_descriptor_set_path(descriptor_path)
+        .compile(&["vanity.proto", "worker.proto"], &["protos"])?;
+
+    Ok(())
+}

libvanity/protos/vanity.proto

@@ -0,0 +1,15 @@
+syntax = "proto3";
+
+package vanity;
+
+message VanityRequest {
+    string pattern = 1;
+}
+
+message VanityResponse {
+    string result = 1;
+}
+
+service VanityService {
+    rpc RequestVanity(VanityRequest) returns (VanityResponse);
+}

libvanity/protos/worker.proto

@@ -0,0 +1,26 @@
+syntax = "proto3";
+
+package worker;
+
+import "vanity.proto";
+
+message NewRequest {
+    uint64 id = 1;
+    string pattern = 2;
+}
+
+message Request {
+    oneof NewOrDone {
+        NewRequest new = 1;
+        uint64 done = 2;
+    }
+}
+
+message Result {
+    uint64 id = 1;
+    string result = 2;
+}
+
+service Worker {
+    rpc Poll(stream Result) returns (stream Request);
+}

libvanity/src/lib.rs

@@ -0,0 +1,9 @@
+pub mod vanity {
+    tonic::include_proto!("vanity");
+}
+
+pub mod worker {
+    tonic::include_proto!("worker");
+}
+
+pub const FILE_DESCRIPTOR_SET: &[u8] = tonic::include_file_descriptor_set!("descriptor");

vanity-worker/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "vanity-worker"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+clap = { version = "4", features = ["derive", "env", "cargo"] }
+
+libvanity = {path = "../libvanity"}
+
+x25519-dalek = "1.2.0"
+rand_core = {version = "0.5.1", features = ["getrandom"]}
+
+tonic = "0.8"
+tokio = {version = "1", features = ["full"]}
+tokio-stream = { version = "0.1.11", features = ["net"] }
+
+base64 = "0.13.0"
+regex = "1.6.0"
+tracing-subscriber = { version = "0.3.16", features = ["json"] }
+tracing = { version = "0.1.37", features = ["log"] }
+futures = "0.3.25"

vanity-worker/src/main.rs

@@ -0,0 +1,76 @@
+use std::collections::HashMap;
+
+use rand_core::OsRng;
+use tokio::sync::RwLock;
+use x25519_dalek::{PublicKey, StaticSecret};
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let client =
+        libvanity::worker::worker_client::WorkerClient::connect("grpc://127.0.0.1:8877").await?;
+
+    let (tx, rx) = tokio::sync::mpsc::channel(128);
+
+    let requests = std::sync::Arc::new(RwLock::new(HashMap::new()));
+
+    let mut another = client.clone();
+    let requests_writer = requests.clone();
+    let pull_thread = tokio::spawn(async move {
+        use libvanity::worker::request::NewOrDone::{Done, New};
+
+        let mut tasks = another
+            .poll(tokio_stream::wrappers::ReceiverStream::new(rx))
+            .await
+            .expect("msg")
+            .into_inner();
+        let requests = requests_writer;
+        while let Some(msg) = tasks.message().await.unwrap() {
+            if let Some(new_or_done) = msg.new_or_done {
+                match new_or_done {
+                    New(request) => {
+                        let mut writer = requests.write().await;
+                        writer.insert(request.id, request.pattern);
+                    }
+                    Done(response) => {
+                        let mut writer = requests.write().await;
+                        writer.remove(&response);
+                    }
+                }
+            }
+        }
+    });
+
+    let worker_threads = Vec::from_iter(1..10).into_iter().map(|id| {
+        let requests = requests.clone();
+        let tx = tx.clone();
+        let _id = id.clone();
+        tokio::spawn(async move {
+            loop {
+                let requests = requests.read().await.clone();
+                if requests.is_empty() {
+                    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+                    continue;
+                }
+
+                let private = StaticSecret::new(&mut OsRng);
+                let public = PublicKey::from(&private);
+
+                for (id, pattern) in requests.iter() {
+                    let matcher = regex::Regex::new(pattern).unwrap();
+                    let matcher = matcher.clone();
+                    let public_b64 = base64::encode(public.as_bytes());
+                    if matcher.is_match(&public_b64) {
+                        let mut result = libvanity::worker::Result::default();
+                        result.id = id.clone();
+                        result.result = base64::encode(private.to_bytes());
+                        tx.send(result).await.expect("no error");
+                    }
+                }
+            }
+        })
+    });
+
+    let _ = tokio::join!(futures::future::join_all(worker_threads), pull_thread);
+
+    Ok(())
+}

vanity/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "vanity"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+clap = { version = "4", features = ["derive", "env", "cargo"] }
+
+tonic = "0.8"
+tokio = {version = "1", features = ["full"]}
+tokio-stream = { version = "0.1.11", features = ["net"] }
+
+
+libvanity = {path = "../libvanity"}

vanity/src/main.rs

@@ -0,0 +1,30 @@
+use clap::Parser;
+use libvanity::vanity::VanityRequest;
+
+#[derive(clap::Parser, Debug)]
+#[clap(
+    about = "ask some worker to generate a private key where public key matches some regular expression"
+)]
+struct Args {
+    #[clap(short = 'c', long = "connect", help = "server address to connect to")]
+    endpoint: String,
+
+    #[clap(help = "pattern of public key")]
+    pattern: String,
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let args = Args::parse();
+
+    let mut client =
+        libvanity::vanity::vanity_service_client::VanityServiceClient::connect(args.endpoint)
+            .await?;
+
+    let mut request = VanityRequest::default();
+    request.pattern = args.pattern;
+    let response = client.request_vanity(request).await?.into_inner();
+    println!("{}", response.result);
+
+    Ok(())
+}

vanityd/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "vanityd"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+clap = { version = "4", features = ["derive", "env", "cargo"] }
+
+tonic = "0.8"
+tokio = {version = "1", features = ["full"]}
+tokio-stream = { version = "0.1.11", features = ["net"] }
+
+futures = "0.3.25"
+
+redis = { version = "0.22", features = ["aio", "tokio-comp"] }
+
+libvanity = {path = "../libvanity"}
+tracing-subscriber = { version = "0.3.16", features = ["json"] }
+tracing = { version = "0.1.37", features = ["log"] }

vanityd/src/consts.rs

@@ -0,0 +1 @@
+pub(crate) const REDIS_PUBSUB_PATTERN_RESULT: &str = "vanity:result:";

vanityd/src/main.rs

@@ -0,0 +1,31 @@
+pub(crate) mod consts;
+pub(crate) mod server;
+pub(crate) mod worker;
+
+use std::net::ToSocketAddrs;
+
+use libvanity::{
+    vanity::vanity_service_server::VanityServiceServer, worker::worker_server::WorkerServer,
+};
+
+pub use libvanity;
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    tracing_subscriber::FmtSubscriber::builder()
+        .with_max_level(tracing::Level::DEBUG)
+        .init();
+
+    let redis = redis::Client::open("redis://127.0.0.1:6379")?;
+
+    let vanity_server = server::VanityService::new(redis.clone());
+    let worker_server = worker::WorkerService::new(redis.clone());
+
+    tonic::transport::Server::builder()
+        .add_service(VanityServiceServer::new(vanity_server))
+        .add_service(WorkerServer::new(worker_server))
+        .serve("127.0.0.1:8877".to_socket_addrs()?.next().unwrap())
+        .await?;
+
+    Ok(())
+}

vanityd/src/server.rs

@@ -0,0 +1,56 @@
+use libvanity::vanity::{vanity_service_server as service_server, VanityRequest, VanityResponse};
+use tokio_stream::StreamExt;
+use tonic::async_trait;
+
+use crate::consts::*;
+
+use redis::AsyncCommands;
+
+pub(crate) struct VanityService {
+    counter: std::sync::atomic::AtomicU64,
+    redis: redis::Client,
+}
+
+impl VanityService {
+    pub fn new(redis: redis::Client) -> Self {
+        Self {
+            redis,
+            counter: std::sync::atomic::AtomicU64::new(0),
+        }
+    }
+}
+
+#[async_trait]
+impl service_server::VanityService for VanityService {
+    async fn request_vanity(
+        &self,
+        req: tonic::Request<VanityRequest>,
+    ) -> Result<tonic::Response<VanityResponse>, tonic::Status> {
+        let id = self
+            .counter
+            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+        let pattern = req.into_inner().pattern;
+
+        let mut redis = self.redis.get_async_connection().await.expect("msg");
+        redis
+            .publish::<&str, String, ()>("vanity:new", format!("{}|{}", id, pattern))
+            .await
+            .expect("should not fail");
+
+        let mut pubsub = redis.into_pubsub();
+        pubsub
+            .subscribe(format!("{}:{}", REDIS_PUBSUB_PATTERN_RESULT, id))
+            .await
+            .expect("msg");
+
+        let mut msgs = pubsub.on_message();
+
+        loop {
+            if let Some(msg) = msgs.next().await {
+                let mut response = VanityResponse::default();
+                response.result = String::from_utf8_lossy(msg.get_payload_bytes()).to_string();
+                return Ok(tonic::Response::new(response));
+            }
+        }
+    }
+}

vanityd/src/worker.rs

@@ -0,0 +1,146 @@
+use futures::StreamExt;
+use tonic::async_trait;
+
+use redis::AsyncCommands;
+
+use crate::consts::*;
+
+pub struct WorkerService {
+    redis: redis::Client,
+}
+
+impl WorkerService {
+    pub fn new(r: redis::Client) -> Self {
+        Self { redis: r }
+    }
+}
+
+#[async_trait]
+impl libvanity::worker::worker_server::Worker for WorkerService {
+    type PollStream =
+        tokio_stream::wrappers::ReceiverStream<tonic::Result<libvanity::worker::Request>>;
+    async fn poll(
+        &self,
+        result: tonic::Request<tonic::Streaming<libvanity::worker::Result>>,
+    ) -> Result<tonic::Response<Self::PollStream>, tonic::Status> {
+        tracing::trace!("entered polling method");
+        let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+        tracing::debug!("getting redis connection");
+        let redis_new_request = match self.redis.get_async_connection().await {
+            Ok(conn) => conn,
+            Err(err) => return Err(tonic::Status::unavailable(err.to_string())),
+        };
+
+        // handle new requests or results
+        let new_tx = tx.clone();
+        tokio::spawn(async move {
+            tracing::debug!("entering puller closure");
+            let tx = new_tx.clone();
+
+            let mut pubsub = redis_new_request.into_pubsub();
+            match pubsub.subscribe("vanity:new").await {
+                Err(err) => {
+                    tx.send(Err(tonic::Status::unavailable(err.to_string())))
+                        .await
+                        .ok();
+                }
+                _ => {}
+            }
+            match pubsub
+                .psubscribe(format!("{}*", REDIS_PUBSUB_PATTERN_RESULT))
+                .await
+            {
+                Err(err) => {
+                    tx.send(Err(tonic::Status::unavailable(err.to_string())))
+                        .await
+                        .ok();
+                }
+                _ => {}
+            }
+            let mut msgs = pubsub.on_message();
+            while let Some(msg) = msgs.next().await {
+                let mut request = libvanity::worker::Request::default();
+                if let Ok(pattern) = msg.get_channel::<String>() {
+                    // result,
+                    tracing::info!("pattern = {}", pattern);
+                    if let Some(key) = pattern.strip_prefix(&REDIS_PUBSUB_PATTERN_RESULT) {
+                        tracing::info!("key = {}", key);
+                        if let Ok(id) = u64::from_str_radix(key.trim_start_matches(":"), 10) {
+                            request.new_or_done =
+                                Some(libvanity::worker::request::NewOrDone::Done(id));
+                            tracing::info!("done -> {}", id.clone());
+                        }
+                    } else {
+                        // new pattern request
+                        let payload: String = msg.get_payload().unwrap();
+                        if let Some((idstr, pattern)) = payload.rsplit_once('|') {
+                            if let Ok(id) = u64::from_str_radix(idstr, 10) {
+                                let mut new_request = libvanity::worker::NewRequest::default();
+                                new_request.id = id.clone();
+                                new_request.pattern = String::from(pattern);
+                                request.new_or_done =
+                                    Some(libvanity::worker::request::NewOrDone::New(
+                                        // New
+                                        new_request,
+                                    ));
+                                tracing::info!("new -> {}", id.clone());
+                            }
+                        }
+                    }
+                }
+                if request.new_or_done == None {
+                    continue;
+                }
+
+                tracing::info!("sending {:?}", request.new_or_done);
+                if let Err(err) = tx.send(Ok(request)).await {
+                    tx.send(Err(tonic::Status::from_error(Box::new(err))))
+                        .await
+                        .ok();
+                };
+            }
+        });
+
+        let mut in_stream = result.into_inner();
+        let redis = self.redis.clone();
+        tokio::spawn(async move {
+            tracing::debug!("entering trans closure");
+            let redis = redis.clone();
+            while let Some(result) = in_stream.next().await {
+                match result {
+                    Ok(result) => {
+                        let mut redis = match redis.get_async_connection().await {
+                            Ok(redis) => redis,
+                            Err(err) => {
+                                tx.send(Err(tonic::Status::unavailable(err.to_string())))
+                                    .await
+                                    .ok();
+                                continue;
+                            }
+                        };
+                        tracing::info!("request done: {}", result.id);
+                        let _: u64 = redis
+                            .publish(
+                                format!("{}:{}", REDIS_PUBSUB_PATTERN_RESULT, result.id),
+                                result.result,
+                            )
+                            .await
+                            .unwrap();
+                    }
+                    Err(err) => {
+                        match tx.send(Err(err)).await {
+                            Ok(_) => (),
+                            Err(_err) => break, // response was droped
+                        }
+                    }
+                }
+            }
+        });
+
+        tracing::debug!("returned a response stream.( working in background )");
+        Ok(tonic::Response::new(
+            tokio_stream::wrappers::ReceiverStream::new(rx),
+        ))
+    }
+}