arr/crates/arr-dl/src/lib.rs

//! Transmission RPC client for download lifecycle operations.

use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use base64::Engine as _;
use reqwest::{StatusCode, Url};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tokio::sync::RwLock;

const SESSION_HEADER: &str = "x-transmission-session-id";
const MAX_SESSION_NEGOTIATIONS: usize = 4;
const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);

/// Where Transmission is to get the torrent from.
///
/// Deliberately not an indexer URL: Transmission has no route to Prowlarr and
/// cannot follow its redirect to a magnet, so arr resolves the link itself and
/// hands over the result (issue #100).
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum TorrentSource {
    Magnet(String),
    /// The bytes of a `.torrent` file, sent inline.
    Metainfo(Vec<u8>),
}

/// Settings applied atomically when a torrent is added.
#[derive(Clone, Debug, PartialEq)]
pub struct AddTorrent {
    pub source: TorrentSource,
    pub label: String,
    pub download_dir: PathBuf,
    pub seed_ratio_limit: f64,
    /// Transmission expresses its idle limit in minutes.
    pub seed_idle_limit_minutes: u64,
}

/// The identity returned by Transmission for an added or existing torrent.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct AddedTorrent {
    pub id: i64,
    pub name: String,
    pub hash: String,
    pub was_duplicate: bool,
}

/// Transmission's torrent activity state.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TorrentState {
    Stopped,
    QueuedToVerify,
    Verifying,
    QueuedToDownload,
    Downloading,
    QueuedToSeed,
    Seeding,
    Unknown(u8),
}

/// Current torrent state, read directly from Transmission.
#[derive(Clone, Debug, PartialEq)]
pub struct Torrent {
    pub id: i64,
    pub name: String,
    pub hash: String,
    pub state: TorrentState,
    pub progress: f64,
    /// Download rate in bytes per second.
    pub download_rate: u64,
    /// Estimated seconds until complete, or `None` when Transmission cannot
    /// estimate it.
    pub eta: Option<i64>,
    /// Transmission's error message, when the torrent has errored.
    pub error: Option<String>,
    /// Transmission's own verdict that this torrent has gone quiet for longer
    /// than its configured stalled window. A rate of zero is not the same
    /// fact: a torrent between peers reads zero for one poll and is fine.
    pub is_stalled: bool,
    pub download_dir: PathBuf,
    pub labels: Vec<String>,
    /// Transmission has stopped this torrent because its configured seeding
    /// ratio or idle limit was reached.
    pub is_finished: bool,
}

/// One file inside a torrent, as Transmission reports it.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct TorrentFile {
    /// Path relative to the torrent's download directory, torrent folder
    /// included.
    pub path: PathBuf,
    /// Size in bytes.
    pub size: u64,
}

/// Where one torrent's data lives on disk.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct TorrentContent {
    pub hash: String,
    pub download_dir: PathBuf,
    pub files: Vec<TorrentFile>,
}

/// A Transmission RPC failure.
#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("invalid Transmission RPC URL: {0}")]
    InvalidUrl(#[from] url::ParseError),
    #[error("download directory is not valid UTF-8")]
    NonUtf8DownloadDir,
    #[error("could not build Transmission HTTP client: {0}")]
    BuildClient(reqwest::Error),
    #[error("Transmission request failed: {0}")]
    Http(#[from] reqwest::Error),
    #[error("Transmission did not return a session ID after {0} attempts")]
    SessionNegotiation(usize),
    #[error("Transmission RPC failed: {0}")]
    Rpc(String),
    #[error("invalid Transmission response: {0}")]
    InvalidResponse(String),
}

/// Client for the unauthenticated Transmission JSON-RPC endpoint.
#[derive(Clone, Debug)]
pub struct TransmissionClient {
    endpoint: Url,
    http: reqwest::Client,
    session_id: Arc<RwLock<Option<String>>>,
}

impl TransmissionClient {
    /// Create a client for a full RPC URL, such as
    /// `http://host:9091/transmission/rpc`.
    ///
    /// # Errors
    ///
    /// Returns an error when `endpoint` is not a valid URL.
    pub fn new(endpoint: &str) -> Result<Self, Error> {
        Ok(Self {
            endpoint: Url::parse(endpoint)?,
            http: reqwest::Client::builder()
                .connect_timeout(CONNECT_TIMEOUT)
                .timeout(REQUEST_TIMEOUT)
                .build()
                .map_err(Error::BuildClient)?,
            session_id: Arc::new(RwLock::new(None)),
        })
    }

    /// Add a torrent with its label, destination, and tracker seed limits.
    ///
    /// # Errors
    ///
    /// Returns an error for transport failures or rejected/malformed RPC
    /// responses.
    pub async fn add_torrent(&self, request: AddTorrent) -> Result<AddedTorrent, Error> {
        let download_dir = request
            .download_dir
            .to_str()
            .ok_or(Error::NonUtf8DownloadDir)?
            .to_owned();
        let label = request.label;
        let seed_ratio_limit = request.seed_ratio_limit;
        let seed_idle_limit = request.seed_idle_limit_minutes;
        let mut arguments = json!({
            "download-dir": download_dir,
            "labels": [label],
            "seedRatioLimit": seed_ratio_limit,
            "seedRatioMode": 1,
            "seedIdleLimit": seed_idle_limit,
            "seedIdleMode": 1,
        });

        match request.source {
            TorrentSource::Magnet(uri) => {
                arguments["filename"] = json!(uri);
            }
            TorrentSource::Metainfo(bytes) => {
                arguments["metainfo"] =
                    json!(base64::engine::general_purpose::STANDARD.encode(bytes));
            }
        }

        let arguments = self.call("torrent-add", arguments).await?;
        let (torrent, was_duplicate) = if let Some(value) = arguments.get("torrent-added") {
            (value.clone(), false)
        } else if let Some(value) = arguments.get("torrent-duplicate") {
            (value.clone(), true)
        } else {
            return Err(Error::InvalidResponse(
                "torrent-add omitted torrent-added and torrent-duplicate".into(),
            ));
        };
        let torrent: RpcAddedTorrent = serde_json::from_value(torrent)
            .map_err(|error| Error::InvalidResponse(error.to_string()))?;

        if was_duplicate {
            self.call(
                "torrent-set",
                json!({
                    "ids": [torrent.id],
                    "labels": [label],
                    "seedRatioLimit": seed_ratio_limit,
                    "seedRatioMode": 1,
                    "seedIdleLimit": seed_idle_limit,
                    "seedIdleMode": 1,
                }),
            )
            .await?;
            self.call(
                "torrent-set-location",
                json!({"ids": [torrent.id], "location": download_dir, "move": true}),
            )
            .await?;
        }

        Ok(AddedTorrent {
            id: torrent.id,
            name: torrent.name,
            hash: torrent.hash,
            was_duplicate,
        })
    }

    /// List torrents with their current state and completion progress.
    ///
    /// No torrent state is cached: every result is authoritative from
    /// Transmission, including the first call after process startup.
    ///
    /// # Errors
    ///
    /// Returns an error for transport failures or rejected/malformed RPC
    /// responses.
    pub async fn list_torrents(&self) -> Result<Vec<Torrent>, Error> {
        let arguments = self
            .call(
                "torrent-get",
                json!({
                    "fields": [
                        "id", "name", "hashString", "status", "percentDone",
                        "rateDownload", "eta", "errorString", "isStalled",
                        "downloadDir", "labels", "isFinished"
                    ]
                }),
            )
            .await?;
        let response: RpcTorrentList = serde_json::from_value(arguments)
            .map_err(|error| Error::InvalidResponse(error.to_string()))?;

        Ok(response.torrents.into_iter().map(Torrent::from).collect())
    }

    /// The file list and download directory of one torrent, by infohash.
    ///
    /// `None` when Transmission no longer knows the hash — the operator may
    /// have removed the torrent by hand.
    ///
    /// # Errors
    ///
    /// Returns an error for transport failures or rejected/malformed RPC
    /// responses.
    pub async fn torrent_content(&self, hash: &str) -> Result<Option<TorrentContent>, Error> {
        let arguments = self
            .call(
                "torrent-get",
                json!({
                    "ids": [hash],
                    "fields": ["hashString", "downloadDir", "files"]
                }),
            )
            .await?;
        let response: RpcContentList = serde_json::from_value(arguments)
            .map_err(|error| Error::InvalidResponse(error.to_string()))?;

        Ok(response
            .torrents
            .into_iter()
            .next()
            .map(|torrent| TorrentContent {
                hash: torrent.hash,
                download_dir: torrent.download_dir,
                files: torrent
                    .files
                    .into_iter()
                    .map(|file| TorrentFile {
                        path: file.name,
                        size: file.length,
                    })
                    .collect(),
            }))
    }

    /// Remove one torrent, optionally deleting its downloaded data.
    ///
    /// # Errors
    ///
    /// Returns an error for transport failures or rejected RPC responses.
    pub async fn remove_torrent(&self, id: i64, delete_data: bool) -> Result<(), Error> {
        self.call(
            "torrent-remove",
            json!({"ids": [id], "delete-local-data": delete_data}),
        )
        .await?;
        Ok(())
    }

    async fn call(&self, method: &str, arguments: Value) -> Result<Value, Error> {
        let body = RpcRequest { method, arguments };

        for _ in 0..MAX_SESSION_NEGOTIATIONS {
            let session_id = self.session_id.read().await.clone();
            let mut request = self.http.post(self.endpoint.clone()).json(&body);
            if let Some(session_id) = session_id {
                request = request.header(SESSION_HEADER, session_id);
            }

            let response = request.send().await?;
            if response.status() == StatusCode::CONFLICT {
                let session_id = response
                    .headers()
                    .get(SESSION_HEADER)
                    .and_then(|value| value.to_str().ok())
                    .ok_or_else(|| {
                        Error::InvalidResponse("409 response omitted session ID".into())
                    })?;
                *self.session_id.write().await = Some(session_id.to_owned());
                continue;
            }

            let response = response.error_for_status()?.json::<RpcResponse>().await?;
            if response.result != "success" {
                return Err(Error::Rpc(response.result));
            }
            return Ok(response.arguments);
        }

        Err(Error::SessionNegotiation(MAX_SESSION_NEGOTIATIONS))
    }
}

#[derive(Debug, Serialize)]
struct RpcRequest<'a> {
    method: &'a str,
    arguments: Value,
}

#[derive(Debug, Deserialize)]
struct RpcResponse {
    result: String,
    arguments: Value,
}

#[derive(Debug, Deserialize)]
struct RpcAddedTorrent {
    id: i64,
    name: String,
    #[serde(rename = "hashString")]
    hash: String,
}

#[derive(Debug, Deserialize)]
struct RpcTorrentList {
    torrents: Vec<RpcTorrent>,
}

#[derive(Debug, Deserialize)]
struct RpcContentList {
    torrents: Vec<RpcContent>,
}

#[derive(Debug, Deserialize)]
struct RpcContent {
    #[serde(rename = "hashString")]
    hash: String,
    #[serde(rename = "downloadDir")]
    download_dir: PathBuf,
    #[serde(default)]
    files: Vec<RpcFile>,
}

#[derive(Debug, Deserialize)]
struct RpcFile {
    name: PathBuf,
    length: u64,
}

#[derive(Debug, Deserialize)]
struct RpcTorrent {
    id: i64,
    name: String,
    #[serde(rename = "hashString")]
    hash: String,
    status: u8,
    #[serde(rename = "percentDone")]
    progress: f64,
    #[serde(rename = "rateDownload", default)]
    download_rate: u64,
    #[serde(default)]
    eta: Option<i64>,
    #[serde(rename = "errorString", default)]
    error: String,
    #[serde(rename = "isStalled", default)]
    is_stalled: bool,
    #[serde(rename = "downloadDir")]
    download_dir: PathBuf,
    #[serde(default)]
    labels: Vec<String>,
    #[serde(rename = "isFinished", default)]
    is_finished: bool,
}

impl From<RpcTorrent> for Torrent {
    fn from(value: RpcTorrent) -> Self {
        let state = match value.status {
            0 => TorrentState::Stopped,
            1 => TorrentState::QueuedToVerify,
            2 => TorrentState::Verifying,
            3 => TorrentState::QueuedToDownload,
            4 => TorrentState::Downloading,
            5 => TorrentState::QueuedToSeed,
            6 => TorrentState::Seeding,
            status => TorrentState::Unknown(status),
        };

        Self {
            id: value.id,
            name: value.name,
            hash: value.hash,
            state,
            progress: value.progress,
            download_rate: value.download_rate,
            eta: value.eta.filter(|eta| *eta >= 0),
            error: (!value.error.is_empty()).then_some(value.error),
            is_stalled: value.is_stalled,
            download_dir: value.download_dir,
            labels: value.labels,
            is_finished: value.is_finished,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;

    use serde_json::json;
    use wiremock::matchers::{body_partial_json, method};
    use wiremock::{Mock, MockServer, ResponseTemplate};

    use super::{AddTorrent, Error, TorrentSource, TorrentState, TransmissionClient};

    #[tokio::test]
    async fn retries_the_session_handshake() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .respond_with(
                ResponseTemplate::new(409).insert_header("x-transmission-session-id", "session"),
            )
            .with_priority(1)
            .up_to_n_times(1)
            .expect(1)
            .mount(&server)
            .await;
        Mock::given(method("POST"))
            .and(body_partial_json(json!({"method": "torrent-get"})))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "result": "success",
                "arguments": {"torrents": []}
            })))
            .with_priority(2)
            .expect(1)
            .mount(&server)
            .await;

        let client = TransmissionClient::new(&server.uri()).expect("client");
        assert!(client.list_torrents().await.expect("list").is_empty());
    }

    #[tokio::test]
    async fn duplicate_reapplies_label_location_and_seed_limits() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(body_partial_json(json!({"method": "torrent-add"})))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "result": "success",
                "arguments": {"torrent-duplicate": {
                    "id": 7, "name": "duplicate", "hashString": "abc"
                }}
            })))
            .expect(1)
            .mount(&server)
            .await;
        Mock::given(method("POST"))
            .and(body_partial_json(json!({
                "method": "torrent-set",
                "arguments": {
                    "ids": [7], "labels": ["movies-main"],
                    "seedRatioLimit": 1.5, "seedIdleLimit": 60
                }
            })))
            .respond_with(success())
            .expect(1)
            .mount(&server)
            .await;
        Mock::given(method("POST"))
            .and(body_partial_json(json!({
                "method": "torrent-set-location",
                "arguments": {"ids": [7], "location": "/downloads", "move": true}
            })))
            .respond_with(success())
            .expect(1)
            .mount(&server)
            .await;

        let client = TransmissionClient::new(&server.uri()).expect("client");
        let added = client
            .add_torrent(AddTorrent {
                source: TorrentSource::Magnet("magnet:?xt=urn:btih:abc".into()),
                label: "movies-main".into(),
                download_dir: PathBuf::from("/downloads"),
                seed_ratio_limit: 1.5,
                seed_idle_limit_minutes: 60,
            })
            .await
            .expect("duplicate");

        assert!(added.was_duplicate);
    }

    #[tokio::test]
    #[cfg(unix)]
    async fn rejects_non_utf8_download_directory() {
        use std::os::unix::ffi::OsStringExt as _;

        let client = TransmissionClient::new("http://127.0.0.1:1").expect("client");
        let error = client
            .add_torrent(AddTorrent {
                source: TorrentSource::Magnet("magnet:?xt=urn:btih:abc".into()),
                label: "movies-main".into(),
                download_dir: PathBuf::from(std::ffi::OsString::from_vec(vec![0xff])),
                seed_ratio_limit: 1.5,
                seed_idle_limit_minutes: 60,
            })
            .await
            .expect_err("invalid path");

        assert!(matches!(error, Error::NonUtf8DownloadDir));
    }

    #[tokio::test]
    async fn preserves_unknown_status_without_hiding_the_list() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "result": "success",
                "arguments": {"torrents": [{
                    "id": 1, "name": "future", "hashString": "abc", "status": 99,
                    "percentDone": 0.5, "downloadDir": "/downloads", "labels": []
                }]}
            })))
            .expect(1)
            .mount(&server)
            .await;

        let client = TransmissionClient::new(&server.uri()).expect("client");
        let torrents = client.list_torrents().await.expect("list");
        assert_eq!(torrents[0].state, TorrentState::Unknown(99));
    }

    #[tokio::test]
    async fn torrent_content_lists_files_and_download_dir() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(body_partial_json(json!({
                "method": "torrent-get",
                "arguments": {"ids": ["abc"]}
            })))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "result": "success",
                "arguments": {"torrents": [{
                    "hashString": "abc",
                    "downloadDir": "/downloads",
                    "files": [
                        {"name": "Movie/Movie.mkv", "length": 100, "bytesCompleted": 100},
                        {"name": "Movie/Movie.nfo", "length": 5, "bytesCompleted": 5}
                    ]
                }]}
            })))
            .expect(1)
            .mount(&server)
            .await;

        let client = TransmissionClient::new(&server.uri()).expect("client");
        let content = client
            .torrent_content("abc")
            .await
            .expect("content")
            .expect("torrent known");

        assert_eq!(content.download_dir, PathBuf::from("/downloads"));
        assert_eq!(content.files.len(), 2);
        assert_eq!(content.files[0].path, PathBuf::from("Movie/Movie.mkv"));
        assert_eq!(content.files[0].size, 100);
    }

    #[tokio::test]
    async fn a_torrent_transmission_forgot_is_none_not_an_error() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "result": "success",
                "arguments": {"torrents": []}
            })))
            .mount(&server)
            .await;

        let client = TransmissionClient::new(&server.uri()).expect("client");
        assert!(client
            .torrent_content("gone")
            .await
            .expect("call succeeds")
            .is_none());
    }

    fn success() -> ResponseTemplate {
        ResponseTemplate::new(200).set_body_json(json!({"result": "success", "arguments": {}}))
    }
}