use std::str::FromStr; use async_trait::async_trait; use music_dht::{ DhtKey, EndpointId, LibraryItem, MAX_RECORDS_PER_RESPONSE, MusicDhtError, MusicDhtStorage, NodeContact, NodeId, SecretKey, StoreDecision, StoredRecord, decide_store, }; use sqlx::{PgPool, Row as _}; const IDENTITY_NAME: &str = "default"; const SCHEMA: &[&str] = &[ "CREATE TABLE IF NOT EXISTS furumusic__federation_identity ( name TEXT PRIMARY KEY, secret_key BYTEA NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL )", "CREATE TABLE IF NOT EXISTS furumusic__federation_local_item ( id BYTEA PRIMARY KEY, normalized_name TEXT NOT NULL, revision BIGINT NOT NULL, deleted BOOLEAN NOT NULL DEFAULT false, updated_at_ms BIGINT NOT NULL, payload BYTEA NOT NULL )", "CREATE INDEX IF NOT EXISTS idx_furumusic_federation_local_item_normalized_name ON furumusic__federation_local_item(normalized_name)", "CREATE TABLE IF NOT EXISTS furumusic__federation_dht_record ( dht_key BYTEA NOT NULL, item_id BYTEA NOT NULL, owner_peer_id TEXT NOT NULL, payload BYTEA NOT NULL, revision BIGINT NOT NULL, deleted BOOLEAN NOT NULL, expires_at_ms BIGINT NOT NULL, PRIMARY KEY (dht_key, item_id, owner_peer_id) )", "CREATE INDEX IF NOT EXISTS idx_furumusic_federation_dht_record_expires_at ON furumusic__federation_dht_record(expires_at_ms)", "CREATE TABLE IF NOT EXISTS furumusic__federation_known_peer ( peer_id TEXT PRIMARY KEY, node_id BYTEA NOT NULL, ticket TEXT NOT NULL, last_seen_ms BIGINT NOT NULL )", "CREATE TABLE IF NOT EXISTS furumusic__federation_content_id_cache ( media_file_id BIGINT PRIMARY KEY, sha256_hash TEXT NOT NULL, content_id TEXT NOT NULL, updated_at TEXT NOT NULL )", "CREATE INDEX IF NOT EXISTS idx_furumusic_federation_content_id_cache_content_id ON furumusic__federation_content_id_cache(content_id)", ]; #[derive(Debug, Clone)] pub struct PostgresFederationStorage { pool: PgPool, } impl PostgresFederationStorage { pub async fn new(pool: PgPool) -> music_dht::Result { let storage = Self { pool }; storage.ensure_schema().await?; Ok(storage) } pub async fn load_or_create_secret_key(&self) -> music_dht::Result { if let Some(bytes) = sqlx::query_scalar::<_, Vec>( "SELECT secret_key FROM furumusic__federation_identity WHERE name = $1", ) .bind(IDENTITY_NAME) .fetch_optional(&self.pool) .await .map_err(db_error)? { return secret_from_bytes(bytes); } let key = SecretKey::generate(); let key_bytes = key.to_bytes(); let now = now_iso(); let inserted = sqlx::query( "INSERT INTO furumusic__federation_identity (name, secret_key, created_at, updated_at) VALUES ($1, $2, $3, $4) ON CONFLICT (name) DO NOTHING", ) .bind(IDENTITY_NAME) .bind(key_bytes.as_slice()) .bind(&now) .bind(&now) .execute(&self.pool) .await .map_err(db_error)? .rows_affected(); if inserted == 1 { return Ok(key); } let bytes = sqlx::query_scalar::<_, Vec>( "SELECT secret_key FROM furumusic__federation_identity WHERE name = $1", ) .bind(IDENTITY_NAME) .fetch_one(&self.pool) .await .map_err(db_error)?; secret_from_bytes(bytes) } async fn ensure_schema(&self) -> music_dht::Result<()> { for sql in SCHEMA { sqlx::query(sql) .execute(&self.pool) .await .map_err(db_error)?; } Ok(()) } } #[async_trait] impl MusicDhtStorage for PostgresFederationStorage { async fn upsert_local_item(&self, item: &LibraryItem) -> music_dht::Result<()> { let payload = postcard::to_stdvec(item).map_err(db_error)?; sqlx::query( "INSERT INTO furumusic__federation_local_item (id, normalized_name, revision, deleted, updated_at_ms, payload) VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT (id) DO UPDATE SET normalized_name = EXCLUDED.normalized_name, revision = EXCLUDED.revision, deleted = EXCLUDED.deleted, updated_at_ms = EXCLUDED.updated_at_ms, payload = EXCLUDED.payload", ) .bind(item.id.as_bytes().as_slice()) .bind(&item.normalized_name) .bind(item.revision as i64) .bind(item.deleted) .bind(item.updated_at_ms as i64) .bind(payload) .execute(&self.pool) .await .map_err(db_error)?; Ok(()) } async fn list_local_items(&self, include_deleted: bool) -> music_dht::Result> { let rows = sqlx::query( "SELECT payload FROM furumusic__federation_local_item WHERE $1 OR deleted = false ORDER BY normalized_name", ) .bind(include_deleted) .fetch_all(&self.pool) .await .map_err(db_error)?; Ok(rows .into_iter() .filter_map(|row| postcard::from_bytes::(&row.get::, _>(0)).ok()) .collect()) } async fn store_dht_record(&self, key: DhtKey, record: StoredRecord) -> music_dht::Result { let mut conn = self.pool.acquire().await.map_err(db_error)?; store_record_in_conn(&mut conn, &key, &record).await } async fn store_dht_records( &self, entries: Vec<(DhtKey, StoredRecord)>, ) -> music_dht::Result> { let mut tx = self.pool.begin().await.map_err(db_error)?; let mut stored = Vec::with_capacity(entries.len()); for (key, record) in &entries { stored.push(store_record_in_conn(&mut tx, key, record).await?); } tx.commit().await.map_err(db_error)?; Ok(stored) } async fn dht_records_by_key( &self, key: DhtKey, now_ms: u64, ) -> music_dht::Result> { let rows = sqlx::query( "SELECT payload FROM furumusic__federation_dht_record WHERE dht_key = $1 AND expires_at_ms > $2 ORDER BY expires_at_ms DESC, item_id LIMIT $3", ) .bind(key.as_bytes().as_slice()) .bind(now_ms as i64) .bind(MAX_RECORDS_PER_RESPONSE as i64) .fetch_all(&self.pool) .await .map_err(db_error)?; Ok(rows .into_iter() .filter_map(|row| postcard::from_bytes::(&row.get::, _>(0)).ok()) .collect()) } async fn delete_expired_records(&self, now_ms: u64) -> music_dht::Result { let result = sqlx::query("DELETE FROM furumusic__federation_dht_record WHERE expires_at_ms <= $1") .bind(now_ms as i64) .execute(&self.pool) .await .map_err(db_error)?; Ok(result.rows_affected() as usize) } async fn upsert_known_peer(&self, contact: &NodeContact) -> music_dht::Result<()> { sqlx::query( "INSERT INTO furumusic__federation_known_peer (peer_id, node_id, ticket, last_seen_ms) VALUES ($1, $2, $3, $4) ON CONFLICT (peer_id) DO UPDATE SET node_id = EXCLUDED.node_id, ticket = EXCLUDED.ticket, last_seen_ms = EXCLUDED.last_seen_ms", ) .bind(contact.peer_id.to_string()) .bind(contact.node_id.as_bytes().as_slice()) .bind(&contact.ticket) .bind(contact.last_seen_ms as i64) .execute(&self.pool) .await .map_err(db_error)?; Ok(()) } async fn delete_known_peer(&self, peer_id: EndpointId) -> music_dht::Result<()> { sqlx::query("DELETE FROM furumusic__federation_known_peer WHERE peer_id = $1") .bind(peer_id.to_string()) .execute(&self.pool) .await .map_err(db_error)?; Ok(()) } async fn load_known_peers(&self) -> music_dht::Result> { let rows = sqlx::query( "SELECT peer_id, node_id, ticket, last_seen_ms FROM furumusic__federation_known_peer", ) .fetch_all(&self.pool) .await .map_err(db_error)?; let mut contacts = Vec::new(); for row in rows { let peer_id: String = row.get(0); let node_id: Vec = row.get(1); let ticket: String = row.get(2); let last_seen_ms: i64 = row.get(3); let Ok(peer_id) = EndpointId::from_str(&peer_id) else { continue; }; let Ok(node_id) = <[u8; 32]>::try_from(node_id.as_slice()) else { continue; }; contacts.push(NodeContact { node_id: NodeId::from_bytes(node_id), peer_id, ticket, last_seen_ms: last_seen_ms as u64, }); } Ok(contacts) } } fn now_iso() -> String { chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string() } fn secret_from_bytes(bytes: Vec) -> music_dht::Result { let bytes: [u8; 32] = bytes .as_slice() .try_into() .map_err(|_| MusicDhtError::Database("stored federation identity is corrupted".into()))?; Ok(SecretKey::from_bytes(&bytes)) } /// Applies one validated record following the revision/tombstone rules. /// Returns `true` if the record was written or refreshed. Runs against a /// pooled connection or an open transaction. async fn store_record_in_conn( conn: &mut sqlx::PgConnection, key: &DhtKey, record: &StoredRecord, ) -> music_dht::Result { let existing = sqlx::query( "SELECT revision, deleted, expires_at_ms FROM furumusic__federation_dht_record WHERE dht_key = $1 AND item_id = $2 AND owner_peer_id = $3", ) .bind(key.as_bytes().as_slice()) .bind(record.item.id.as_bytes().as_slice()) .bind(record.item.owner.to_string()) .fetch_optional(&mut *conn) .await .map_err(db_error)? .map(|row| { ( row.get::(0) as u64, row.get::(1), row.get::(2) as u64, ) }); match decide_store(existing, record) { StoreDecision::Ignore => return Ok(false), StoreDecision::Write | StoreDecision::RefreshExpiry(_) => {} } let payload = postcard::to_stdvec(record).map_err(db_error)?; sqlx::query( "INSERT INTO furumusic__federation_dht_record (dht_key, item_id, owner_peer_id, payload, revision, deleted, expires_at_ms) VALUES ($1, $2, $3, $4, $5, $6, $7) ON CONFLICT (dht_key, item_id, owner_peer_id) DO UPDATE SET payload = EXCLUDED.payload, revision = EXCLUDED.revision, deleted = EXCLUDED.deleted, expires_at_ms = EXCLUDED.expires_at_ms", ) .bind(key.as_bytes().as_slice()) .bind(record.item.id.as_bytes().as_slice()) .bind(record.item.owner.to_string()) .bind(payload) .bind(record.item.revision as i64) .bind(record.item.deleted) .bind(record.expires_at_ms as i64) .execute(&mut *conn) .await .map_err(db_error)?; Ok(true) } fn db_error(err: impl std::fmt::Display) -> MusicDhtError { MusicDhtError::Database(err.to_string()) }