diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 034728c..e3e6271 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -154,10 +154,27 @@ therefore use dedicated ALPNs while sharing identity and connectivity. Music similarity is another shared extension protocol. `music-dht` owns its versioned ALPN, bounded request/response models, validation, and byte-stream framing. Applications own model selection, user consent, audio preprocessing, -embedding generation and storage, nearest-neighbor search, peer selection, and -result presentation. Compatible clients can therefore use different local -implementations while exchanging vectors only when their exact profile -fingerprints match. +embedding generation and storage, nearest-neighbor search, participation +policy, and result presentation. Compatible clients can therefore use +different local implementations while exchanging vectors only when their exact +profile fingerprints match. + +Similarity peer selection uses a second, schema-independent Kademlia overlay. +It derives a model-neutral 256-bit SimHash, maps fixed subsets into twelve +two-level LSH tables, and publishes one summary per owner and occupied coarse +bucket. The coarse key is routed through the overlay; compact fine-bucket +suffixes, one anonymous representative signature, and the owner's self-contained +connection ticket are carried in the value. This bounds a peer's presence in a +hot bucket independently of how many tracks it owns, makes multi-hop results +dialable, and avoids a canonical calibration file. + +Routing summaries are signed by the owner's persistent transport identity. +Replicas validate the owner signature, network, profile, derived key, bucket +bounds, issue time and TTL before storing or forwarding a record. The routing +cache uses a separate SQLite database because it is replaceable federation +state, not application embedding storage. Its ALPN is additive: peers that do +not implement it continue to use the unchanged `music-dht-v5` catalog network, +while upgraded peers bootstrap and gossip their own routing contacts. Capability discovery is the narrow exception to schema isolation. The bounded, self-versioned `furumi/capabilities/1` stream still validates the federation diff --git a/Cargo.lock b/Cargo.lock index 1d6def7..2a5cfe7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -801,7 +801,7 @@ checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" [[package]] name = "federation-net" -version = "0.2.0" +version = "0.3.0" dependencies = [ "anyhow", "blake3", @@ -1967,7 +1967,7 @@ dependencies = [ [[package]] name = "music-dht" -version = "0.3.1" +version = "0.4.0" dependencies = [ "anyhow", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 3db68c4..15e0c99 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ rust-version = "1.97" repository = "https://gt.hexor.cy/ab/frid" [workspace.dependencies] -federation-net = { path = "crates/federation-net", version = "0.2.0" } +federation-net = { path = "crates/federation-net", version = "0.3.0" } iroh = "1" iroh-base = "1" iroh-tickets = "1" diff --git a/README.md b/README.md index 37b0bef..831b532 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,8 @@ The workspace has two layers: Mainline DHT. - [`music-dht`](crates/music-dht) builds a Kademlia-style distributed music directory on top of that transport, including catalog discovery, content-id - lookup, direct byte streams, and shared Furumi wire types. + lookup, signed similarity-LSH routing, direct byte streams, and shared Furumi + wire types. Each participant is a client, router, and storage peer. Automatic rendezvous has no Frid-operated bootstrap service, while self-contained peer tickets diff --git a/crates/federation-net/Cargo.toml b/crates/federation-net/Cargo.toml index 3275363..b1b9990 100644 --- a/crates/federation-net/Cargo.toml +++ b/crates/federation-net/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "federation-net" -version = "0.2.0" +version = "0.3.0" description = "Generic peer-to-peer networking engine built on Iroh" readme = "README.md" documentation = "https://docs.rs/federation-net" diff --git a/crates/federation-net/src/engine.rs b/crates/federation-net/src/engine.rs index b759163..6a191bf 100644 --- a/crates/federation-net/src/engine.rs +++ b/crates/federation-net/src/engine.rs @@ -8,7 +8,7 @@ use std::time::Duration; use iroh::endpoint::{Connection, RecvStream, SendStream, VarInt, presets}; use iroh::protocol::{AcceptError, ProtocolHandler, Router}; -use iroh::{Endpoint, EndpointAddr, EndpointId, SecretKey}; +use iroh::{Endpoint, EndpointAddr, EndpointId, SecretKey, Signature}; use serde::Serialize; use serde::de::DeserializeOwned; use tokio::sync::{Semaphore, mpsc, watch}; @@ -1041,6 +1041,15 @@ impl NetworkEngine { self.shared.endpoint.id() } + /// Signs application data with this peer's persistent transport identity. + /// + /// The secret key never leaves the engine. Protocols can attach the + /// returned signature to records that may be forwarded by other peers; + /// recipients verify it against [`Self::endpoint_id`]. + pub fn sign_identity(&self, message: &[u8]) -> Signature { + self.shared.endpoint.secret_key().sign(message) + } + /// Returns the network this engine participates in. pub fn network_id(&self) -> NetworkId { self.shared.config.network_id diff --git a/crates/federation-net/src/lib.rs b/crates/federation-net/src/lib.rs index fdc77db..75ca714 100644 --- a/crates/federation-net/src/lib.rs +++ b/crates/federation-net/src/lib.rs @@ -78,6 +78,6 @@ pub use ticket::{PeerTicket, TICKET_VERSION}; // Re-exported Iroh types that appear in the public API. pub use iroh::endpoint::{RecvStream, SendStream}; -pub use iroh::{EndpointAddr, EndpointId, SecretKey}; +pub use iroh::{EndpointAddr, EndpointId, SecretKey, Signature}; // Re-exported so applications can use the generic ticket helpers. pub use iroh_tickets::Ticket; diff --git a/crates/music-dht/Cargo.toml b/crates/music-dht/Cargo.toml index e3c5419..2a084e0 100644 --- a/crates/music-dht/Cargo.toml +++ b/crates/music-dht/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "music-dht" -version = "0.3.1" +version = "0.4.0" description = "Distributed music library search: a Kademlia-style DHT on top of federation-net" readme = "README.md" documentation = "https://docs.rs/music-dht" diff --git a/crates/music-dht/README.md b/crates/music-dht/README.md index 2156fc3..2c7f604 100644 --- a/crates/music-dht/README.md +++ b/crates/music-dht/README.md @@ -91,6 +91,22 @@ Results may carry a shared 128-bit SimHash of their embedding. Clients can use its Hamming distance to suppress near-duplicate recordings across peers without transmitting every result vector. +`music_dht::similarity_lsh` and `music_dht::similarity_dht` add decentralized +peer routing without putting models or vector storage into this crate. Every +client computes the same 256-bit SimHash, groups its local signatures into a +fixed two-level LSH (`12 × 20` bits, split into a 10-bit DHT key and a compact +10-bit suffix), and publishes one anonymous summary per peer and coarse bucket. +The summaries contain no track metadata, only the owner's self-contained route, +and are signed by its existing iroh identity, so a forwarding peer cannot alter +or impersonate them. + +The routing overlay has its own schema-independent ALPN. Old clients therefore +remain compatible with ordinary catalog federation while upgraded peers find +one another and replicate expiring LSH summaries without a coordinator or a +global calibration file. Applications provide local routing signatures through +`SimilarityDht::sync_local_signatures`, use `SimilarityDht::find_peers` before +the direct similarity stream, and continue to perform exact search locally. + ## Trusted-device sync and listening history `music_dht::device_sync` is the canonical wire contract shared by Furumi diff --git a/crates/music-dht/src/capabilities.rs b/crates/music-dht/src/capabilities.rs index 556ef3d..94d27a6 100644 --- a/crates/music-dht/src/capabilities.rs +++ b/crates/music-dht/src/capabilities.rs @@ -33,6 +33,8 @@ pub const MUSIC_DHT_ID: &str = "music_dht"; pub const CATALOG_ID: &str = "catalog"; /// Stable protocol identifier for music-similarity streams. pub const SIMILARITY_ID: &str = "similarity"; +/// Stable protocol identifier for DHT-routed similarity discovery. +pub const SIMILARITY_DHT_ID: &str = "similarity_dht"; /// Stable protocol identifier for personal-device synchronization. pub const DEVICE_SYNC_ID: &str = "device_sync"; /// Stable protocol identifier for Jam playback control. @@ -73,6 +75,10 @@ impl CapabilityManifest { SIMILARITY_ID.to_string(), crate::similarity::SIMILARITY_PROTOCOL_VERSION, ); + protocols.insert( + SIMILARITY_DHT_ID.to_string(), + crate::similarity_lsh::SIMILARITY_DHT_PROTOCOL_VERSION, + ); protocols.insert( DEVICE_SYNC_ID.to_string(), crate::device_sync::DEVICE_SYNC_PROTOCOL_VERSION, @@ -212,6 +218,7 @@ mod tests { MUSIC_DHT_ID, CATALOG_ID, SIMILARITY_ID, + SIMILARITY_DHT_ID, DEVICE_SYNC_ID, JAM_ID, ] { diff --git a/crates/music-dht/src/lib.rs b/crates/music-dht/src/lib.rs index 0c3f578..77dea57 100644 --- a/crates/music-dht/src/lib.rs +++ b/crates/music-dht/src/lib.rs @@ -20,6 +20,9 @@ //! the exact same audio bytes. //! * [`similarity`] defines the bounded, model-neutral stream contract used by //! clients that independently generate compatible music embeddings. +//! * [`similarity_lsh`] turns those embeddings into deterministic compact +//! routing summaries, while [`similarity_dht`] stores and discovers signed +//! peer summaries without owning model inference or an embedding database. //! * Records are replicated to the `K` nodes whose ids are XOR-closest to //! each key. Publishers pick the targets from their routing table and send //! batched store requests (one pipeline per peer), so even a large library @@ -94,6 +97,10 @@ mod request; mod routing; mod service; pub mod similarity; +pub mod similarity_dht; +pub mod similarity_lsh; + +pub use similarity_lsh::SimilarityRouteEntry; pub use config::{ DEFAULT_EXPIRE_INTERVAL, DEFAULT_LOOKUP_TIMEOUT, DEFAULT_REPUBLISH_INTERVAL, diff --git a/crates/music-dht/src/service.rs b/crates/music-dht/src/service.rs index 09d9aae..979f693 100644 --- a/crates/music-dht/src/service.rs +++ b/crates/music-dht/src/service.rs @@ -5,8 +5,8 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use federation_net::{ - ByteStream, EndpointAddr, EndpointId, NetworkConfig, NetworkEngine, PeerTicket, SchemaId, - SecretKey, StreamAcceptor, + ByteStream, EndpointAddr, EndpointId, NetworkConfig, NetworkEngine, NetworkId, PeerTicket, + SchemaId, SecretKey, Signature, StreamAcceptor, }; use tokio::sync::mpsc; use tokio::task::JoinHandle; @@ -335,6 +335,17 @@ impl MusicDhtService { self.node.endpoint_id } + /// Returns the federation network this service participates in. + pub fn network_id(&self) -> NetworkId { + self.node.config.network_id + } + + /// Signs a forwarded application record with this peer's persistent + /// transport identity without exposing the secret key. + pub fn sign_identity(&self, message: &[u8]) -> Signature { + self.node.engine.sign_identity(message) + } + /// Returns the DHT identifier of this peer. pub fn node_id(&self) -> NodeId { self.node.node_id @@ -402,6 +413,24 @@ impl MusicDhtService { .map_err(Into::into) } + /// Opens a raw byte stream using a self-contained peer ticket. + /// + /// This is used by schema-independent overlays whose contacts can be + /// learned independently of the main music-DHT routing table. + pub async fn open_stream_to(&self, ticket: &PeerTicket, alpn: &[u8]) -> Result { + self.node.ensure_running()?; + if ticket.network_id != self.node.config.network_id { + return Err(MusicDhtError::Network( + "peer ticket belongs to another network".to_string(), + )); + } + self.node + .engine + .open_stream(ticket.endpoint_addr.clone(), alpn) + .await + .map_err(Into::into) + } + /// Synchronizes the published library with `specs`: the desired set of /// items this peer wants to share. /// diff --git a/crates/music-dht/src/similarity_dht.rs b/crates/music-dht/src/similarity_dht.rs new file mode 100644 index 0000000..09ce9d5 --- /dev/null +++ b/crates/music-dht/src/similarity_dht.rs @@ -0,0 +1,1144 @@ +//! Decentralized storage and lookup of signed similarity-LSH peer summaries. +//! +//! This is a schema-independent overlay carried by `music-dht` byte streams. +//! It reuses peer identities, tickets and Kademlia contact shapes, but keeps +//! replaceable routing summaries separate from durable music metadata. + +use std::collections::{HashMap, HashSet}; +use std::path::Path; +use std::str::FromStr; +use std::sync::{Arc, Mutex, MutexGuard, PoisonError, RwLock}; +use std::time::Duration; + +use federation_net::{ByteStream, EndpointId, NetworkId, PeerTicket, StreamAcceptor}; +use futures::stream::{self, FuturesUnordered, StreamExt}; +use rusqlite::{Connection, params}; +use serde::{Deserialize, Serialize}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt}; +use tokio::sync::Semaphore; +use tracing::{debug, info, warn}; + +use crate::MusicDhtService; +use crate::error::{MusicDhtError, Result}; +use crate::record::now_ms; +use crate::routing::{ALPHA, K, MAX_LOOKUP_REQUESTS, NodeContact, NodeId, RoutingTable, distance}; +use crate::similarity_lsh::{ + ROUTING_RECORD_TTL, ROUTING_TABLES, SIMILARITY_DHT_ALPN, SIMILARITY_DHT_PROTOCOL_VERSION, + SignedSimilarityRoute, SimilarityDhtKey, SimilarityRoutePayload, SimilarityRouteSpec, + build_route_specs, routing_query, routing_signature_distance, +}; + +/// Maximum signed summaries returned for one coarse bucket. +pub const MAX_SIMILARITY_DHT_RECORDS: usize = 256; +/// Maximum summaries in a single replicated store request. +pub const MAX_SIMILARITY_STORE_RECORDS: usize = 64; +/// Maximum encoded request. Store batches are split below this bound. +pub const MAX_SIMILARITY_DHT_REQUEST_BYTES: usize = 512 * 1024; +/// Maximum encoded response. A full bucket of worst-case records is bounded. +pub const MAX_SIMILARITY_DHT_RESPONSE_BYTES: usize = 16 * 1024 * 1024; +/// Periodic bootstrap interval for discovering upgraded routing peers. +pub const SIMILARITY_BOOTSTRAP_INTERVAL: Duration = Duration::from_secs(10 * 60); +/// Signed records are refreshed well before the shared 12-hour TTL. +pub const SIMILARITY_REPUBLISH_INTERVAL: Duration = Duration::from_secs(4 * 60 * 60); +/// Replica expiry sweep interval. +pub const SIMILARITY_EXPIRE_INTERVAL: Duration = Duration::from_secs(60); + +const MAX_BOOTSTRAP_PEERS: usize = 64; +const BOOTSTRAP_CONCURRENCY: usize = 8; +const PUBLISH_CONCURRENCY: usize = 8; +const LOOKUP_TIMEOUT: Duration = Duration::from_secs(15); +const REQUEST_TIMEOUT: Duration = Duration::from_secs(5); +const MAX_CLOCK_SKEW: Duration = Duration::from_secs(5 * 60); +const RECORD_FIXED_BYTES_ESTIMATE: usize = 192; +const ENTRY_BYTES_ESTIMATE: usize = 36; +const MAX_BATCH_BYTES_ESTIMATE: usize = 384 * 1024; + +const STORAGE_SCHEMA: &str = " +CREATE TABLE IF NOT EXISTS similarity_dht_records ( + dht_key BLOB NOT NULL, + owner_peer_id TEXT NOT NULL, + issued_at_ms INTEGER NOT NULL, + expires_at_ms INTEGER NOT NULL, + payload BLOB NOT NULL, + PRIMARY KEY (dht_key, owner_peer_id) +); + +CREATE INDEX IF NOT EXISTS idx_similarity_dht_records_expiry + ON similarity_dht_records(expires_at_ms); + +CREATE TABLE IF NOT EXISTS similarity_dht_peers ( + peer_id TEXT PRIMARY KEY, + node_id BLOB NOT NULL, + ticket TEXT NOT NULL, + last_seen_ms INTEGER NOT NULL +); +"; + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct RoutingRequest { + version: u16, + requester: NodeContact, + operation: RoutingOperation, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +enum RoutingOperation { + FindNode { target: [u8; 32] }, + FindValue { key: SimilarityDhtKey }, + StoreBatch { records: Vec }, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct RoutingResponse { + version: u16, + ok: bool, + error: Option, + nodes: Vec, + records: Vec, + stored: u32, +} + +impl RoutingResponse { + fn success(nodes: Vec, records: Vec, stored: u32) -> Self { + Self { + version: SIMILARITY_DHT_PROTOCOL_VERSION, + ok: true, + error: None, + nodes, + records, + stored, + } + } + + fn error(message: impl Into) -> Self { + Self { + version: SIMILARITY_DHT_PROTOCOL_VERSION, + ok: false, + error: Some(message.into()), + nodes: Vec::new(), + records: Vec::new(), + stored: 0, + } + } +} + +/// Statistics for one local-summary synchronization or republish. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct SimilarityDhtPublishStats { + /// Peer summary records built from the local signature set. + pub records: usize, + /// Distinct coarse DHT keys published. + pub keys: usize, + /// Remote routing nodes that accepted at least one replica. + pub remote_nodes: usize, + /// Whether at least one local replica was retained. + pub local_replica: bool, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +struct PeerScore { + distance: u32, + collisions: u16, +} + +#[derive(Clone)] +struct RoutingDatabase { + conn: Arc>, +} + +impl std::fmt::Debug for RoutingDatabase { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("RoutingDatabase") + .finish_non_exhaustive() + } +} + +impl RoutingDatabase { + async fn open(path: &Path) -> Result { + let path = path.to_path_buf(); + let conn = tokio::task::spawn_blocking(move || -> Result { + let conn = Connection::open(&path).map_err(|error| { + MusicDhtError::Database(format!("failed to open {}: {error}", path.display())) + })?; + conn.execute_batch(STORAGE_SCHEMA).map_err(|error| { + MusicDhtError::Database(format!( + "failed to initialize similarity DHT storage: {error}" + )) + })?; + Ok(conn) + }) + .await + .map_err(|error| MusicDhtError::Database(format!("database task panicked: {error}")))??; + Ok(Self { + conn: Arc::new(Mutex::new(conn)), + }) + } + + async fn call(&self, operation: F) -> Result + where + F: FnOnce(&Connection) -> rusqlite::Result + Send + 'static, + T: Send + 'static, + { + let conn = Arc::clone(&self.conn); + tokio::task::spawn_blocking(move || { + let guard = lock(&conn); + operation(&guard).map_err(|error| MusicDhtError::Database(error.to_string())) + }) + .await + .map_err(|error| MusicDhtError::Database(format!("database task panicked: {error}")))? + } + + async fn store_records(&self, records: Vec) -> Result> { + self.call(move |conn| { + let transaction = conn.unchecked_transaction()?; + let mut stored = Vec::with_capacity(records.len()); + for record in records { + let key = record + .key() + .map_err(|error| rusqlite::Error::ToSqlConversionFailure(Box::new(error)))?; + let owner = record.payload.owner.to_string(); + let issued_at_ms = record.payload.issued_at_ms; + let expires_at_ms = + issued_at_ms.saturating_add(ROUTING_RECORD_TTL.as_millis() as u64); + let payload = postcard::to_stdvec(&record) + .map_err(|error| rusqlite::Error::ToSqlConversionFailure(Box::new(error)))?; + let changed = transaction.execute( + "INSERT INTO similarity_dht_records ( + dht_key, owner_peer_id, issued_at_ms, expires_at_ms, payload + ) VALUES (?1, ?2, ?3, ?4, ?5) + ON CONFLICT(dht_key, owner_peer_id) DO UPDATE SET + issued_at_ms = excluded.issued_at_ms, + expires_at_ms = excluded.expires_at_ms, + payload = excluded.payload + WHERE excluded.issued_at_ms >= similarity_dht_records.issued_at_ms", + params![ + key.as_bytes().as_slice(), + owner, + issued_at_ms as i64, + expires_at_ms as i64, + payload, + ], + )? > 0; + stored.push(changed); + } + transaction.commit()?; + Ok(stored) + }) + .await + } + + async fn records_by_key( + &self, + key: SimilarityDhtKey, + at_ms: u64, + ) -> Result> { + self.call(move |conn| { + let mut statement = conn.prepare( + "SELECT payload FROM similarity_dht_records + WHERE dht_key = ?1 AND expires_at_ms > ?2 + ORDER BY expires_at_ms DESC, owner_peer_id + LIMIT ?3", + )?; + let rows = statement.query_map( + params![ + key.as_bytes().as_slice(), + at_ms as i64, + MAX_SIMILARITY_DHT_RECORDS as i64 + ], + |row| row.get::<_, Vec>(0), + )?; + let mut records = Vec::new(); + for row in rows { + if let Ok(record) = postcard::from_bytes::(&row?) { + records.push(record); + } + } + Ok(records) + }) + .await + } + + async fn delete_expired(&self, at_ms: u64) -> Result { + self.call(move |conn| { + conn.execute( + "DELETE FROM similarity_dht_records WHERE expires_at_ms <= ?1", + [at_ms as i64], + ) + }) + .await + } + + async fn upsert_peer(&self, contact: NodeContact) -> Result<()> { + self.call(move |conn| { + conn.execute( + "INSERT OR REPLACE INTO similarity_dht_peers + (peer_id, node_id, ticket, last_seen_ms) + VALUES (?1, ?2, ?3, ?4)", + params![ + contact.peer_id.to_string(), + contact.node_id.as_bytes().as_slice(), + contact.ticket, + contact.last_seen_ms as i64, + ], + )?; + Ok(()) + }) + .await + } + + async fn load_peers(&self) -> Result> { + self.call(move |conn| { + let mut statement = + conn.prepare("SELECT peer_id, ticket, last_seen_ms FROM similarity_dht_peers")?; + let rows = statement.query_map([], |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, String>(1)?, + row.get::<_, i64>(2)?, + )) + })?; + let mut peers = Vec::new(); + for row in rows { + let (peer, ticket, last_seen_ms) = row?; + let Ok(peer_id) = EndpointId::from_str(&peer) else { + continue; + }; + peers.push(NodeContact { + node_id: NodeId::from_endpoint(&peer_id), + peer_id, + ticket, + last_seen_ms: last_seen_ms.max(0) as u64, + }); + } + Ok(peers) + }) + .await + } +} + +/// Reusable similarity-routing DHT attached to a [`MusicDhtService`]. +pub struct SimilarityDht { + service: Arc, + database: RoutingDatabase, + network_id: NetworkId, + own_contact: NodeContact, + routing: Mutex, + local_specs: RwLock>, + publish_gate: Semaphore, +} + +impl std::fmt::Debug for SimilarityDht { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("SimilarityDht") + .field("endpoint_id", &self.own_contact.peer_id) + .field("known_peers", &lock(&self.routing).len()) + .finish_non_exhaustive() + } +} + +impl SimilarityDht { + /// Opens the replaceable routing cache and restores previously learned + /// similarity-routing contacts. The caller owns the serving and + /// maintenance tasks so they share the application's shutdown lifecycle. + pub async fn open( + service: Arc, + database_path: impl AsRef, + ) -> Result> { + let database = RoutingDatabase::open(database_path.as_ref()).await?; + let network_id = service.network_id(); + let ticket = service.ticket().await?.to_string(); + let owner = service.endpoint_id(); + let own_contact = NodeContact { + node_id: NodeId::from_endpoint(&owner), + peer_id: owner, + ticket, + last_seen_ms: now_ms(), + }; + let mut routing = RoutingTable::new(own_contact.node_id); + for contact in database.load_peers().await? { + if validate_contact(&contact, &network_id, Some(owner)).is_ok() { + routing.upsert(contact); + } + } + Ok(Arc::new(Self { + service, + database, + network_id, + own_contact, + routing: Mutex::new(routing), + local_specs: RwLock::new(Vec::new()), + publish_gate: Semaphore::new(1), + })) + } + + /// Number of currently known peers that implement the routing overlay. + pub fn known_peers(&self) -> usize { + lock(&self.routing).len() + } + + /// Serves authenticated routing requests until `acceptor` closes. + pub async fn serve(self: Arc, mut acceptor: StreamAcceptor) { + while let Some(stream) = acceptor.accept().await { + let this = Arc::clone(&self); + tokio::spawn(async move { + let peer = stream.peer_id; + if let Err(error) = this.serve_one(stream).await { + debug!(%peer, %error, "similarity DHT request failed"); + } + }); + } + } + + /// Periodically discovers upgraded peers, republishes local summaries and + /// removes expired replicas. The first bootstrap runs immediately. + pub async fn maintenance(self: Arc) { + let mut bootstrap = tokio::time::interval(SIMILARITY_BOOTSTRAP_INTERVAL); + let mut republish = tokio::time::interval(SIMILARITY_REPUBLISH_INTERVAL); + let mut expire = tokio::time::interval(SIMILARITY_EXPIRE_INTERVAL); + bootstrap.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + republish.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + expire.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + // `republish` must wait for either a local sync or a full interval. + republish.tick().await; + loop { + tokio::select! { + _ = bootstrap.tick() => { + let before = self.known_peers(); + if let Err(error) = self.bootstrap().await { + debug!(%error, "similarity DHT bootstrap failed"); + } else { + let has_local_specs = !read(&self.local_specs).is_empty(); + if self.known_peers() > before && has_local_specs + && let Err(error) = self.publish_local().await + { + warn!(%error, "similarity DHT publication after bootstrap failed"); + } + } + } + _ = republish.tick() => { + let has_local_specs = !read(&self.local_specs).is_empty(); + if has_local_specs + && let Err(error) = self.publish_local().await + { + warn!(%error, "similarity DHT republish failed"); + } + } + _ = expire.tick() => { + match self.database.delete_expired(now_ms()).await { + Ok(0) => {} + Ok(count) => debug!(count, "expired similarity DHT replicas removed"), + Err(error) => warn!(%error, "similarity DHT expiry sweep failed"), + } + } + } + } + } + + /// Replaces the locally published active profile with summaries built + /// from all local routing signatures, then publishes them immediately. + pub async fn sync_local_signatures( + &self, + profile_id: String, + signatures: Vec<[u8; 32]>, + ) -> Result { + let specs = + tokio::task::spawn_blocking(move || build_route_specs(&profile_id, &signatures)) + .await + .map_err(|error| { + MusicDhtError::Protocol(format!("LSH build task panicked: {error}")) + })??; + *write(&self.local_specs) = specs; + self.bootstrap().await?; + self.publish_local().await + } + + /// Stops advertising local summaries. Existing replicas disappear by the + /// shared TTL; no network-wide delete broadcast is required. + pub fn clear_local_signatures(&self) { + write(&self.local_specs).clear(); + } + + /// Finds and ranks self-contained routes to peers likely to own tracks + /// close to `vector`. + pub async fn find_peers( + &self, + profile_id: &str, + vector: &[f32], + limit: usize, + ) -> Result> { + if limit == 0 { + return Ok(Vec::new()); + } + let query = routing_query(vector)?; + if self.known_peers() == 0 { + self.bootstrap().await?; + } + let mut lookups = Vec::with_capacity(ROUTING_TABLES * 2); + let mut table_probes = Vec::with_capacity(ROUTING_TABLES); + for table in 0..ROUTING_TABLES as u8 { + let probes = query.probes(profile_id, table)?; + for primary_bucket in &probes.primary_buckets { + lookups.push(( + table, + *primary_bucket, + SimilarityDhtKey::derive(&self.network_id, profile_id, table, *primary_bucket)?, + )); + } + table_probes.push(probes); + } + + let outcomes = + stream::iter(lookups.into_iter().map(|(table, primary, key)| async move { + (table, primary, self.lookup(key).await) + })) + .buffer_unordered(8) + .collect::>() + .await; + + let mut scores: HashMap = HashMap::new(); + for (table, primary, outcome) in outcomes { + let records = match outcome { + Ok(records) => records, + Err(error) => { + debug!(table, primary, %error, "similarity DHT bucket lookup failed"); + continue; + } + }; + let probes = &table_probes[table as usize]; + for record in records { + let payload = &record.payload; + if payload.owner == self.own_contact.peer_id + || payload.profile_id != profile_id + || payload.table != table + || payload.primary_bucket != primary + { + continue; + } + let representatives = probes.suffix_buckets.iter().filter_map(|suffix| { + payload + .entries + .binary_search_by_key(suffix, |entry| entry.suffix) + .ok() + .map(|index| payload.entries[index].representative) + }); + for representative in representatives { + let candidate_distance = + routing_signature_distance(query.signature(), &representative); + scores + .entry(payload.owner) + .and_modify(|(score, ticket, issued_at_ms)| { + score.distance = score.distance.min(candidate_distance); + score.collisions = score.collisions.saturating_add(1); + if payload.issued_at_ms > *issued_at_ms { + ticket.clone_from(&payload.owner_ticket); + *issued_at_ms = payload.issued_at_ms; + } + }) + .or_insert(( + PeerScore { + distance: candidate_distance, + collisions: 1, + }, + payload.owner_ticket.clone(), + payload.issued_at_ms, + )); + } + } + } + let mut peers: Vec<_> = scores.into_iter().collect(); + peers.sort_by_key(|(peer, (score, _, _))| { + (score.distance, std::cmp::Reverse(score.collisions), *peer) + }); + peers.truncate(limit); + peers + .into_iter() + .map(|(_, (_, ticket, _))| { + PeerTicket::from_str(&ticket) + .map_err(|error| MusicDhtError::InvalidTicket(error.to_string())) + }) + .collect() + } + + async fn serve_one(&self, mut stream: ByteStream) -> Result<()> { + let authenticated_peer = stream.peer_id; + let request: RoutingRequest = + read_postcard(&mut stream.recv, MAX_SIMILARITY_DHT_REQUEST_BYTES).await?; + let response = match self.handle_request(authenticated_peer, request).await { + Ok(response) => response, + Err(error) => RoutingResponse::error(error.to_string()), + }; + write_postcard( + &mut stream.send, + &response, + MAX_SIMILARITY_DHT_RESPONSE_BYTES, + ) + .await?; + stream.send.finish().map_err(network_error)?; + let _ = tokio::time::timeout(Duration::from_secs(2), stream.send.stopped()).await; + Ok(()) + } + + async fn handle_request( + &self, + authenticated_peer: EndpointId, + request: RoutingRequest, + ) -> Result { + if request.version != SIMILARITY_DHT_PROTOCOL_VERSION { + return Err(protocol_error("unsupported similarity DHT request version")); + } + if request.requester.peer_id != authenticated_peer { + return Err(protocol_error("similarity DHT requester identity mismatch")); + } + self.learn_contact(request.requester).await?; + match request.operation { + RoutingOperation::FindNode { target } => Ok(RoutingResponse::success( + self.closest_contacts(&target, K), + Vec::new(), + 0, + )), + RoutingOperation::FindValue { key } => { + let records = self.database.records_by_key(key, now_ms()).await?; + Ok(RoutingResponse::success( + self.closest_contacts(key.as_bytes(), K), + records, + 0, + )) + } + RoutingOperation::StoreBatch { records } => { + if records.is_empty() || records.len() > MAX_SIMILARITY_STORE_RECORDS { + return Err(protocol_error("invalid similarity DHT store batch")); + } + let now = now_ms(); + for record in &records { + validate_record(record, &self.network_id, now)?; + } + let stored = self + .database + .store_records(records) + .await? + .into_iter() + .filter(|stored| *stored) + .count(); + Ok(RoutingResponse::success( + Vec::new(), + Vec::new(), + stored as u32, + )) + } + } + } + + async fn bootstrap(&self) -> Result<()> { + let known_feature: HashSet = lock(&self.routing) + .contacts() + .into_iter() + .map(|contact| contact.peer_id) + .collect(); + let mut candidates = lock(&self.routing).contacts(); + candidates.extend( + self.service + .known_peers() + .into_iter() + .filter(|contact| !known_feature.contains(&contact.peer_id)), + ); + let mut seen = HashSet::new(); + candidates.retain(|contact| { + contact.peer_id != self.own_contact.peer_id && seen.insert(contact.peer_id) + }); + candidates.truncate(MAX_BOOTSTRAP_PEERS); + let target = self.own_contact.node_id.as_bytes().to_owned(); + let responses = stream::iter(candidates.into_iter().map(|contact| async move { + let response = self + .exchange(&contact, RoutingOperation::FindNode { target }) + .await; + (contact, response) + })) + .buffer_unordered(BOOTSTRAP_CONCURRENCY) + .collect::>() + .await; + for (contact, response) in responses { + let response = match response { + Ok(response) => response, + Err(error) => { + debug!(peer = %contact.peer_id, %error, "similarity DHT bootstrap peer unavailable"); + continue; + } + }; + self.learn_contact(contact).await?; + self.learn_response_nodes(response.nodes).await; + } + Ok(()) + } + + async fn publish_local(&self) -> Result { + let _permit = self + .publish_gate + .acquire() + .await + .map_err(|_| protocol_error("similarity DHT publisher closed"))?; + let specs = read(&self.local_specs).clone(); + if specs.is_empty() { + return Ok(SimilarityDhtPublishStats::default()); + } + let issued_at_ms = now_ms(); + let mut stats = SimilarityDhtPublishStats { + records: specs.len(), + keys: specs.len(), + ..SimilarityDhtPublishStats::default() + }; + let mut accepted_peers = HashSet::new(); + for wave in specs.chunks(64) { + let mut local = Vec::new(); + let mut remote: HashMap)> = + HashMap::new(); + for spec in wave { + let payload = SimilarityRoutePayload::from_spec( + spec.clone(), + self.network_id, + self.own_contact.peer_id, + self.own_contact.ticket.clone(), + issued_at_ms, + )?; + let record = SignedSimilarityRoute::sign(payload, |bytes| { + self.service.sign_identity(bytes) + })?; + let key = record.key()?; + let (store_local, contacts) = self.replica_targets(&key); + if store_local { + local.push(record.clone()); + } + for contact in contacts { + remote + .entry(contact.peer_id) + .or_insert_with(|| (contact, Vec::new())) + .1 + .push(record.clone()); + } + } + if !local.is_empty() { + self.database.store_records(local).await?; + stats.local_replica = true; + } + let sends = stream::iter(remote.into_values().map(|(contact, records)| async move { + let accepted = self.store_on_peer(&contact, records).await.unwrap_or(false); + (contact.peer_id, accepted) + })) + .buffer_unordered(PUBLISH_CONCURRENCY) + .collect::>() + .await; + for (peer, accepted) in sends { + if accepted { + accepted_peers.insert(peer); + } + } + } + stats.remote_nodes = accepted_peers.len(); + info!( + records = stats.records, + nodes = stats.remote_nodes, + "similarity DHT summaries published" + ); + Ok(stats) + } + + async fn store_on_peer( + &self, + contact: &NodeContact, + records: Vec, + ) -> Result { + let mut accepted = false; + for batch in chunk_store_records(records) { + let response = self + .exchange(contact, RoutingOperation::StoreBatch { records: batch }) + .await?; + accepted |= response.stored > 0; + } + Ok(accepted) + } + + async fn lookup(&self, key: SimilarityDhtKey) -> Result> { + let local = self.database.records_by_key(key, now_ms()).await?; + let mut records: HashMap = HashMap::new(); + merge_records(&mut records, local, &self.network_id, key); + + let mut candidates = lock(&self.routing).closest(key.as_bytes(), K); + let mut known: HashSet<_> = candidates.iter().map(|contact| contact.peer_id).collect(); + let mut queried = HashSet::new(); + let mut sent = 0usize; + let deadline = tokio::time::Instant::now() + LOOKUP_TIMEOUT; + while tokio::time::Instant::now() < deadline && sent < MAX_LOOKUP_REQUESTS { + candidates.sort_by_key(|contact| distance(contact.node_id.as_bytes(), key.as_bytes())); + let batch: Vec<_> = candidates + .iter() + .filter(|contact| !queried.contains(&contact.peer_id)) + .take(ALPHA.min(MAX_LOOKUP_REQUESTS - sent)) + .cloned() + .collect(); + if batch.is_empty() { + break; + } + sent += batch.len(); + queried.extend(batch.iter().map(|contact| contact.peer_id)); + let mut futures: FuturesUnordered<_> = batch + .iter() + .map(|contact| async move { + ( + contact, + tokio::time::timeout( + REQUEST_TIMEOUT, + self.exchange(contact, RoutingOperation::FindValue { key }), + ) + .await, + ) + }) + .collect(); + let mut found_in_round = false; + while let Ok(Some((_, response))) = + tokio::time::timeout_at(deadline, futures.next()).await + { + let Ok(Ok(response)) = response else { + continue; + }; + if !response.records.is_empty() { + found_in_round = true; + } + merge_records(&mut records, response.records, &self.network_id, key); + for node in response.nodes { + if validate_contact(&node, &self.network_id, Some(self.own_contact.peer_id)) + .is_ok() + && known.insert(node.peer_id) + { + self.learn_contact(node.clone()).await?; + candidates.push(node); + } + } + } + if found_in_round { + break; + } + } + Ok(records.into_values().collect()) + } + + async fn exchange( + &self, + contact: &NodeContact, + operation: RoutingOperation, + ) -> Result { + validate_contact(contact, &self.network_id, Some(self.own_contact.peer_id))?; + let ticket = PeerTicket::from_str(&contact.ticket) + .map_err(|error| MusicDhtError::InvalidTicket(error.to_string()))?; + let mut stream = self + .service + .open_stream_to(&ticket, SIMILARITY_DHT_ALPN) + .await?; + let request = RoutingRequest { + version: SIMILARITY_DHT_PROTOCOL_VERSION, + requester: self.fresh_own_contact(), + operation, + }; + write_postcard(&mut stream.send, &request, MAX_SIMILARITY_DHT_REQUEST_BYTES).await?; + stream.send.finish().map_err(network_error)?; + let response: RoutingResponse = + read_postcard(&mut stream.recv, MAX_SIMILARITY_DHT_RESPONSE_BYTES).await?; + validate_response(&response)?; + if !response.ok { + return Err(protocol_error( + response + .error + .as_deref() + .unwrap_or("routing request refused"), + )); + } + Ok(response) + } + + fn replica_targets(&self, key: &SimilarityDhtKey) -> (bool, Vec) { + let mut contacts = lock(&self.routing).contacts(); + contacts.push(self.fresh_own_contact()); + contacts.sort_by_key(|contact| distance(contact.node_id.as_bytes(), key.as_bytes())); + contacts.dedup_by_key(|contact| contact.peer_id); + contacts.truncate(K); + let local = contacts + .iter() + .any(|contact| contact.peer_id == self.own_contact.peer_id); + contacts.retain(|contact| contact.peer_id != self.own_contact.peer_id); + (local, contacts) + } + + fn closest_contacts(&self, target: &[u8; 32], count: usize) -> Vec { + let mut contacts = lock(&self.routing).contacts(); + contacts.push(self.fresh_own_contact()); + contacts.sort_by_key(|contact| distance(contact.node_id.as_bytes(), target)); + contacts.dedup_by_key(|contact| contact.peer_id); + contacts.truncate(count); + contacts + } + + fn fresh_own_contact(&self) -> NodeContact { + let mut contact = self.own_contact.clone(); + contact.last_seen_ms = now_ms(); + contact + } + + async fn learn_response_nodes(&self, nodes: Vec) { + for contact in nodes { + if let Err(error) = self.learn_contact(contact).await { + debug!(%error, "invalid similarity DHT contact ignored"); + } + } + } + + async fn learn_contact(&self, mut contact: NodeContact) -> Result<()> { + validate_contact(&contact, &self.network_id, Some(self.own_contact.peer_id))?; + contact.node_id = NodeId::from_endpoint(&contact.peer_id); + // A remote peer does not get to extend its own lifetime by supplying + // an arbitrary future timestamp. Successful authenticated contact is + // the observation that matters locally. + contact.last_seen_ms = now_ms(); + let is_new = lock(&self.routing).upsert(contact.clone()); + self.database.upsert_peer(contact.clone()).await?; + if is_new { + debug!(peer = %contact.peer_id, "learned similarity DHT peer"); + } + Ok(()) + } +} + +fn validate_record( + record: &SignedSimilarityRoute, + network_id: &NetworkId, + at_ms: u64, +) -> Result<()> { + record.verify(network_id)?; + let future_limit = at_ms.saturating_add(MAX_CLOCK_SKEW.as_millis() as u64); + let expiry = record + .payload + .issued_at_ms + .saturating_add(ROUTING_RECORD_TTL.as_millis() as u64); + if record.payload.issued_at_ms > future_limit || expiry <= at_ms { + return Err(protocol_error("invalid similarity DHT record time")); + } + Ok(()) +} + +fn validate_contact( + contact: &NodeContact, + network_id: &NetworkId, + excluded: Option, +) -> Result<()> { + if excluded == Some(contact.peer_id) + || contact.node_id != NodeId::from_endpoint(&contact.peer_id) + || contact.ticket.len() > 16 * 1024 + { + return Err(protocol_error("invalid similarity DHT contact")); + } + let ticket = PeerTicket::from_str(&contact.ticket) + .map_err(|error| MusicDhtError::InvalidTicket(error.to_string()))?; + if ticket.endpoint_id() != contact.peer_id || &ticket.network_id != network_id { + return Err(protocol_error("similarity DHT contact ticket mismatch")); + } + Ok(()) +} + +fn validate_response(response: &RoutingResponse) -> Result<()> { + if response.version != SIMILARITY_DHT_PROTOCOL_VERSION + || response.nodes.len() > K + || response.records.len() > MAX_SIMILARITY_DHT_RECORDS + || response + .error + .as_ref() + .is_some_and(|error| error.len() > 1024) + || (response.ok && response.error.is_some()) + || (!response.ok && response.error.is_none()) + { + return Err(protocol_error("invalid similarity DHT response")); + } + Ok(()) +} + +fn merge_records( + destination: &mut HashMap, + records: Vec, + network_id: &NetworkId, + expected_key: SimilarityDhtKey, +) { + let at_ms = now_ms(); + for record in records { + if validate_record(&record, network_id, at_ms).is_err() + || record.key().ok() != Some(expected_key) + { + continue; + } + let owner = record.payload.owner; + match destination.get(&owner) { + Some(existing) if existing.payload.issued_at_ms > record.payload.issued_at_ms => {} + _ => { + destination.insert(owner, record); + } + } + } +} + +fn chunk_store_records(records: Vec) -> Vec> { + let mut chunks = Vec::new(); + let mut current = Vec::new(); + let mut current_bytes = 0usize; + for record in records { + let estimate = RECORD_FIXED_BYTES_ESTIMATE + + record.payload.owner_ticket.len() + + record.payload.entries.len() * ENTRY_BYTES_ESTIMATE; + if !current.is_empty() + && (current.len() >= MAX_SIMILARITY_STORE_RECORDS + || current_bytes.saturating_add(estimate) > MAX_BATCH_BYTES_ESTIMATE) + { + chunks.push(std::mem::take(&mut current)); + current_bytes = 0; + } + current_bytes = current_bytes.saturating_add(estimate); + current.push(record); + } + if !current.is_empty() { + chunks.push(current); + } + chunks +} + +async fn write_postcard( + writer: &mut W, + value: &T, + maximum: usize, +) -> Result<()> { + let payload = postcard::to_stdvec(value).map_err(protocol_error)?; + if payload.len() > maximum { + return Err(protocol_error("similarity DHT message is too large")); + } + writer.write_all(&payload).await.map_err(network_error) +} + +async fn read_postcard Deserialize<'de>>( + reader: &mut R, + maximum: usize, +) -> Result { + let mut payload = Vec::new(); + reader + .take(maximum as u64 + 1) + .read_to_end(&mut payload) + .await + .map_err(network_error)?; + if payload.is_empty() || payload.len() > maximum { + return Err(protocol_error("invalid similarity DHT message size")); + } + postcard::from_bytes(&payload).map_err(protocol_error) +} + +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(PoisonError::into_inner) +} + +fn read(lock: &RwLock) -> std::sync::RwLockReadGuard<'_, T> { + lock.read().unwrap_or_else(PoisonError::into_inner) +} + +fn write(lock: &RwLock) -> std::sync::RwLockWriteGuard<'_, T> { + lock.write().unwrap_or_else(PoisonError::into_inner) +} + +fn protocol_error(error: impl std::fmt::Display) -> MusicDhtError { + MusicDhtError::Protocol(error.to_string()) +} + +fn network_error(error: impl std::fmt::Display) -> MusicDhtError { + MusicDhtError::Network(error.to_string()) +} + +#[cfg(test)] +mod tests { + use federation_net::{ + EndpointAddr, PROTOCOL_VERSION, PeerTicket, SchemaId, SecretKey, TICKET_VERSION, + }; + + use super::*; + use crate::SimilarityRouteEntry; + use crate::similarity_lsh::{SimilarityRoutePayload, build_route_specs}; + + fn ticket(network: NetworkId, owner: EndpointId) -> String { + PeerTicket { + ticket_version: TICKET_VERSION, + protocol_version: PROTOCOL_VERSION, + network_id: network, + schema_id: SchemaId::from_name("similarity-test"), + endpoint_addr: EndpointAddr::new(owner), + } + .to_string() + } + + #[test] + fn store_batches_obey_count_and_estimated_byte_bounds() { + let network = NetworkId::from_name("batch"); + let owner = SecretKey::from_bytes(&[8; 32]); + let spec = SimilarityRouteSpec { + profile_id: "sim1:test".into(), + table: 0, + primary_bucket: 0, + entries: (0..512) + .map(|suffix| SimilarityRouteEntry { + suffix, + representative: [suffix as u8; 32], + }) + .collect(), + }; + let records: Vec<_> = (0..100) + .map(|issued| { + let payload = SimilarityRoutePayload::from_spec( + spec.clone(), + network, + owner.public(), + ticket(network, owner.public()), + issued + 1, + ) + .unwrap(); + SignedSimilarityRoute::sign(payload, |bytes| owner.sign(bytes)).unwrap() + }) + .collect(); + let chunks = chunk_store_records(records); + assert!(chunks.len() > 1); + assert!( + chunks + .iter() + .all(|chunk| chunk.len() <= MAX_SIMILARITY_STORE_RECORDS) + ); + } + + #[tokio::test] + async fn routing_database_keeps_newest_signed_owner_record() { + let temp = tempfile::tempdir().unwrap(); + let database = RoutingDatabase::open(&temp.path().join("routing.sqlite3")) + .await + .unwrap(); + let network = NetworkId::from_name("db"); + let owner = SecretKey::from_bytes(&[9; 32]); + let spec = build_route_specs("sim1:test", &[[0x11; 32]]) + .unwrap() + .remove(0); + let make = |issued_at_ms| { + let payload = SimilarityRoutePayload::from_spec( + spec.clone(), + network, + owner.public(), + ticket(network, owner.public()), + issued_at_ms, + ) + .unwrap(); + SignedSimilarityRoute::sign(payload, |bytes| owner.sign(bytes)).unwrap() + }; + let newer = make(now_ms()); + let older = make(now_ms().saturating_sub(1)); + let key = newer.key().unwrap(); + database.store_records(vec![newer.clone()]).await.unwrap(); + database.store_records(vec![older]).await.unwrap(); + assert_eq!(database.records_by_key(key, 1).await.unwrap(), vec![newer]); + } +} diff --git a/crates/music-dht/src/similarity_lsh.rs b/crates/music-dht/src/similarity_lsh.rs new file mode 100644 index 0000000..dfd8fc0 --- /dev/null +++ b/crates/music-dht/src/similarity_lsh.rs @@ -0,0 +1,665 @@ +//! Deterministic, model-neutral LSH routing for federated similarity search. +//! +//! Applications still own embedding generation and exact local search. This +//! module turns normalized embeddings into compact routing signatures and +//! groups them into bounded, signed peer summaries suitable for a DHT. + +use std::collections::{BTreeMap, HashSet}; +use std::sync::OnceLock; +use std::time::Duration; + +use federation_net::{EndpointId, NetworkId, PeerTicket, Signature}; +use serde::{Deserialize, Serialize}; + +use crate::error::{MusicDhtError, Result}; + +/// ALPN of the schema-independent similarity-routing overlay. +pub const SIMILARITY_DHT_ALPN: &[u8] = b"furumi-fd/similarity-dht/1"; +/// Current similarity-routing wire and record version. +pub const SIMILARITY_DHT_PROTOCOL_VERSION: u16 = 1; +/// Stable identifier advertised in capability manifests. +pub const SIMILARITY_DHT_ID: &str = "similarity_dht"; +/// Size of the routing SimHash. It is separate from the shorter result +/// signature used only for near-duplicate filtering. +pub const ROUTING_SIGNATURE_BYTES: usize = 32; +/// Number of independent LSH tables published by every peer. +pub const ROUTING_TABLES: usize = 12; +/// Bits in the fine bucket of one table. +pub const ROUTING_BITS: usize = 20; +/// Prefix bits forming the DHT key. The remaining bits are carried compactly +/// inside the peer's signed record. +pub const ROUTING_PRIMARY_BITS: usize = 10; +/// Fine-bucket suffix bits stored inside a primary-bucket record. +pub const ROUTING_SUFFIX_BITS: usize = ROUTING_BITS - ROUTING_PRIMARY_BITS; +/// Query-side primary buckets per table: exact plus the lowest-margin +/// one-bit neighbor. +pub const ROUTING_PRIMARY_PROBES: usize = 2; +/// Query-side fine suffixes checked inside a record: exact plus eight +/// lowest-margin one-bit neighbors. +pub const ROUTING_SUFFIX_PROBES: usize = 9; +/// Maximum profile fingerprint length accepted on the routing wire. +pub const MAX_ROUTING_PROFILE_BYTES: usize = 128; +/// Maximum self-contained owner ticket accepted in a routing record. +pub const MAX_ROUTING_TICKET_BYTES: usize = 16 * 1024; +/// Maximum fine buckets in one peer summary. Ten suffix bits make this a +/// natural hard bound independent of library size. +pub const MAX_ROUTING_ENTRIES: usize = 1 << ROUTING_SUFFIX_BITS; +/// Routing records expire together with ordinary active library records. +pub const ROUTING_RECORD_TTL: Duration = crate::ACTIVE_RECORD_TTL; + +const ROUTING_SIGNATURE_BITS: usize = ROUTING_SIGNATURE_BYTES * 8; +const SIGNATURE_DOMAIN: &[u8] = b"frid-similarity-simhash-v1"; +const TABLE_DOMAIN: &[u8] = b"frid-similarity-lsh-table-v1\0"; +const KEY_DOMAIN: &[u8] = b"frid-similarity-lsh-key-v1\0"; +const RECORD_DOMAIN: &[u8] = b"frid-similarity-lsh-record-v1\0"; + +/// A 256-bit point in the similarity-routing DHT key space. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)] +pub struct SimilarityDhtKey([u8; 32]); + +impl SimilarityDhtKey { + /// Derives a network-, profile-, table- and primary-bucket-specific key. + pub fn derive( + network_id: &NetworkId, + profile_id: &str, + table: u8, + primary_bucket: u16, + ) -> Result { + validate_profile(profile_id)?; + if table as usize >= ROUTING_TABLES + || primary_bucket as usize >= (1 << ROUTING_PRIMARY_BITS) + { + return Err(protocol_error("invalid similarity DHT bucket")); + } + let mut hasher = blake3::Hasher::new(); + hasher.update(KEY_DOMAIN); + hasher.update(network_id.as_bytes()); + hasher.update(&(profile_id.len() as u16).to_le_bytes()); + hasher.update(profile_id.as_bytes()); + hasher.update(&[table]); + hasher.update(&primary_bucket.to_le_bytes()); + Ok(Self(*hasher.finalize().as_bytes())) + } + + /// Creates a key from raw bytes received on the wire. + pub fn from_bytes(bytes: [u8; 32]) -> Self { + Self(bytes) + } + + /// Returns the raw DHT key bytes. + pub fn as_bytes(&self) -> &[u8; 32] { + &self.0 + } +} + +/// Query-side routing projection. Margins are intentionally local-only: they +/// choose useful neighboring probes but are never sent over the network. +#[derive(Clone)] +pub struct SimilarityRoutingQuery { + signature: [u8; ROUTING_SIGNATURE_BYTES], + margins: [f32; ROUTING_SIGNATURE_BITS], +} + +impl std::fmt::Debug for SimilarityRoutingQuery { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("SimilarityRoutingQuery") + .field("signature", &self.signature) + .finish_non_exhaustive() + } +} + +impl SimilarityRoutingQuery { + /// Compact routing signature of the normalized query embedding. + pub fn signature(&self) -> &[u8; ROUTING_SIGNATURE_BYTES] { + &self.signature + } + + /// Returns the two coarse DHT buckets and nine fine suffixes to probe for + /// one table. + pub fn probes(&self, profile_id: &str, table: u8) -> Result { + let positions = table_bit_positions(profile_id, table)?; + let primary_positions = &positions[..ROUTING_PRIMARY_BITS]; + let suffix_positions = &positions[ROUTING_PRIMARY_BITS..]; + let primary = extract_bits(&self.signature, primary_positions); + let suffix = extract_bits(&self.signature, suffix_positions); + + let mut primary_by_margin: Vec<(usize, f32)> = primary_positions + .iter() + .enumerate() + .map(|(bucket_bit, signature_bit)| (bucket_bit, self.margins[*signature_bit as usize])) + .collect(); + primary_by_margin.sort_by(|left, right| left.1.total_cmp(&right.1)); + let mut primary_buckets = vec![primary]; + for (bit, _) in primary_by_margin + .into_iter() + .take(ROUTING_PRIMARY_PROBES.saturating_sub(1)) + { + primary_buckets.push(primary ^ (1 << bit)); + } + + let mut suffix_by_margin: Vec<(usize, f32)> = suffix_positions + .iter() + .enumerate() + .map(|(bucket_bit, signature_bit)| (bucket_bit, self.margins[*signature_bit as usize])) + .collect(); + suffix_by_margin.sort_by(|left, right| left.1.total_cmp(&right.1)); + let mut suffix_buckets = vec![suffix]; + for (bit, _) in suffix_by_margin + .into_iter() + .take(ROUTING_SUFFIX_PROBES.saturating_sub(1)) + { + suffix_buckets.push(suffix ^ (1 << bit)); + } + + Ok(SimilarityTableProbes { + table, + primary_buckets, + suffix_buckets, + }) + } +} + +/// Bounded query probes for one LSH table. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SimilarityTableProbes { + /// LSH table number. + pub table: u8, + /// Exact and one low-margin neighboring DHT bucket. + pub primary_buckets: Vec, + /// Exact and eight low-margin fine buckets checked locally. + pub suffix_buckets: Vec, +} + +/// One fine bucket and its anonymous representative routing signature. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct SimilarityRouteEntry { + /// Ten-bit suffix inside the record's primary bucket. + pub suffix: u16, + /// A stable representative of tracks in this peer's fine bucket. + pub representative: [u8; ROUTING_SIGNATURE_BYTES], +} + +/// Unsigned local summary built from a peer's embeddings. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SimilarityRouteSpec { + /// Exact model/preprocessing compatibility fingerprint. + pub profile_id: String, + /// LSH table number. + pub table: u8, + /// Coarse ten-bit DHT bucket. + pub primary_bucket: u16, + /// Sorted, unique fine-bucket representatives. + pub entries: Vec, +} + +impl SimilarityRouteSpec { + /// Derives this summary's DHT key for `network_id`. + pub fn key(&self, network_id: &NetworkId) -> Result { + SimilarityDhtKey::derive( + network_id, + &self.profile_id, + self.table, + self.primary_bucket, + ) + } +} + +/// Signed immutable payload owned by one peer. Forwarders cannot change its +/// profile, buckets, representatives or issue time without invalidating the +/// Ed25519 identity signature. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SimilarityRoutePayload { + /// Record format version. + pub version: u16, + /// Network scope; prevents replay into another federation. + pub network_id: NetworkId, + /// Peer that owns and signed the summary. + pub owner: EndpointId, + /// Self-contained route to the owner. This is required because a lookup + /// may discover a peer that is absent from the caller's routing table. + pub owner_ticket: String, + /// Exact model/preprocessing compatibility fingerprint. + pub profile_id: String, + /// LSH table number. + pub table: u8, + /// Coarse ten-bit DHT bucket. + pub primary_bucket: u16, + /// Sorted, unique fine-bucket representatives. + pub entries: Vec, + /// Unix time at which the owner issued this version. + pub issued_at_ms: u64, +} + +impl SimilarityRoutePayload { + /// Builds a payload from a validated local summary. + pub fn from_spec( + spec: SimilarityRouteSpec, + network_id: NetworkId, + owner: EndpointId, + owner_ticket: String, + issued_at_ms: u64, + ) -> Result { + let payload = Self { + version: SIMILARITY_DHT_PROTOCOL_VERSION, + network_id, + owner, + owner_ticket, + profile_id: spec.profile_id, + table: spec.table, + primary_bucket: spec.primary_bucket, + entries: spec.entries, + issued_at_ms, + }; + payload.validate(&payload.network_id)?; + Ok(payload) + } + + /// Derives the DHT key this payload must be stored under. + pub fn key(&self) -> Result { + SimilarityDhtKey::derive( + &self.network_id, + &self.profile_id, + self.table, + self.primary_bucket, + ) + } + + /// Validates all untrusted bounds and the expected network scope. + pub fn validate(&self, expected_network: &NetworkId) -> Result<()> { + if self.version != SIMILARITY_DHT_PROTOCOL_VERSION { + return Err(protocol_error( + "unsupported similarity route record version", + )); + } + if &self.network_id != expected_network { + return Err(protocol_error( + "similarity route belongs to another network", + )); + } + if self.owner_ticket.len() > MAX_ROUTING_TICKET_BYTES { + return Err(protocol_error("invalid similarity route owner ticket")); + } + let ticket = self + .owner_ticket + .parse::() + .map_err(|_| protocol_error("invalid similarity route owner ticket"))?; + if ticket.endpoint_id() != self.owner || &ticket.network_id != expected_network { + return Err(protocol_error("similarity route owner ticket mismatch")); + } + validate_profile(&self.profile_id)?; + if self.table as usize >= ROUTING_TABLES + || self.primary_bucket as usize >= (1 << ROUTING_PRIMARY_BITS) + || self.entries.is_empty() + || self.entries.len() > MAX_ROUTING_ENTRIES + || self.issued_at_ms == 0 + { + return Err(protocol_error("invalid similarity route bounds")); + } + let mut previous = None; + for entry in &self.entries { + if entry.suffix as usize >= (1 << ROUTING_SUFFIX_BITS) + || previous.is_some_and(|value| value >= entry.suffix) + { + return Err(protocol_error("invalid similarity route entries")); + } + previous = Some(entry.suffix); + } + Ok(()) + } +} + +/// End-to-end signed routing summary. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SignedSimilarityRoute { + /// Immutable owner payload. + pub payload: SimilarityRoutePayload, + /// Ed25519 signature made by `payload.owner`. + pub signature: Signature, +} + +impl SignedSimilarityRoute { + /// Signs a validated payload using an identity-backed signer. + pub fn sign( + payload: SimilarityRoutePayload, + signer: impl FnOnce(&[u8]) -> Signature, + ) -> Result { + payload.validate(&payload.network_id)?; + let bytes = signing_bytes(&payload)?; + Ok(Self { + payload, + signature: signer(&bytes), + }) + } + + /// Verifies payload bounds, network scope and the owner's signature. + pub fn verify(&self, expected_network: &NetworkId) -> Result<()> { + self.payload.validate(expected_network)?; + let bytes = signing_bytes(&self.payload)?; + self.payload + .owner + .verify(&bytes, &self.signature) + .map_err(|_| protocol_error("invalid similarity route signature")) + } + + /// DHT key under which this record must be stored. + pub fn key(&self) -> Result { + self.payload.key() + } +} + +/// Computes the stable 256-bit routing signature of a normalized embedding. +pub fn routing_signature(vector: &[f32]) -> Result<[u8; ROUTING_SIGNATURE_BYTES]> { + Ok(routing_query(vector)?.signature) +} + +/// Computes a routing signature and local projection margins for multi-probe +/// lookup. No model-specific constants or calibration data are used. +pub fn routing_query(vector: &[f32]) -> Result { + validate_vector(vector)?; + let signs = hyperplane_signs(); + let mut projections = [0.0f32; ROUTING_SIGNATURE_BITS]; + for (dimension, value) in vector.iter().copied().enumerate() { + for (projection, sign) in projections.iter_mut().zip(&signs[dimension]) { + *projection += value * *sign; + } + } + let mut signature = [0u8; ROUTING_SIGNATURE_BYTES]; + let mut margins = [0.0f32; ROUTING_SIGNATURE_BITS]; + for (bit, projection) in projections.into_iter().enumerate() { + if projection >= 0.0 { + signature[bit / 8] |= 1 << (bit % 8); + } + margins[bit] = projection.abs(); + } + Ok(SimilarityRoutingQuery { signature, margins }) +} + +/// Builds one bounded peer summary per occupied `(table, primary bucket)`. +/// Multiple tracks in the same fine bucket collapse to one deterministic +/// representative, so a large library cannot flood a hot DHT key. +pub fn build_route_specs( + profile_id: &str, + signatures: &[[u8; ROUTING_SIGNATURE_BYTES]], +) -> Result> { + validate_profile(profile_id)?; + if signatures.is_empty() { + return Ok(Vec::new()); + } + let positions = (0..ROUTING_TABLES) + .map(|table| table_bit_positions(profile_id, table as u8)) + .collect::>>()?; + let mut groups: BTreeMap<(u8, u16), BTreeMap> = + BTreeMap::new(); + for signature in signatures { + for (table, table_positions) in positions.iter().enumerate() { + let primary = extract_bits(signature, &table_positions[..ROUTING_PRIMARY_BITS]); + let suffix = extract_bits(signature, &table_positions[ROUTING_PRIMARY_BITS..]); + groups + .entry((table as u8, primary)) + .or_default() + .entry(suffix) + .and_modify(|representative| { + if signature < representative { + *representative = *signature; + } + }) + .or_insert(*signature); + } + } + Ok(groups + .into_iter() + .map(|((table, primary_bucket), entries)| SimilarityRouteSpec { + profile_id: profile_id.to_string(), + table, + primary_bucket, + entries: entries + .into_iter() + .map(|(suffix, representative)| SimilarityRouteEntry { + suffix, + representative, + }) + .collect(), + }) + .collect()) +} + +/// Hamming distance between two routing signatures. +pub fn routing_signature_distance( + left: &[u8; ROUTING_SIGNATURE_BYTES], + right: &[u8; ROUTING_SIGNATURE_BYTES], +) -> u32 { + left.iter() + .zip(right) + .map(|(left, right)| (left ^ right).count_ones()) + .sum() +} + +fn hyperplane_signs() -> &'static Vec<[f32; ROUTING_SIGNATURE_BITS]> { + static SIGNS: OnceLock> = OnceLock::new(); + SIGNS.get_or_init(|| { + (0..crate::similarity::MAX_SIMILARITY_DIMENSIONS) + .map(|dimension| { + let mut hasher = blake3::Hasher::new(); + hasher.update(SIGNATURE_DOMAIN); + hasher.update(&(dimension as u64).to_le_bytes()); + let digest = hasher.finalize(); + std::array::from_fn(|bit| { + if digest.as_bytes()[bit / 8] & (1 << (bit % 8)) != 0 { + 1.0 + } else { + -1.0 + } + }) + }) + .collect() + }) +} + +fn table_bit_positions(profile_id: &str, table: u8) -> Result<[u8; ROUTING_BITS]> { + validate_profile(profile_id)?; + if table as usize >= ROUTING_TABLES { + return Err(protocol_error("invalid similarity routing table")); + } + let mut selected = Vec::with_capacity(ROUTING_BITS); + let mut seen = HashSet::with_capacity(ROUTING_BITS); + let mut counter = 0u32; + while selected.len() < ROUTING_BITS { + let mut hasher = blake3::Hasher::new(); + hasher.update(TABLE_DOMAIN); + hasher.update(profile_id.as_bytes()); + hasher.update(&(table as u16).to_le_bytes()); + hasher.update(&counter.to_le_bytes()); + counter += 1; + for candidate in hasher.finalize().as_bytes() { + if seen.insert(*candidate) { + selected.push(*candidate); + if selected.len() == ROUTING_BITS { + break; + } + } + } + } + selected + .try_into() + .map_err(|_| protocol_error("failed to derive similarity routing table")) +} + +fn extract_bits(signature: &[u8; ROUTING_SIGNATURE_BYTES], positions: &[u8]) -> u16 { + positions + .iter() + .enumerate() + .fold(0u16, |value, (bucket_bit, signature_bit)| { + let bit = (signature[*signature_bit as usize / 8] >> (*signature_bit as usize % 8)) & 1; + value | (u16::from(bit) << bucket_bit) + }) +} + +fn signing_bytes(payload: &SimilarityRoutePayload) -> Result> { + let encoded = postcard::to_stdvec(payload).map_err(protocol_error)?; + let mut bytes = Vec::with_capacity(RECORD_DOMAIN.len() + encoded.len()); + bytes.extend_from_slice(RECORD_DOMAIN); + bytes.extend_from_slice(&encoded); + Ok(bytes) +} + +fn validate_profile(profile_id: &str) -> Result<()> { + if profile_id.is_empty() || profile_id.len() > MAX_ROUTING_PROFILE_BYTES { + Err(protocol_error("invalid similarity routing profile id")) + } else { + Ok(()) + } +} + +fn validate_vector(vector: &[f32]) -> Result<()> { + if vector.is_empty() + || vector.len() > crate::similarity::MAX_SIMILARITY_DIMENSIONS + || !vector.iter().all(|value| value.is_finite()) + { + return Err(protocol_error("invalid vector for similarity routing")); + } + let norm = vector.iter().map(|value| value * value).sum::().sqrt(); + if !norm.is_finite() || (norm - 1.0).abs() > 0.05 { + return Err(protocol_error( + "similarity routing vector is not L2-normalized", + )); + } + Ok(()) +} + +fn protocol_error(error: impl std::fmt::Display) -> MusicDhtError { + MusicDhtError::Protocol(error.to_string()) +} + +#[cfg(test)] +mod tests { + use federation_net::{ + EndpointAddr, PROTOCOL_VERSION, PeerTicket, SchemaId, SecretKey, TICKET_VERSION, + }; + + use super::*; + + fn normalized(values: &mut [f32]) { + let norm = values.iter().map(|value| value * value).sum::().sqrt(); + for value in values { + *value /= norm; + } + } + + fn ticket(network: NetworkId, owner: EndpointId) -> String { + PeerTicket { + ticket_version: TICKET_VERSION, + protocol_version: PROTOCOL_VERSION, + network_id: network, + schema_id: SchemaId::from_name("similarity-test"), + endpoint_addr: EndpointAddr::new(owner), + } + .to_string() + } + + #[test] + fn routing_signature_is_stable_and_extends_result_simhash() { + let mut vector = vec![0.0; 64]; + for (index, value) in vector.iter_mut().enumerate() { + *value = index as f32 - 31.5; + } + normalized(&mut vector); + let routing = routing_signature(&vector).unwrap(); + let compact = crate::similarity::embedding_signature(&vector).unwrap(); + assert_eq!(&routing[..compact.len()], &compact); + assert_eq!(routing, routing_signature(&vector).unwrap()); + } + + #[test] + fn route_specs_are_bounded_sorted_and_deterministic() { + let signatures = [[0u8; 32], [0xff; 32], [0u8; 32]]; + let first = build_route_specs("sim1:test", &signatures).unwrap(); + let second = build_route_specs("sim1:test", &signatures).unwrap(); + assert_eq!(first, second); + assert!(!first.is_empty()); + assert!(first.len() <= ROUTING_TABLES * signatures.len()); + assert!(first.iter().all(|spec| { + !spec.entries.is_empty() + && spec.entries.len() <= MAX_ROUTING_ENTRIES + && spec + .entries + .windows(2) + .all(|pair| pair[0].suffix < pair[1].suffix) + })); + } + + #[test] + fn signed_route_cannot_be_forged_or_replayed_between_networks() { + let network = NetworkId::from_name("a"); + let other = NetworkId::from_name("b"); + let key = SecretKey::from_bytes(&[7; 32]); + let spec = build_route_specs("sim1:test", &[[0x55; 32]]) + .unwrap() + .remove(0); + let payload = SimilarityRoutePayload::from_spec( + spec, + network, + key.public(), + ticket(network, key.public()), + 42, + ) + .unwrap(); + let signed = SignedSimilarityRoute::sign(payload, |bytes| key.sign(bytes)).unwrap(); + signed.verify(&network).unwrap(); + assert!(signed.verify(&other).is_err()); + + let mut forged = signed.clone(); + forged.payload.entries[0].representative[0] ^= 1; + assert!(forged.verify(&network).is_err()); + } + + #[test] + fn route_payload_rejects_unbounded_wire_fields() { + let network = NetworkId::from_name("bounds"); + let owner = SecretKey::from_bytes(&[6; 32]).public(); + let spec = build_route_specs("sim1:test", &[[0x33; 32]]) + .unwrap() + .remove(0); + let payload = + SimilarityRoutePayload::from_spec(spec, network, owner, ticket(network, owner), 1) + .unwrap(); + + let mut oversized_profile = payload.clone(); + oversized_profile.profile_id = "x".repeat(MAX_ROUTING_PROFILE_BYTES + 1); + assert!(oversized_profile.validate(&network).is_err()); + + let mut oversized_entries = payload; + oversized_entries.entries = (0..=MAX_ROUTING_ENTRIES) + .map(|suffix| SimilarityRouteEntry { + suffix: suffix as u16, + representative: [0; ROUTING_SIGNATURE_BYTES], + }) + .collect(); + assert!(oversized_entries.validate(&network).is_err()); + + let mut mismatched_ticket = oversized_profile; + mismatched_ticket.profile_id = "sim1:test".into(); + mismatched_ticket.owner_ticket = ticket(network, SecretKey::from_bytes(&[5; 32]).public()); + assert!(mismatched_ticket.validate(&network).is_err()); + } + + #[test] + fn query_probes_are_bounded_and_include_exact_bucket() { + let query = routing_query(&[0.5; 4]).unwrap(); + let probes = query.probes("sim1:test", 0).unwrap(); + assert_eq!(probes.primary_buckets.len(), ROUTING_PRIMARY_PROBES); + assert_eq!(probes.suffix_buckets.len(), ROUTING_SUFFIX_PROBES); + assert!( + probes + .primary_buckets + .iter() + .all(|bucket| *bucket < (1 << ROUTING_PRIMARY_BITS)) + ); + assert!( + probes + .suffix_buckets + .iter() + .all(|bucket| *bucket < (1 << ROUTING_SUFFIX_BITS)) + ); + } +} diff --git a/crates/music-dht/tests/similarity_dht.rs b/crates/music-dht/tests/similarity_dht.rs new file mode 100644 index 0000000..fee1f9e --- /dev/null +++ b/crates/music-dht/tests/similarity_dht.rs @@ -0,0 +1,108 @@ +use std::sync::Arc; +use std::time::Duration; + +use music_dht::similarity_dht::SimilarityDht; +use music_dht::similarity_lsh::{SIMILARITY_DHT_ALPN, routing_signature}; +use music_dht::{MusicDhtConfig, MusicDhtService, NetworkId}; + +const TEST_DIRECT_ALPN: &[u8] = b"music-dht-test/similarity-owner/1"; + +async fn start_node( + directory: &std::path::Path, + network: NetworkId, +) -> (Arc, tokio::task::JoinHandle<()>) { + std::fs::create_dir_all(directory).unwrap(); + let config = MusicDhtConfig::builder() + .data_dir(directory) + .network_id(network) + .schema_independent_stream_protocol(SIMILARITY_DHT_ALPN) + .schema_independent_stream_protocol(TEST_DIRECT_ALPN) + .request_timeout(Duration::from_secs(2)) + .lookup_timeout(Duration::from_secs(5)) + .transport_timeout(Duration::from_secs(5)) + .dial_timeout(Duration::from_secs(2)) + .build() + .unwrap(); + let (service, mut events) = MusicDhtService::start(config).await.unwrap(); + let task = tokio::spawn(async move { while events.recv().await.is_some() {} }); + (Arc::new(service), task) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn signed_lsh_summary_routes_a_query_to_another_peer() { + let temp = tempfile::tempdir().unwrap(); + let network = NetworkId::from_name("similarity-dht-integration"); + let (first, first_events) = start_node(&temp.path().join("first"), network).await; + let (second, second_events) = start_node(&temp.path().join("second"), network).await; + + let first_acceptor = first.stream_acceptor(SIMILARITY_DHT_ALPN).unwrap(); + let second_acceptor = second.stream_acceptor(SIMILARITY_DHT_ALPN).unwrap(); + let mut second_direct_acceptor = second.stream_acceptor(TEST_DIRECT_ALPN).unwrap(); + let first_routing = SimilarityDht::open( + Arc::clone(&first), + temp.path().join("first-routing.sqlite3"), + ) + .await + .unwrap(); + let second_routing = SimilarityDht::open( + Arc::clone(&second), + temp.path().join("second-routing.sqlite3"), + ) + .await + .unwrap(); + let first_serve = tokio::spawn(Arc::clone(&first_routing).serve(first_acceptor)); + let second_serve = tokio::spawn(Arc::clone(&second_routing).serve(second_acceptor)); + + second.connect(first.ticket().await.unwrap()).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + while first.known_peers().is_empty() || second.known_peers().is_empty() { + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await + .unwrap(); + + let vector = vec![0.5; 4]; + let profile = "sim1:integration"; + let stats = second_routing + .sync_local_signatures( + profile.to_string(), + vec![routing_signature(&vector).unwrap()], + ) + .await + .unwrap(); + assert_eq!(stats.records, 12); + assert!(stats.local_replica); + assert_eq!(stats.remote_nodes, 1); + + let peers = tokio::time::timeout( + Duration::from_secs(10), + first_routing.find_peers(profile, &vector, 16), + ) + .await + .unwrap() + .unwrap(); + assert_eq!( + peers.first().map(|ticket| ticket.endpoint_id()), + Some(second.endpoint_id()) + ); + let mut outbound = first + .open_stream_to(peers.first().unwrap(), TEST_DIRECT_ALPN) + .await + .unwrap(); + let inbound = tokio::time::timeout(Duration::from_secs(5), second_direct_acceptor.accept()) + .await + .unwrap() + .unwrap(); + assert_eq!(inbound.peer_id, first.endpoint_id()); + outbound.send.finish().unwrap(); + drop(inbound); + drop(outbound); + + first_serve.abort(); + second_serve.abort(); + first.shutdown().await.unwrap(); + second.shutdown().await.unwrap(); + first_events.abort(); + second_events.abort(); +}