test
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
target/
|
||||
blobs/
|
||||
identity.key
|
||||
Generated
+4443
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,21 @@
|
||||
[package]
|
||||
name = "lattice-proto"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[dependencies]
|
||||
anyhow = "1.0"
|
||||
blake3 = "1.5"
|
||||
bytes = "1.5"
|
||||
clap = { version = "4.5", features = ["derive"] }
|
||||
ed25519-dalek = { version = "2.1", features = ["rand_core"] }
|
||||
futures = "0.3"
|
||||
hex = "0.4"
|
||||
iroh = { version = "0.95.1", features = ["discovery-local-network"] }
|
||||
iroh-gossip = "0.95.0"
|
||||
iroh-tickets = "0.2.0"
|
||||
rand = "0.8"
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
tokio = { version = "1.36", features = ["full"] }
|
||||
tracing-subscriber = { version = "0.3.22", features = ["env-filter"] }
|
||||
@@ -0,0 +1,434 @@
|
||||
use serde::{Serialize, Deserialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::sync::{mpsc, Notify};
|
||||
use ed25519_dalek::{Signer, Verifier, Signature};
|
||||
use iroh::{Endpoint, SecretKey, PublicKey};
|
||||
use iroh::discovery::mdns::MdnsDiscovery;
|
||||
use iroh_gossip::net::Gossip;
|
||||
use iroh_tickets::endpoint::EndpointTicket;
|
||||
use iroh_gossip::proto::TopicId;
|
||||
use iroh::protocol::Router;
|
||||
use futures::StreamExt;
|
||||
use std::io::Write;
|
||||
|
||||
// --- Data Structures ---
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone)]
|
||||
pub struct Entry {
|
||||
pub path: String,
|
||||
pub content_hash: String,
|
||||
pub author: String,
|
||||
pub timestamp: u64,
|
||||
pub signature: Vec<u8>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone)]
|
||||
pub enum Message {
|
||||
Update(Entry),
|
||||
WantBlob(String),
|
||||
BlobData(String, Vec<u8>),
|
||||
}
|
||||
|
||||
// --- Logic ---
|
||||
|
||||
type Db = Arc<Mutex<HashMap<String, Entry>>>;
|
||||
|
||||
pub struct Node {
|
||||
pub keypair: ed25519_dalek::SigningKey,
|
||||
pub db: Db,
|
||||
pub blobs_dir: std::path::PathBuf,
|
||||
pub gossip_sender: Option<mpsc::Sender<Message>>,
|
||||
pub neighbors: Arc<Mutex<std::collections::HashSet<PublicKey>>>,
|
||||
pub pending_requests: Arc<Mutex<HashMap<String, Arc<Notify>>>>,
|
||||
}
|
||||
|
||||
impl Node {
|
||||
pub fn new(blobs_dir: std::path::PathBuf) -> Self {
|
||||
let key_path = std::path::Path::new("identity.key");
|
||||
let keypair = if key_path.exists() {
|
||||
let bytes = std::fs::read(key_path).expect("Failed to read key");
|
||||
ed25519_dalek::SigningKey::from_bytes(bytes.as_slice().try_into().unwrap())
|
||||
} else {
|
||||
let pk = ed25519_dalek::SigningKey::generate(&mut rand::rngs::OsRng);
|
||||
std::fs::write(key_path, pk.to_bytes()).expect("Failed to write key");
|
||||
pk
|
||||
};
|
||||
|
||||
if !blobs_dir.exists() {
|
||||
std::fs::create_dir_all(&blobs_dir).expect("Failed to create blobs dir");
|
||||
}
|
||||
|
||||
Self {
|
||||
keypair,
|
||||
db: Arc::new(Mutex::new(HashMap::new())),
|
||||
blobs_dir,
|
||||
gossip_sender: None,
|
||||
neighbors: Arc::new(Mutex::new(std::collections::HashSet::new())),
|
||||
pending_requests: Arc::new(Mutex::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn load_peers(&self) -> Vec<PublicKey> {
|
||||
let peers_file = self.blobs_dir.parent().unwrap().join("peers.txt");
|
||||
if !peers_file.exists() {
|
||||
return Vec::new();
|
||||
}
|
||||
let content = std::fs::read_to_string(&peers_file).unwrap_or_default();
|
||||
content.lines()
|
||||
.filter_map(|line| std::str::FromStr::from_str(line).ok())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn save_peer(&self, peer: PublicKey) {
|
||||
let peers_file = self.blobs_dir.parent().unwrap().join("peers.txt");
|
||||
// Avoid duplicates in file trivially by checking if already known in memory?
|
||||
// But restart clears memory.
|
||||
// Just append. Ideally we read all and check, but performance.
|
||||
// Let's just append.
|
||||
let mut file = std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(peers_file)
|
||||
.expect("Failed to open peers file");
|
||||
writeln!(file, "{}", peer).expect("Failed to write peer");
|
||||
}
|
||||
|
||||
pub fn pubkey_hex(&self) -> String {
|
||||
hex::encode(self.keypair.verifying_key().to_bytes())
|
||||
}
|
||||
|
||||
pub fn sign_entry(&self, path: &str, hash: &str, timestamp: u64) -> Vec<u8> {
|
||||
let mut msg = Vec::new();
|
||||
msg.extend_from_slice(path.as_bytes());
|
||||
msg.extend_from_slice(hash.as_bytes());
|
||||
msg.extend_from_slice(×tamp.to_le_bytes());
|
||||
self.keypair.sign(&msg).to_bytes().to_vec()
|
||||
}
|
||||
|
||||
pub fn verify_entry(entry: &Entry) -> bool {
|
||||
let pubkey_bytes = match hex::decode(&entry.author) {
|
||||
Ok(b) => b,
|
||||
Err(_) => return false,
|
||||
};
|
||||
let pubkey_arr: [u8; 32] = match pubkey_bytes.as_slice().try_into() {
|
||||
Ok(a) => a,
|
||||
Err(_) => return false,
|
||||
};
|
||||
let pubkey = match ed25519_dalek::VerifyingKey::from_bytes(&pubkey_arr) {
|
||||
Ok(pk) => pk,
|
||||
Err(_) => return false,
|
||||
};
|
||||
|
||||
let mut msg = Vec::new();
|
||||
msg.extend_from_slice(entry.path.as_bytes());
|
||||
msg.extend_from_slice(entry.content_hash.as_bytes());
|
||||
msg.extend_from_slice(&entry.timestamp.to_le_bytes());
|
||||
|
||||
let sig_arr: [u8; 64] = match entry.signature.as_slice().try_into() {
|
||||
Ok(a) => a,
|
||||
Err(_) => return false,
|
||||
};
|
||||
let sig = Signature::from_bytes(&sig_arr);
|
||||
|
||||
pubkey.verify(&msg, &sig).is_ok()
|
||||
}
|
||||
|
||||
pub fn put_local(&self, path: String, content: String) -> Entry {
|
||||
let hash = blake3::hash(content.as_bytes()).to_hex().to_string();
|
||||
let blob_path = self.blobs_dir.join(&hash);
|
||||
std::fs::write(blob_path, content).expect("Failed to write blob");
|
||||
|
||||
let timestamp = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_micros() as u64;
|
||||
|
||||
let signature = self.sign_entry(&path, &hash, timestamp);
|
||||
|
||||
let entry = Entry {
|
||||
path: path.clone(),
|
||||
content_hash: hash.clone(),
|
||||
author: self.pubkey_hex(),
|
||||
timestamp,
|
||||
signature,
|
||||
};
|
||||
|
||||
println!("[LOCAL] Writing {} -> {}", path, hash);
|
||||
let mut db = self.db.lock().unwrap();
|
||||
db.insert(path, entry.clone());
|
||||
|
||||
entry
|
||||
}
|
||||
|
||||
pub fn process_update(&self, entry: Entry) -> bool {
|
||||
if !Self::verify_entry(&entry) {
|
||||
println!("[WARN] Invalid Signature for {}", entry.path);
|
||||
return false;
|
||||
}
|
||||
|
||||
let mut db = self.db.lock().unwrap();
|
||||
if let Some(existing) = db.get(&entry.path) {
|
||||
if entry.timestamp <= existing.timestamp {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
println!("[SYNC] Received Update: {} -> {}", entry.path, entry.content_hash);
|
||||
db.insert(entry.path.clone(), entry);
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
// --- Main ---
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
tracing_subscriber::fmt::init();
|
||||
|
||||
let blobs_dir = std::path::PathBuf::from("blobs");
|
||||
let node = Arc::new(Node::new(blobs_dir));
|
||||
println!("Lattice Node Started. Identity: {}", node.pubkey_hex());
|
||||
|
||||
// Convert Node key to Iroh SecretKey
|
||||
let secret_key = SecretKey::from_bytes(&node.keypair.to_bytes());
|
||||
let node_id = secret_key.public();
|
||||
|
||||
// Configure MDNS
|
||||
// Note: If build() returns Result, use ?
|
||||
let mdns = MdnsDiscovery::builder()
|
||||
.build(node_id)?;
|
||||
|
||||
// Bind endpoint
|
||||
let endpoint = Endpoint::builder()
|
||||
.secret_key(secret_key)
|
||||
.discovery(mdns)
|
||||
.bind()
|
||||
.await?;
|
||||
|
||||
println!("Iroh Node ID: {}", endpoint.secret_key().public());
|
||||
|
||||
// Spawn Gossip
|
||||
let gossip = Gossip::builder().spawn(endpoint.clone());
|
||||
|
||||
// Spawn Router (The Accept Loop handling Gossip ALPN)
|
||||
let router = Router::builder(endpoint.clone())
|
||||
.accept(iroh_gossip::ALPN, gossip.clone())
|
||||
.spawn(); // CORRECTED: No await.
|
||||
|
||||
|
||||
// Create and print Ticket
|
||||
let my_addr = endpoint.addr(); // CORRECTED: No ?
|
||||
|
||||
let ticket = EndpointTicket::new(my_addr);
|
||||
println!("My Ticket: {}", ticket);
|
||||
|
||||
let topic_bytes = blake3::hash(b"lattice-test-net").as_bytes().to_owned();
|
||||
let topic_id = TopicId::from_bytes(topic_bytes);
|
||||
|
||||
// Join
|
||||
println!("Joining Gossip Topic: lattice-test-net...");
|
||||
|
||||
// Load persisted peers
|
||||
let initial_peers = node.load_peers();
|
||||
if !initial_peers.is_empty() {
|
||||
println!("Loaded {} persisted peers.", initial_peers.len());
|
||||
// Populate neighbors set specifically for display if desired,
|
||||
// though NeighborUp will fire when connection is actually established.
|
||||
}
|
||||
|
||||
let (sink, mut stream) = gossip.subscribe(topic_id, initial_peers).await?.split();
|
||||
println!("Joined Gossip Topic: lattice-test-net");
|
||||
|
||||
// Command Loop
|
||||
let (tx_cmd, mut rx_cmd) = mpsc::channel::<String>(100);
|
||||
|
||||
// Stdin Task
|
||||
tokio::spawn(async move {
|
||||
let stdin = std::io::stdin();
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
if stdin.read_line(&mut line).is_ok() {
|
||||
if tx_cmd.send(line).await.is_err() { break; }
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
print!("> ");
|
||||
std::io::stdout().flush()?;
|
||||
|
||||
// Event Loop
|
||||
loop {
|
||||
tokio::select! {
|
||||
Some(line) = rx_cmd.recv() => {
|
||||
let node = node.clone();
|
||||
let sink = sink.clone();
|
||||
let endpoint = endpoint.clone();
|
||||
let gossip = gossip.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let parts: Vec<&str> = line.trim().split_whitespace().collect();
|
||||
match parts.as_slice() {
|
||||
["put", path, content] => {
|
||||
let entry = node.put_local(path.to_string(), content.to_string());
|
||||
let msg = Message::Update(entry);
|
||||
if let Ok(msg_bytes) = serde_json::to_vec(&msg) {
|
||||
sink.broadcast(msg_bytes.into()).await.ok();
|
||||
println!("[GOSSIP] Broadcasted update.");
|
||||
}
|
||||
},
|
||||
["get", path] => {
|
||||
let entry_opt = {
|
||||
let db = node.db.lock().unwrap();
|
||||
db.get(*path).cloned()
|
||||
};
|
||||
|
||||
if let Some(entry) = entry_opt {
|
||||
let blob_path = node.blobs_dir.join(&entry.content_hash);
|
||||
if blob_path.exists() {
|
||||
if let Ok(content) = std::fs::read_to_string(blob_path) {
|
||||
println!("Content: {}", content);
|
||||
}
|
||||
} else {
|
||||
println!("[MISSING BLOB] Need hash: {}", entry.content_hash);
|
||||
|
||||
let notify = Arc::new(Notify::new());
|
||||
node.pending_requests.lock().unwrap().insert(entry.content_hash.clone(), notify.clone());
|
||||
|
||||
let msg = Message::WantBlob(entry.content_hash.clone());
|
||||
if let Ok(msg_bytes) = serde_json::to_vec(&msg) {
|
||||
sink.broadcast(msg_bytes.into()).await.ok();
|
||||
println!("Requesting blob... waiting...");
|
||||
|
||||
// Wait for notification with timeout
|
||||
match tokio::time::timeout(std::time::Duration::from_secs(5), notify.notified()).await {
|
||||
Ok(_) => {
|
||||
// check again
|
||||
let blob_path = node.blobs_dir.join(&entry.content_hash);
|
||||
if let Ok(content) = std::fs::read_to_string(blob_path) {
|
||||
println!("Content: {}", content);
|
||||
} else {
|
||||
println!("Error reading received blob.");
|
||||
}
|
||||
},
|
||||
Err(_) => {
|
||||
println!("Timeout waiting for blob.");
|
||||
// Clean up
|
||||
node.pending_requests.lock().unwrap().remove(&entry.content_hash);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
println!("Path not found.");
|
||||
}
|
||||
},
|
||||
["list"] => {
|
||||
let db = node.db.lock().unwrap();
|
||||
for (k, v) in db.iter() {
|
||||
println!("{} -> {} (ts: {})", k, v.content_hash, v.timestamp);
|
||||
}
|
||||
},
|
||||
["connect", arg] => {
|
||||
use std::str::FromStr;
|
||||
let topic_bytes = blake3::hash(b"lattice-test-net").as_bytes().to_owned();
|
||||
let topic_id = TopicId::from_bytes(topic_bytes);
|
||||
|
||||
if let Ok(ticket) = EndpointTicket::from_str(arg) {
|
||||
println!("Connecting via Ticket...");
|
||||
let addr = iroh::EndpointAddr::from(ticket);
|
||||
let peer_id = addr.id;
|
||||
|
||||
// Connect first (ensure transport)
|
||||
match endpoint.connect(addr, iroh_gossip::ALPN).await {
|
||||
Ok(_) => {
|
||||
println!("Connected! Adding to Gossip...");
|
||||
match gossip.subscribe(topic_id, vec![peer_id]).await {
|
||||
Ok(_) => println!("Subscribed peer to gossip topic!"),
|
||||
Err(e) => println!("Gossip subscribe error: {}", e),
|
||||
}
|
||||
},
|
||||
Err(e) => println!("Connection failed: {}", e),
|
||||
}
|
||||
} else if let Ok(peer_id) = PublicKey::from_str(arg) {
|
||||
println!("Connecting via Node ID (Discovery) to: {}", peer_id);
|
||||
match endpoint.connect(peer_id, iroh_gossip::ALPN).await {
|
||||
Ok(_) => {
|
||||
println!("Connected! Adding to Gossip...");
|
||||
match gossip.subscribe(topic_id, vec![peer_id]).await {
|
||||
Ok(_) => println!("Subscribed peer to gossip topic!"),
|
||||
Err(e) => println!("Gossip subscribe error: {}", e),
|
||||
}
|
||||
},
|
||||
Err(e) => println!("Connection failed: {}", e),
|
||||
}
|
||||
} else {
|
||||
println!("Invalid ID or Ticket");
|
||||
}
|
||||
},
|
||||
["peers"] => {
|
||||
let neighbors = node.neighbors.lock().unwrap();
|
||||
println!("Connected Gossip Peers: {}", neighbors.len());
|
||||
for peer in neighbors.iter() {
|
||||
println!("- {}", peer);
|
||||
}
|
||||
},
|
||||
["quit"] => std::process::exit(0), // Can't break loop from spawn
|
||||
_ => println!("Unknown command. Usage: put <path> <content> | get <path> | list | connect <ticket_or_node_id> | peers"),
|
||||
}
|
||||
print!("> ");
|
||||
std::io::stdout().flush().ok();
|
||||
});
|
||||
},
|
||||
Some(res) = stream.next() => {
|
||||
let event = res?;
|
||||
match event {
|
||||
iroh_gossip::api::Event::Received(msg) => {
|
||||
if let Ok(message) = serde_json::from_slice::<Message>(&msg.content) {
|
||||
match message {
|
||||
Message::Update(entry) => { node.process_update(entry); },
|
||||
Message::WantBlob(hash) => {
|
||||
let blob_path = node.blobs_dir.join(&hash);
|
||||
if blob_path.exists() {
|
||||
if let Ok(content) = std::fs::read_to_string(&blob_path) {
|
||||
let reply = Message::BlobData(hash, content.into_bytes());
|
||||
if let Ok(reply_bytes) = serde_json::to_vec(&reply) {
|
||||
sink.broadcast(reply_bytes.into()).await.ok();
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
Message::BlobData(hash, data) => {
|
||||
let blob_path = node.blobs_dir.join(&hash);
|
||||
if !blob_path.exists() {
|
||||
std::fs::write(blob_path, data).expect("Failed to write blob");
|
||||
println!("[SYNC] Received Blob: {}", hash);
|
||||
|
||||
// Notify waiters
|
||||
if let Some(notify) = node.pending_requests.lock().unwrap().remove(&hash) {
|
||||
notify.notify_waiters();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
iroh_gossip::api::Event::NeighborUp(node_id) => {
|
||||
println!("[GOSSIP] Neighbor Up: {}", node_id);
|
||||
if node.neighbors.lock().unwrap().insert(node_id) {
|
||||
node.save_peer(node_id);
|
||||
}
|
||||
},
|
||||
iroh_gossip::api::Event::NeighborDown(node_id) => {
|
||||
println!("[GOSSIP] Neighbor Down: {}", node_id);
|
||||
node.neighbors.lock().unwrap().remove(&node_id);
|
||||
},
|
||||
_ => {} // Handle Lagged or other future variants
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
router.shutdown().await?;
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user