From 12fb704461b7baf897a6aca2e0c8fb213fee157c Mon Sep 17 00:00:00 2001 From: RyanSquared Date: Mon, 7 Aug 2023 12:37:05 -0500 Subject: [PATCH] mnemonic-hash-checker: move from postgres to bloom filter --- mnemonic-hash-checker/Cargo.lock | 299 ++--------------------------- mnemonic-hash-checker/Cargo.toml | 3 +- mnemonic-hash-checker/src/bloom.rs | 43 +++++ mnemonic-hash-checker/src/main.rs | 88 +++------ 4 files changed, 83 insertions(+), 350 deletions(-) create mode 100644 mnemonic-hash-checker/src/bloom.rs diff --git a/mnemonic-hash-checker/Cargo.lock b/mnemonic-hash-checker/Cargo.lock index 5b16512..d415c70 100644 --- a/mnemonic-hash-checker/Cargo.lock +++ b/mnemonic-hash-checker/Cargo.lock @@ -130,10 +130,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" [[package]] -name = "base64" -version = "0.21.2" +name = "bit-vec" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "604178f6c5c21f02dc555784810edfb88d34ac2c73b2eae109655649ee73ce3d" +checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" [[package]] name = "bitflags" @@ -157,10 +157,15 @@ dependencies = [ ] [[package]] -name = "byteorder" -version = "1.4.3" +name = "bloomfilter" +version = "1.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610" +checksum = "8129c0ab340c1b0caf6dbc587e814d04ba811e336dcf8fc268c04e047428ebb0" +dependencies = [ + "bit-vec", + "getrandom", + "siphasher", +] [[package]] name = "bytes" @@ -266,40 +271,6 @@ dependencies = [ "typenum", ] -[[package]] -name = "deadpool" -version = "0.9.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "421fe0f90f2ab22016f32a9881be5134fdd71c65298917084b0c7477cbc3856e" -dependencies = [ - "async-trait", - "deadpool-runtime", - "num_cpus", - "retain_mut", - "tokio", -] - -[[package]] -name = "deadpool-postgres" -version = "0.10.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "836a24a9d49deefe610b8b60c767a7412e9a931d79a89415cd2d2d71630ca8d7" -dependencies = [ - "deadpool", - "log", - "tokio", - "tokio-postgres", -] - -[[package]] -name = "deadpool-runtime" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eaa37046cc0f6c3cc6090fbdbf73ef0b8ef4cfcc37f6befc0020f63e8cf121e1" -dependencies = [ - "tokio", -] - [[package]] name = "digest" version = "0.10.7" @@ -308,7 +279,6 @@ checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ "block-buffer", "crypto-common", - "subtle", ] [[package]] @@ -342,12 +312,6 @@ dependencies = [ "once_cell", ] -[[package]] -name = "fallible-iterator" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" - [[package]] name = "fnv" version = "1.0.7" @@ -370,7 +334,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "955518d47e09b25bbebc7a18df10b81f0c766eaf4c4f1cccef2fca5f2a4fb5f2" dependencies = [ "futures-core", - "futures-sink", ] [[package]] @@ -379,23 +342,6 @@ version = "0.3.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4bca583b7e26f571124fe5b7561d49cb2868d79116cfa0eefce955557c6fee8c" -[[package]] -name = "futures-macro" -version = "0.3.28" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89ca545a94061b6365f2c7355b4b32bd20df3ff95f02da9329b34ccc3bd6ee72" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.28", -] - -[[package]] -name = "futures-sink" -version = "0.3.28" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f43be4fe21a13b9781a69afa4985b0f6ee0e1afab2c6f454a8cf30e2b2237b6e" - [[package]] name = "futures-task" version = "0.3.28" @@ -409,8 +355,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26b01e40b772d54cf6c6d721c1d1abd0647a0106a12ecaa1c186273392a69533" dependencies = [ "futures-core", - "futures-macro", - "futures-sink", "futures-task", "pin-project-lite", "pin-utils", @@ -450,7 +394,7 @@ version = "0.3.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f3e372db8e5c0d213e0cd0b9be18be2aca3d44cf2fe30a9d46a65581cd454584" dependencies = [ - "base64 0.13.1", + "base64", "bitflags 1.3.2", "bytes", "headers-core", @@ -481,15 +425,6 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "443144c8cdadd93ebf52ddb4056d257f5b52c04d3c804e657d19eb73fc33668b" -[[package]] -name = "hmac" -version = "0.12.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" -dependencies = [ - "digest", -] - [[package]] name = "http" version = "0.2.9" @@ -546,7 +481,7 @@ dependencies = [ "httpdate", "itoa", "pin-project-lite", - "socket2 0.4.9", + "socket2", "tokio", "tower-service", "tracing", @@ -625,15 +560,6 @@ version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed1202b2a6f884ae56f04cff409ab315c5ce26b5e58d7412e484f01fd52f52ef" -[[package]] -name = "md-5" -version = "0.10.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6365506850d44bff6e2fbcb5176cf63650e48bd45ef2fe2665ae1570e0f4b9ca" -dependencies = [ - "digest", -] - [[package]] name = "memchr" version = "2.5.0" @@ -671,12 +597,11 @@ name = "mnemonic-hash-checker" version = "0.1.0" dependencies = [ "axum", + "bloomfilter", "clap", "color-eyre", - "deadpool-postgres", "serde", "tokio", - "tokio-postgres", "tower-http", "tracing", "tracing-error", @@ -765,24 +690,6 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b2a4787296e9989611394c33f193f676704af1686e70b8f8033ab5ba9a35a94" -[[package]] -name = "phf" -version = "0.11.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ade2d8b8f33c7333b51bcf0428d37e217e9f32192ae4772156f65063b8ce03dc" -dependencies = [ - "phf_shared", -] - -[[package]] -name = "phf_shared" -version = "0.11.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90fcb95eef784c2ac79119d1dd819e162b5da872ce6f3c3abe1e8ca1c082f72b" -dependencies = [ - "siphasher", -] - [[package]] name = "pin-project" version = "1.1.3" @@ -815,41 +722,6 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" -[[package]] -name = "postgres-protocol" -version = "0.6.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78b7fa9f396f51dffd61546fd8573ee20592287996568e6175ceb0f8699ad75d" -dependencies = [ - "base64 0.21.2", - "byteorder", - "bytes", - "fallible-iterator", - "hmac", - "md-5", - "memchr", - "rand", - "sha2", - "stringprep", -] - -[[package]] -name = "postgres-types" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f028f05971fe20f512bcc679e2c10227e57809a3af86a7606304435bc8896cd6" -dependencies = [ - "bytes", - "fallible-iterator", - "postgres-protocol", -] - -[[package]] -name = "ppv-lite86" -version = "0.2.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" - [[package]] name = "proc-macro-error" version = "1.0.4" @@ -892,36 +764,6 @@ dependencies = [ "proc-macro2", ] -[[package]] -name = "rand" -version = "0.8.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" -dependencies = [ - "libc", - "rand_chacha", - "rand_core", -] - -[[package]] -name = "rand_chacha" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" -dependencies = [ - "ppv-lite86", - "rand_core", -] - -[[package]] -name = "rand_core" -version = "0.6.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" -dependencies = [ - "getrandom", -] - [[package]] name = "redox_syscall" version = "0.3.5" @@ -975,12 +817,6 @@ version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5ea92a5b6195c6ef2a0295ea818b312502c6fc94dde986c5553242e18fd4ce2" -[[package]] -name = "retain_mut" -version = "0.1.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4389f1d5789befaf6029ebd9f7dac4af7f7e3d61b69d4f30e2ac02b57e7712b0" - [[package]] name = "rustc-demangle" version = "0.1.23" @@ -1082,17 +918,6 @@ dependencies = [ "digest", ] -[[package]] -name = "sha2" -version = "0.10.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "479fb9d862239e610720565ca91403019f2f00410f1864c5aa7479b950a76ed8" -dependencies = [ - "cfg-if", - "cpufeatures", - "digest", -] - [[package]] name = "sharded-slab" version = "0.1.4" @@ -1142,38 +967,12 @@ dependencies = [ "winapi", ] -[[package]] -name = "socket2" -version = "0.5.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2538b18701741680e0322a2302176d3253a35388e2e62f172f64f4f16605f877" -dependencies = [ - "libc", - "windows-sys", -] - -[[package]] -name = "stringprep" -version = "0.1.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db3737bde7edce97102e0e2b15365bf7a20bfdb5f60f4f9e8d7004258a51a8da" -dependencies = [ - "unicode-bidi", - "unicode-normalization", -] - [[package]] name = "strsim" version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" -[[package]] -name = "subtle" -version = "2.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81cdd64d312baedb58e21336b31bc043b77e01cc99033ce76ef539f78e965ebc" - [[package]] name = "syn" version = "1.0.109" @@ -1221,21 +1020,6 @@ dependencies = [ "once_cell", ] -[[package]] -name = "tinyvec" -version = "1.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87cc5ceb3875bb20c2890005a4e226a4651264a5c75edb2421b52861a0a0cb50" -dependencies = [ - "tinyvec_macros", -] - -[[package]] -name = "tinyvec_macros" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" - [[package]] name = "tokio" version = "1.29.1" @@ -1251,7 +1035,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2 0.4.9", + "socket2", "tokio-macros", "tracing", "windows-sys", @@ -1268,44 +1052,6 @@ dependencies = [ "syn 2.0.28", ] -[[package]] -name = "tokio-postgres" -version = "0.7.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e89f6234aa8fd43779746012fcf53603cdb91fdd8399aa0de868c2d56b6dde1" -dependencies = [ - "async-trait", - "byteorder", - "bytes", - "fallible-iterator", - "futures-channel", - "futures-util", - "log", - "parking_lot", - "percent-encoding", - "phf", - "pin-project-lite", - "postgres-protocol", - "postgres-types", - "socket2 0.5.3", - "tokio", - "tokio-util", -] - -[[package]] -name = "tokio-util" -version = "0.7.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "806fe8c2c87eccc8b3267cbae29ed3ab2d0bd37fca70ab622e46aaa9375ddb7d" -dependencies = [ - "bytes", - "futures-core", - "futures-sink", - "pin-project-lite", - "tokio", - "tracing", -] - [[package]] name = "tower" version = "0.4.13" @@ -1451,27 +1197,12 @@ version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "497961ef93d974e23eb6f433eb5fe1b7930b659f06d12dec6fc44a8f554c0bba" -[[package]] -name = "unicode-bidi" -version = "0.3.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92888ba5573ff080736b3648696b70cafad7d250551175acbaa4e0385b3e1460" - [[package]] name = "unicode-ident" version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "301abaae475aa91687eb82514b328ab47a211a533026cb25fc3e519b86adfc3c" -[[package]] -name = "unicode-normalization" -version = "0.1.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c5713f0fc4b5db668a2ac63cdb7bb4469d8c9fed047b1d0292cc7b0ce2ba921" -dependencies = [ - "tinyvec", -] - [[package]] name = "valuable" version = "0.1.0" diff --git a/mnemonic-hash-checker/Cargo.toml b/mnemonic-hash-checker/Cargo.toml index 69d1d52..bab549f 100644 --- a/mnemonic-hash-checker/Cargo.toml +++ b/mnemonic-hash-checker/Cargo.toml @@ -7,12 +7,11 @@ edition = "2021" [dependencies] axum = { version = "0.6.20", features = ["headers", "macros", "tracing"] } +bloomfilter = "1.0.9" clap = { version = "=4.0.32", features = ["derive", "env"] } color-eyre = "0.6.2" -deadpool-postgres = "0.10.5" serde = { version = "1.0.181", features = ["serde_derive"] } tokio = { version = "1.29.1", features = ["full", "fs", "tracing"] } -tokio-postgres = "0.7.8" tower-http = { version = "0.4.3", features = ["trace", "catch-panic"] } tracing = "0.1.37" tracing-error = "0.2.0" diff --git a/mnemonic-hash-checker/src/bloom.rs b/mnemonic-hash-checker/src/bloom.rs new file mode 100644 index 0000000..5f20b23 --- /dev/null +++ b/mnemonic-hash-checker/src/bloom.rs @@ -0,0 +1,43 @@ +use std::{fs::File, io::{BufReader, Read}, path::Path}; + +use color_eyre::eyre::Result; +use bloomfilter::Bloom; + +pub fn load(file: &Path) -> Result> { + let file = File::open(file)?; + let length = file.metadata().unwrap().len(); + + let mut buf = BufReader::new(file); + + let mut num_bits: [u8; 8] = [0; 8]; + buf.read_exact(&mut num_bits)?; + + let mut num_hash_fun: [u8; 4] = [0; 4]; + buf.read_exact(&mut num_hash_fun)?; + + let mut sk00: [u8; 8] = [0; 8]; + buf.read_exact(&mut sk00)?; + let mut sk01: [u8; 8] = [0; 8]; + buf.read_exact(&mut sk01)?; + let mut sk10: [u8; 8] = [0; 8]; + buf.read_exact(&mut sk10)?; + let mut sk11: [u8; 8] = [0; 8]; + buf.read_exact(&mut sk11)?; + + let number_of_bits: u64 = u64::from_be_bytes(num_bits); + let number_of_hash_functions: u32 = u32::from_be_bytes(num_hash_fun); + let sip_keys: [(u64, u64); 2] = [ + (u64::from_be_bytes(sk00), (u64::from_be_bytes(sk01))), + (u64::from_be_bytes(sk10), (u64::from_be_bytes(sk11))), + ]; + + let mut bitmap = vec![0; (length - 8 - 4 - 32) as usize]; + buf.read_exact(&mut bitmap)?; + + Ok(Bloom::from_existing( + &bitmap, + number_of_bits, + number_of_hash_functions, + sip_keys, + )) +} diff --git a/mnemonic-hash-checker/src/main.rs b/mnemonic-hash-checker/src/main.rs index caff246..b369e59 100644 --- a/mnemonic-hash-checker/src/main.rs +++ b/mnemonic-hash-checker/src/main.rs @@ -1,4 +1,5 @@ use std::net::SocketAddr; +use std::path::PathBuf; use std::sync::Arc; use axum::{ @@ -9,13 +10,12 @@ use axum::{ Json, Router, }; +use bloomfilter::Bloom; + use color_eyre::eyre::Result; use clap::Parser; -use deadpool_postgres::{Config, ManagerConfig, Pool, RecyclingMethod, Runtime}; -use tokio_postgres::NoTls; - use serde::Deserialize; use tower_http::catch_panic::CatchPanicLayer; @@ -25,39 +25,21 @@ use tracing_subscriber::filter::{EnvFilter, LevelFilter}; use tracing_subscriber::fmt::format::FmtSpan; use tracing_subscriber::prelude::*; -use tracing::{debug, info}; +use tracing::info; + +mod bloom; #[derive(Parser)] struct CliConfig { #[clap(long, env)] - bind_address: Option, - - #[clap(long, env, default_value = "postgres")] - postgres_dbname: String, - - #[clap(long, env, default_value = "localhost")] - postgres_host: String, + bloom_filter_path: PathBuf, #[clap(long, env)] - postgres_password: String, - - #[clap(long, env, default_value = "5432")] - postgres_port: u16, - - #[clap(long, env, default_value = "postgres")] - postgres_user: String, - - #[clap(long, env, default_value = "mnemonics")] - database_table: String, - - #[clap(long, env, default_value = "hash")] - database_column: String, + bind_address: Option, } struct AppState { - pool: Pool, - database_table: String, - database_column: String, + bloom_filter: Bloom } fn setup_registry() { @@ -79,19 +61,9 @@ async fn index() -> impl IntoResponse { } #[tracing::instrument(skip(state))] -async fn check_hash(hash: &str, state: &AppState) -> Result { - let client = state.pool.get().await?; - let column = &state.database_column; - let table = &state.database_table; - let formatted_query = format!("SELECT {column} FROM {table} WHERE hash = $1"); - let query = client.prepare_cached(formatted_query.as_str()).await?; - let rows = client.query(&query, &[&hash]).await?; - if let Some(row) = rows.get(0) { - let retrieved_hash: String = row.try_get(0)?; - Ok(hash == retrieved_hash) - } else { - Ok(false) - } +async fn check_hash(hash: &str, state: &AppState) -> bool { + let bloom_filter = &state.bloom_filter; + bloom_filter.check(&hash.to_string()) } #[derive(Deserialize)] @@ -104,17 +76,13 @@ struct CheckQuery { async fn check_hash_slug( query: Query, State(state): State>, -) -> (StatusCode, Json>) { +) -> (StatusCode, Json) { let result = check_hash(&query.sha256, &state).await; - let status_code = match result.as_ref() { - Ok(true) => StatusCode::OK, - Ok(false) => StatusCode::NOT_FOUND, - Err(e) => { - debug!(%e, "Error while performing lookup"); - StatusCode::INTERNAL_SERVER_ERROR - } + let status_code = match state.bloom_filter.check(&query.sha256.to_string()) { + true => StatusCode::OK, + false => StatusCode::NOT_FOUND, }; - (status_code, Json(result.ok())) + (status_code, Json(result)) } #[tokio::main] @@ -123,28 +91,20 @@ async fn main() -> Result<()> { color_eyre::install()?; let cli_config = CliConfig::parse(); - let mut config = Config::new(); - config.dbname = Some(cli_config.postgres_dbname); - config.host = Some(cli_config.postgres_host); - config.password = Some(cli_config.postgres_password); - config.port = Some(cli_config.postgres_port); - config.user = Some(cli_config.postgres_user); - config.manager = Some(ManagerConfig { - recycling_method: RecyclingMethod::Fast, - }); - let pool = config.create_pool(Some(Runtime::Tokio1), NoTls)?; let addr = cli_config .bind_address .unwrap_or_else(|| (std::net::Ipv4Addr::new(127, 0, 0, 1), 8000).into()); + info!("loading bloom filter"); + + let state = AppState { + bloom_filter: bloom::load(cli_config.bloom_filter_path.as_ref())? + }; + let app = Router::new() .route("/", get(index)) .route("/check", get(check_hash_slug)) - .with_state(Arc::new(AppState { - pool, - database_table: cli_config.database_table, - database_column: cli_config.database_column, - })) + .with_state(Arc::new(state)) .layer(CatchPanicLayer::new()) .layer(TraceLayer::new_for_http());