SoulStream/crates/soulseek-client/src/share.rs

//! Share index: the files we make available to the network. Built by
//! walking configured folders; queried by the search and browse responders.
//!
//! Remote paths use backslashes and a virtual root so we never leak local
//! absolute paths. Each shared folder is mounted under its own base name,
//! e.g. local `D:\Music\Album` → remote `Music\Album\track.flac`.

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::RwLock;

use soulseek_proto::types::{Directory, File};

const AUDIO_EXTS: [&str; 10] = [
    "flac", "mp3", "m4a", "ogg", "opus", "wav", "aac", "wma", "ape", "aiff",
];

/// One shared file: its remote (virtual) path and where it lives locally.
#[derive(Debug, Clone)]
struct SharedFile {
    remote_path: String, // backslash-separated
    local_path: PathBuf,
    size: u64,
    extension: String,
    lowercase: String, // remote_path lowercased, for matching
}

#[derive(Default)]
struct Index {
    files: Vec<SharedFile>,
    /// remote folder (backslashes) → indices into `files`.
    folders: BTreeMap<String, Vec<usize>>,
}

pub struct ShareIndex {
    index: RwLock<Index>,
}

impl Default for ShareIndex {
    fn default() -> Self {
        ShareIndex {
            index: RwLock::new(Index::default()),
        }
    }
}

impl ShareIndex {
    pub fn new() -> Self {
        Self::default()
    }

    /// Rebuild the index from the given local folders. Returns (dirs, files).
    pub fn rebuild(&self, roots: &[PathBuf]) -> (u32, u32) {
        let mut index = Index::default();
        for root in roots {
            let base = root
                .file_name()
                .map(|n| n.to_string_lossy().into_owned())
                .unwrap_or_else(|| "Shared".into());
            walk(root, &base, root, &mut index);
        }
        let dirs = index.folders.len() as u32;
        let files = index.files.len() as u32;
        *self.index.write().unwrap() = index;
        (dirs, files)
    }

    pub fn counts(&self) -> (u32, u32) {
        let idx = self.index.read().unwrap();
        (idx.folders.len() as u32, idx.files.len() as u32)
    }

    pub fn file_count(&self) -> u32 {
        self.index.read().unwrap().files.len() as u32
    }

    /// Match a search query: every whitespace token must appear somewhere
    /// in the remote path (case-insensitive). Returns proto `File`s.
    /// `excluded` phrases veto a result if any appears in the path.
    pub fn search(&self, query: &str, excluded: &[String], limit: usize) -> Vec<File> {
        let terms: Vec<String> = query
            .split_whitespace()
            .map(|t| t.to_lowercase())
            .filter(|t| !t.is_empty())
            .collect();
        if terms.is_empty() {
            return Vec::new();
        }
        let idx = self.index.read().unwrap();
        let mut out = Vec::new();
        for f in &idx.files {
            if terms.iter().all(|t| f.lowercase.contains(t.as_str()))
                && !excluded
                    .iter()
                    .any(|p| !p.is_empty() && f.lowercase.contains(&p.to_lowercase()))
            {
                out.push(to_proto_file(f));
                if out.len() >= limit {
                    break;
                }
            }
        }
        out
    }

    /// Resolve a remote path a peer asked to download to its local file.
    pub fn resolve(&self, remote_path: &str) -> Option<(PathBuf, u64)> {
        let idx = self.index.read().unwrap();
        idx.files
            .iter()
            .find(|f| f.remote_path.eq_ignore_ascii_case(remote_path))
            .map(|f| (f.local_path.clone(), f.size))
    }

    /// Full directory listing for a browse response.
    pub fn browse(&self) -> Vec<Directory> {
        let idx = self.index.read().unwrap();
        idx.folders
            .iter()
            .map(|(name, indices)| Directory {
                name: name.clone(),
                files: indices
                    .iter()
                    .map(|&i| to_proto_file(&idx.files[i]))
                    .collect(),
            })
            .collect()
    }

    /// Files in one folder (for FolderContentsResponse).
    pub fn folder(&self, name: &str) -> Option<Directory> {
        let idx = self.index.read().unwrap();
        idx.folders.get(name).map(|indices| Directory {
            name: name.to_string(),
            files: indices
                .iter()
                .map(|&i| to_proto_file(&idx.files[i]))
                .collect(),
        })
    }
}

fn to_proto_file(f: &SharedFile) -> File {
    // Search/browse responses send just the leaf filename in the File entry
    // for search, but the full path for browse. We store the remote path
    // and let the responder decide; here we return the remote path as the
    // filename (peers expect a full backslash path in search responses).
    File {
        filename: f.remote_path.clone(),
        size: f.size,
        extension: f.extension.clone(),
        attributes: Vec::new(),
    }
}

fn walk(dir: &Path, base: &str, root: &Path, index: &mut Index) {
    let Ok(entries) = std::fs::read_dir(dir) else {
        return;
    };
    for entry in entries.flatten() {
        let path = entry.path();
        if path.is_dir() {
            // The download staging pen is not part of anyone's library —
            // half-finished transfers are not shareable records.
            if path.file_name().is_some_and(|n| n == "(Cache)") {
                continue;
            }
            walk(&path, base, root, index);
        } else if is_audio(&path) {
            let Ok(meta) = entry.metadata() else { continue };
            let remote_path = to_remote_path(&path, base, root);
            let folder = remote_path
                .rsplit_once('\\')
                .map(|(d, _)| d.to_string())
                .unwrap_or_else(|| base.to_string());
            let ext = path
                .extension()
                .map(|e| e.to_string_lossy().to_lowercase())
                .unwrap_or_default();
            let file = SharedFile {
                lowercase: remote_path.to_lowercase(),
                remote_path,
                local_path: path,
                size: meta.len(),
                extension: ext,
            };
            let idx = index.files.len();
            index.folders.entry(folder).or_default().push(idx);
            index.files.push(file);
        }
    }
}

/// `<base>\<relative path with backslashes>`.
fn to_remote_path(path: &Path, base: &str, root: &Path) -> String {
    let rel = path.strip_prefix(root).unwrap_or(path);
    let rel_str: String = rel
        .components()
        .map(|c| c.as_os_str().to_string_lossy().into_owned())
        .collect::<Vec<_>>()
        .join("\\");
    format!("{base}\\{rel_str}")
}

fn is_audio(path: &Path) -> bool {
    path.extension()
        .and_then(|e| e.to_str())
        .map(|e| AUDIO_EXTS.contains(&e.to_ascii_lowercase().as_str()))
        .unwrap_or(false)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn index_search_and_resolve() {
        let dir = std::env::temp_dir().join("ssplus-share-test");
        let album = dir.join("Nevermind");
        std::fs::create_dir_all(&album).unwrap();
        std::fs::write(album.join("01 Lithium.flac"), b"fake").unwrap();
        std::fs::write(album.join("02 Polly.mp3"), b"fake2").unwrap();
        std::fs::write(album.join("cover.jpg"), b"img").unwrap(); // ignored

        let share = ShareIndex::new();
        let (dirs, files) = share.rebuild(&[dir.clone()]);
        assert_eq!(files, 2);
        assert!(dirs >= 1);

        // Matches on both terms across the path.
        let hits = share.search("nevermind lithium", &[], 10);
        assert_eq!(hits.len(), 1);
        assert!(hits[0].filename.ends_with("01 Lithium.flac"));

        // Excluded phrase vetoes.
        let vetoed = share.search("lithium", &["nevermind".into()], 10);
        assert!(vetoed.is_empty());

        // Resolve remote → local.
        let remote = &hits[0].filename;
        let (local, size) = share.resolve(remote).unwrap();
        assert!(local.ends_with("01 Lithium.flac"));
        assert_eq!(size, 4);

        // Browse lists the folder.
        let browse = share.browse();
        assert!(browse
            .iter()
            .any(|d| d.name.contains("Nevermind") && d.files.len() == 2));
    }
}