Wallet Blacklist System

# Download latest blacklist
curl -o blacklist.jsonl https://allenhark.com/blacklist.jsonl

// file: src/blacklist.rs
// Blacklist with hash-based lookups (matches blacklist.ts)

use anyhow::{Context, Result};
use dashmap::DashSet;
use fs2::FileExt;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::Duration;
use tokio::fs::{self, File, OpenOptions};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};

#[derive(Debug, Serialize, Deserialize)]
pub struct Entry {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub addr: Option<String>,
    pub ts: u64,
    pub hash: String,
}

// 🚀 OPTIMIZATION: Pre-computed hex lookup table for fast encoding
const HEX_LOOKUP: [char; 16] = [
    '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f',
];

/// 🚀 OPTIMIZATION: Fast non-cryptographic hash
/// Uses optimized djb2 hash - minimal operations, maximum speed
/// Matches the TypeScript implementation exactly
pub fn fast_hash_8_bytes(input: &str) -> String {
    let bytes = input.as_bytes();
    let mut h1: u32 = 5381; // djb2 initial value
    let mut h2: u32 = 5381;

    // Process string with minimal operations
    for (i, &byte) in bytes.iter().enumerate() {
        let c = byte as u32;
        h1 = h1.wrapping_mul(33) ^ c; // h1 * 33 XOR c
        h2 = h2.wrapping_mul(33) ^ (c + i as u32); // Mix with position for better distribution
    }

    // Hex encoding: direct string building with lookup table
    // Build hex string directly (single allocation)
    let mut result = String::with_capacity(16);
    result.push(HEX_LOOKUP[((h2 >> 28) & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h2 >> 24) & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h2 >> 20) & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h2 >> 16) & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h2 >> 12) & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h2 >> 8) & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h2 >> 4) & 0xF) as usize]);
    result.push(HEX_LOOKUP[(h2 & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h1 >> 28) & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h1 >> 24) & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h1 >> 20) & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h1 >> 16) & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h1 >> 12) & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h1 >> 8) & 0xF) as usize]);
    result.push(HEX_LOOKUP[((h1 >> 4) & 0xF) as usize]);
    result.push(HEX_LOOKUP[(h1 & 0xF) as usize]);

    result
}

pub struct Blacklist {
    set: DashSet<String>, // stores short-hash hex (16 chars = 8 bytes) - lock-free reads!
    filepath: PathBuf,
    lockfile_path: PathBuf, // Path to lock file for cross-process synchronization
    snapshot_in_progress: bool,
    // 🚀 OPTIMIZATION: Cache hash computations (LRU-style, most recent queries)
    hash_cache: HashMap<String, String>,
    cache_size: usize,
    // Mapping from hash -> address to preserve addresses during snapshot
    hash_to_addr: HashMap<String, String>,
    // Mapping from hash -> timestamp to preserve timestamps during snapshot
    hash_to_ts: HashMap<String, u64>,
}

impl Blacklist {
    pub fn new(filepath: impl AsRef<Path>) -> Self {
        let filepath = filepath.as_ref().to_path_buf();
        let lockfile_path = filepath.with_extension("lock");
        Self {
            set: DashSet::new(),
            filepath,
            lockfile_path,
            snapshot_in_progress: false,
            hash_cache: HashMap::new(),
            cache_size: 1000, // Keep last 1000 hashes in cache (matches TS)
            hash_to_addr: HashMap::new(),
            hash_to_ts: HashMap::new(),
        }
    }

    /// Acquire a shared (read) lock on the blacklist file
    /// Returns a guard that releases the lock when dropped
    async fn acquire_read_lock(&self) -> Result<std::fs::File> {
        // Ensure lock file exists
        if let Some(parent) = self.lockfile_path.parent() {
            fs::create_dir_all(parent).await.ok();
        }
        
        let lockfile = std::fs::OpenOptions::new()
            .create(true)
            .read(true)
            .write(true)
            .open(&self.lockfile_path)
            .context("Failed to open lock file")?;
        
        // Try to acquire shared lock with timeout (5 seconds)
        let start = std::time::Instant::now();
        loop {
            match lockfile.try_lock_shared() {
                Ok(_) => return Ok(lockfile),
                Err(_) => {
                    if start.elapsed() > Duration::from_secs(5) {
                        return Err(anyhow::anyhow!("Timeout waiting for read lock on blacklist"));
                    }
                    tokio::time::sleep(Duration::from_millis(50)).await;
                }
            }
        }
    }

    /// Acquire an exclusive (write) lock on the blacklist file
    /// Returns a guard that releases the lock when dropped
    async fn acquire_write_lock(&self) -> Result<std::fs::File> {
        // Ensure lock file exists
        if let Some(parent) = self.lockfile_path.parent() {
            fs::create_dir_all(parent).await.ok();
        }
        
        let lockfile = std::fs::OpenOptions::new()
            .create(true)
            .read(true)
            .write(true)
            .open(&self.lockfile_path)
            .context("Failed to open lock file")?;
        
        // Try to acquire exclusive lock with timeout (5 seconds)
        let start = std::time::Instant::now();
        loop {
            match lockfile.try_lock_exclusive() {
                Ok(_) => return Ok(lockfile),
                Err(_) => {
                    if start.elapsed() > Duration::from_secs(5) {
                        return Err(anyhow::anyhow!("Timeout waiting for write lock on blacklist"));
                    }
                    tokio::time::sleep(Duration::from_millis(50)).await;
                }
            }
        }
    }

    /// 🚀 OPTIMIZATION: Check membership (cached hash first, then compute)
    /// This is the hot path - optimized for maximum speed with lock-free reads
    pub fn has(&self, addr: &str) -> bool {
        // Fast path: check cache first for hash (lock-free read)
        if let Some(cached_hash) = self.hash_cache.get(addr) {
            return self.set.contains(cached_hash); // Lock-free lookup
        }

        // Compute hash if not cached
        let hash = fast_hash_8_bytes(addr);

        // Check set directly (still lock-free)
        self.set.contains(&hash)
    }

    /// Load file into memory (streaming, safe for large files)
    /// Uses shared lock to allow concurrent reads
    pub async fn load(&mut self) -> Result<()> {
        // Acquire shared lock for reading
        let _lock_guard = self.acquire_read_lock().await?;
        
        // Ensure parent directory exists
        if let Some(parent) = self.filepath.parent() {
            fs::create_dir_all(parent).await.ok();
        }

        // Check if file exists
        if !self.filepath.exists() {
            // Create empty file
            File::create(&self.filepath).await?;
            return Ok(());
        }

        let file = File::open(&self.filepath).await?;
        let reader = BufReader::new(file);
        let mut lines = reader.lines();

        while let Some(line) = lines.next_line().await? {
            if line.is_empty() {
                continue;
            }

            // Try to parse as Entry
            match serde_json::from_str::<Entry>(&line) {
                Ok(entry) => {
                    let hash = entry.hash.clone();
                    self.set.insert(hash.clone());
                    // Preserve address and timestamp mapping if available
                    if let Some(addr) = entry.addr {
                        self.hash_to_addr.insert(hash.clone(), addr);
                    }
                    if entry.ts > 0 {
                        self.hash_to_ts.insert(hash, entry.ts);
                    }
                }
                Err(_) => {
                    // If corrupted line, try simple fallback: treat line as plain address
                    let trimmed = line.trim();
                    if !trimmed.is_empty() {
                        // Assume it's an address string; convert to short-hash
                        let hash = fast_hash_8_bytes(trimmed);
                        self.set.insert(hash.clone());
                        self.hash_to_addr.insert(hash.clone(), trimmed.to_string());
                        // Set timestamp to current time for entries without timestamp
                        let ts = std::time::SystemTime::now()
                            .duration_since(std::time::UNIX_EPOCH)
                            .unwrap()
                            .as_millis() as u64;
                        self.hash_to_ts.insert(hash, ts);
                    }
                }
            }
        }

        Ok(())
    }

    /// Add address if missing; append to jsonl file
    /// Uses exclusive lock to prevent concurrent writes
    pub async fn add(&mut self, addr: &str) -> Result<bool> {
        // Acquire exclusive lock for writing
        let _lock_guard = self.acquire_write_lock().await?;
        
        let hash = self.addr_to_short_hash(addr);
        
        // Double-check after acquiring lock (another process might have added it)
        if self.set.contains(&hash) {
            return Ok(false); // already present
        }

        // Add to in-memory set first (optimistic)
        self.set.insert(hash.clone());

        // Store hash -> address and timestamp mappings to preserve during snapshot
        let ts = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_millis() as u64;
        self.hash_to_addr.insert(hash.clone(), addr.to_string());
        self.hash_to_ts.insert(hash.clone(), ts);

        let entry = Entry {
            addr: Some(addr.to_string()),
            ts,
            hash: hash.clone(),
        };

        let line = serde_json::to_string(&entry)? + "\n";

        // Write with exclusive lock held
        match OpenOptions::new()
            .create(true)
            .append(true)
            .open(&self.filepath)
            .await
        {
            Ok(mut file) => {
                file.write_all(line.as_bytes()).await?;
                file.flush().await?;
                Ok(true)
            }
            Err(err) => {
                // If append failed, remove from set so it stays consistent
                self.set.remove(&entry.hash);
                Err(err.into())
            }
        }
    }

    /// 🚀 OPTIMIZATION: Hash with caching for frequently checked addresses
    /// This is a hot path - optimized for maximum speed
    fn addr_to_short_hash(&mut self, addr: &str) -> String {
        // Check cache first (O(1) lookup)
        if let Some(cached) = self.hash_cache.get(addr) {
            return cached.clone();
        }

        // Compute hash
        let hash = fast_hash_8_bytes(addr);

        // Update cache (LRU: remove oldest if cache full)
        if self.hash_cache.len() >= self.cache_size {
            // Remove first entry (FIFO, simple LRU approximation)
            if let Some(first_key) = self.hash_cache.keys().next().cloned() {
                self.hash_cache.remove(&first_key);
            }
        }
        self.hash_cache.insert(addr.to_string(), hash.clone());

        hash
    }
}

// Test and benchmark the blacklist implementation
// Run with: cargo run --bin test_blacklist --release

use std::time::Instant;

mod blacklist {
    include!("../../src/blacklist.rs");
}

use blacklist::Blacklist;

#[tokio::main]
async fn main() {
    println!("🧪 Blacklist Performance Tests\n");

    // Test 1: Basic functionality
    println!("📋 Test 1: Basic Functionality");
    test_basic_functionality().await;

    // Test 2: Performance benchmark
    println!("\n📋 Test 2: Performance Benchmark");
    test_performance().await;

    println!("\n✅ All tests passed!");
}

async fn test_basic_functionality() {
    let temp_file = "/tmp/test_blacklist_basic.jsonl";
    let mut bl = Blacklist::new(temp_file);

    // Load empty
    bl.load().await.expect("Failed to load");
    println!("  ✓ Loaded empty blacklist");

    // Add address
    let addr1 = "9pSo69eqU1fEmE5bvnSzcZ4uNdfaRMg2bWW7UaXGruzv";
    let added = bl.add(addr1).await.expect("Failed to add");
    assert!(added, "Should add new address");
    println!("  ✓ Added address: {}", addr1);

    // Check it exists
    assert!(bl.has(addr1), "Should find added address");
    println!("  ✓ Found address in blacklist");

    // Cleanup
    let _ = tokio::fs::remove_file(temp_file).await;
}

async fn test_performance() {
    let temp_file = "/tmp/test_blacklist_perf.jsonl";
    let mut bl = Blacklist::new(temp_file);

    // Generate test addresses
    let test_addresses: Vec<String> = (0..10_000)
        .map(|i| format!("TestAddress{:08}", i))
        .collect();

    // Benchmark: Add addresses
    println!("  Adding 10,000 addresses...");
    let start = Instant::now();
    for addr in &test_addresses {
        bl.add(addr).await.expect("Failed to add");
    }
    let add_duration = start.elapsed();
    println!(
        "  ✓ Added 10,000 addresses in {:.2}ms ({:.2}μs per add)",
        add_duration.as_secs_f64() * 1000.0,
        add_duration.as_micros() as f64 / 10_000.0
    );

    // Benchmark: Check membership (hot cache)
    println!("\n  Checking same 1,000 addresses repeatedly (hot cache)...");
    let hot_addresses = &test_addresses[0..1000];
    let start = Instant::now();
    let iterations = 100;
    for _ in 0..iterations {
        for addr in hot_addresses {
            assert!(bl.has(addr));
        }
    }
    let check_hot_duration = start.elapsed();
    let total_checks = 1000 * iterations;
    println!(
        "  ✓ Checked {} addresses in {:.2}ms ({:.2}ns per check)",
        total_checks,
        check_hot_duration.as_secs_f64() * 1000.0,
        check_hot_duration.as_nanos() as f64 / total_checks as f64
    );

    // Cleanup
    let _ = tokio::fs::remove_file(temp_file).await;
}

// add_blacklist.rs - Command-line tool to add addresses to blacklist
// Usage: cargo run --bin add_blacklist -- --address <ADDRESS>

use anyhow::{Context, Result};
use colored::*;
use std::env;

mod blacklist {
    include!("../blacklist.rs");
}

use blacklist::Blacklist;

#[tokio::main]
async fn main() -> Result<()> {
    let args: Vec<String> = env::args().collect();

    // Parse command-line arguments
    let mut address: Option<String> = None;
    let mut blacklist_path: Option<String> = None;

    let mut i = 1;
    while i < args.len() {
        match args[i].as_str() {
            "--address" | "-a" => {
                if i + 1 < args.len() {
                    address = Some(args[i + 1].clone());
                    i += 2;
                } else {
                    eprintln!("{}", "❌ --address requires an address argument".red());
                    print_usage();
                    std::process::exit(1);
                }
            }
            "--file" | "-f" => {
                if i + 1 < args.len() {
                    blacklist_path = Some(args[i + 1].clone());
                    i += 2;
                } else {
                    eprintln!("{}", "❌ --file requires a file path argument".red());
                    print_usage();
                    std::process::exit(1);
                }
            }
            "--help" | "-h" => {
                print_usage();
                std::process::exit(0);
            }
            _ => {
                // If it starts with --, it's an unknown flag
                if args[i].starts_with("--") {
                    eprintln!("{}", format!("❌ Unknown argument: {}", args[i]).red());
                    print_usage();
                    std::process::exit(1);
                }
                // Otherwise, treat as address (positional argument)
                if address.is_none() {
                    address = Some(args[i].clone());
                }
                i += 1;
            }
        }
    }

    // Validate address is provided
    let addr = match address {
        Some(addr) => addr,
        None => {
            eprintln!("{}", "❌ Error: Address is required".red());
            print_usage();
            std::process::exit(1);
        }
    };

    // Use provided path or default
    let filepath = blacklist_path.unwrap_or_else(|| {
        env::var("BLACKLIST_PATH").unwrap_or_else(|_| "./blacklist.jsonl".to_string())
    });

    println!(
        "{}",
        format!("📝 Adding address to blacklist: {}", addr).cyan()
    );
    println!("{}", format!("📁 Blacklist file: {}", filepath).dimmed());

    // Load existing blacklist
    let mut blacklist = Blacklist::new(&filepath);

    println!("{}", "🔄 Loading existing blacklist...".cyan());
    blacklist.load().await.context("Failed to load blacklist")?;

    let initial_size = blacklist.size();
    println!(
        "{}",
        format!("✅ Loaded {} existing entries", initial_size).green()
    );

    // Check if already exists
    if blacklist.has(&addr) {
        println!(
            "{}",
            format!("ℹ️  Address {} is already in blacklist", addr).yellow()
        );
        return Ok(());
    }

    // Add address
    println!("{}", format!("➕ Adding address: {}", addr).cyan());
    match blacklist.add(&addr).await {
        Ok(added) => {
            if added {
                let new_size = blacklist.size();
                println!(
                    "{}",
                    format!("✅ Successfully added address to blacklist",).green()
                );
                println!(
                    "{}",
                    format!("📊 Blacklist size: {} → {}", initial_size, new_size).cyan()
                );
            } else {
                println!(
                    "{}",
                    format!("⚠️  Address was not added (may have been added concurrently)",)
                        .yellow()
                );
            }
        }
        Err(e) => {
            eprintln!("{}", format!("❌ Failed to add address: {}", e).red());
            return Err(e);
        }
    }

    Ok(())
}

fn print_usage() {
    println!("{}", "Usage:".bold());
    println!("  cargo run --bin add_blacklist -- --address <ADDRESS> [OPTIONS]");
    println!("\n{}", "Arguments:".bold());
    println!("  --address, -a <ADDRESS>    Solana address to add to blacklist");
    println!("  --file, -f <PATH>          Path to blacklist file (default: ./blacklist.jsonl)");
    println!("  --help, -h                 Show this help message");
    println!("\n{}", "Examples:".bold());
    println!(
        "  cargo run --bin add_blacklist -- --address GafDS9b8ZF95cNNvEeHpBq3kjthdKpf4RdhAz62MLPtZ"
    );
}

use blacklist::Blacklist;

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize blacklist
    let mut blacklist = Blacklist::new("./blacklist.jsonl");
    blacklist.load().await?;
    
    println!("✅ Loaded {} blacklisted wallets", blacklist.size());

    // In your snipe bot loop
    loop {
        let token_creator = get_new_token_creator().await?;
        
        // In-memory check, no network call
        if blacklist.has(&token_creator) {
            println!("⚠️  Skipping token from blacklisted wallet: {}", token_creator);
            continue;
        }
        
        // Proceed with sniping logic
        snipe_token(&token_creator).await?;
    }
}

// Reload blacklist daily at 00:00 UTC
tokio::spawn(async move {
    loop {
        // Wait until next midnight UTC
        let now = chrono::Utc::now();
        let next_midnight = (now + chrono::Duration::days(1))
            .date()
            .and_hms(0, 0, 0);
        let wait_duration = (next_midnight - now).to_std().unwrap();
        
        tokio::time::sleep(wait_duration).await;
        
        // Reload blacklist
        if let Err(e) = blacklist.reload().await {
            eprintln!("Failed to reload blacklist: {}", e);
        } else {
            println!("✅ Blacklist reloaded with latest updates");
        }
    }
});

[dependencies]
anyhow = "1.0"
dashmap = "5.5"
fs2 = "0.4"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
tokio = { version = "1.35", features = ["full"] }

[dev-dependencies]
colored = "2.1"