apt_transport/

transport.rs

use std::collections::HashMap;

use crate::{
    Error,
    message::{Message, MessageType},
    message_stream::MessageStream,
};
use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader};

/// Asynchronous event stream for APT transport implementations.
pub struct AptStream {
    inner: Box<dyn AsyncBufRead + Unpin + Send>,
    message_stream: MessageStream,
    output_stream: Box<dyn AsyncWrite + Unpin + Send>,
    has_initialized: bool,
}

impl AptStream {
    /// Create a new `AsyncAptStream` with stdin/stdout as the underlying streams.
    pub fn new() -> Self {
        Self {
            inner: Box::new(BufReader::new(tokio::io::stdin())),
            message_stream: MessageStream::default(),
            output_stream: Box::new(tokio::io::stdout()),
            has_initialized: false,
        }
    }

    /// Sets the input stream for reading APT requests.
    ///
    /// Must only be set before any requests are processed
    /// (i.e. the first call to `.next()`), and should
    /// ultimately be routed to the stdin of the transport process
    /// that `apt` has spawned.
    pub fn with_input_stream(
        mut self,
        input_stream: Box<dyn AsyncBufRead + Unpin + Send>,
    ) -> Result<Self, Error> {
        if self.has_initialized {
            return Err(Error::StreamAlreadyInitialized);
        }
        self.inner = Box::new(input_stream);
        Ok(self)
    }

    /// Sets the output stream for all response writes.
    ///
    /// Must only be set before any requests are processed
    /// (i.e. the first call to `.next()`), and should
    /// ultimately be routed to the stdout of the transport process
    /// that `apt` has spawned.
    pub fn with_output_stream(
        mut self,
        output_stream: Box<dyn AsyncWrite + Unpin + Send>,
    ) -> Result<Self, Error> {
        if self.has_initialized {
            return Err(Error::StreamAlreadyInitialized);
        }
        self.output_stream = output_stream;
        Ok(self)
    }

    /// Pulls the next APT request from the stream.
    pub async fn next<'a>(&'a mut self) -> Result<Option<AptRequest<'a>>, Error> {
        if !self.has_initialized {
            self.has_initialized = true;

            // Synthesize a capabilities request at the start of the stream
            log::debug!("synthesizing capabilities request");
            let capabilities_req = AptRequest::Capabilities(CapabilitiesRequest { this: self });

            return Ok(Some(capabilities_req));
        }

        // Read lines until we get a complete message
        let mut line = String::new();
        loop {
            log::trace!("reading line from input stream");
            line.clear();
            let nread = self.inner.as_mut().read_line(&mut line).await?;
            log::trace!(
                "read {} bytes from input stream: {:?}",
                nread,
                line.as_bytes()
            );

            if let Some(message_result) = self.message_stream.push_line(line.as_bytes()) {
                log::debug!("complete message received");
                log::trace!("complete message received: {:?}", message_result);

                match message_result {
                    Ok(message) => {
                        let apt_request = AptRequest::<'a>::try_from_message((self, message)).await;

                        return match apt_request {
                            Ok(req) => {
                                log::debug!("APT request parse OK");
                                Ok(Some(req))
                            }
                            Err(e) => {
                                log::debug!("APT request parse error: {e:#?}");
                                Err(e)
                            }
                        };
                    }
                    Err(e) => {
                        log::debug!("APT message parse error: {e:?}");
                        return Err(e);
                    }
                }
            } else if nread == 0 {
                // EOF reached
                log::debug!("EOF reached on input stream");
                return Ok(None);
            }
        }
    }
}

/// An APT request.
pub enum AptRequest<'a> {
    /// A request for capabilities
    Capabilities(CapabilitiesRequest<'a>),
    /// A configuration request
    Configuration(ConfigRequest),
    /// A URI acquire request
    Uri(UriRequest<'a>),
}

/// An APT capabilities request.
pub struct CapabilitiesRequest<'a> {
    this: &'a mut AptStream,
}

impl<'a> CapabilitiesRequest<'a> {
    /// Creates a response for the capabilities request.
    #[must_use = "responses must be sent to the APT client with `.send().await`"]
    pub fn respond(self) -> CapabilitiesResponse<'a> {
        CapabilitiesResponse::new(self.this)
    }
}

/// A builder for APT capabilities responses.
pub struct CapabilitiesResponse<'a> {
    this: &'a mut AptStream,
    single_instance: bool,
    send_config: bool,
    pipeline: bool,
    local_only: bool,
    removable: bool,
    needs_cleanup: bool,
    version: String,
}

impl CapabilitiesResponse<'_> {
    /// Creates a new builder.
    #[must_use]
    fn new(this: &mut AptStream) -> CapabilitiesResponse<'_> {
        CapabilitiesResponse {
            this,
            single_instance: false,
            send_config: false,
            pipeline: false,
            local_only: false,
            removable: false,
            needs_cleanup: false,
            version: env!("CARGO_PKG_VERSION").to_string(),
        }
    }

    /// Whether or not to advertise single-instance support.
    #[must_use = "responses must be sent with `.send().await`"]
    pub fn single_instance(mut self, enabled: bool) -> Self {
        self.single_instance = enabled;
        self
    }

    /// Whether or not to advertise send-config support.
    #[must_use = "responses must be sent with `.send().await`"]
    pub fn send_config(mut self, enabled: bool) -> Self {
        self.send_config = enabled;
        self
    }

    /// Sets the version string to advertise.
    ///
    /// By default this is set to the `apt-transport` crate version, though
    /// this can be (should be) overridden to match the transport implementation version:
    ///
    /// ```ignore
    /// response.version(env!("CARGO_PKG_VERSION")).send().await?;
    /// ```
    #[must_use = "responses must be sent with `.send().await`"]
    pub fn version<S: Into<String>>(mut self, version: S) -> Self {
        self.version = version.into();
        self
    }

    /// Whether or not to advertise pipeline support.
    #[must_use = "responses must be sent with `.send().await`"]
    pub fn pipeline(mut self, enabled: bool) -> Self {
        self.pipeline = enabled;
        self
    }

    // Whether or not to advertise local-only support.
    #[must_use = "responses must be sent with `.send().await`"]
    pub fn local_only(mut self, enabled: bool) -> Self {
        self.local_only = enabled;
        self
    }

    /// Whether or not to advertise removable support.
    #[must_use = "responses must be sent with `.send().await`"]
    pub fn removable(mut self, enabled: bool) -> Self {
        self.removable = enabled;
        self
    }

    /// Whether or not to advertise needs-cleanup support.
    #[must_use = "responses must be sent with `.send().await`"]
    pub fn needs_cleanup(mut self, enabled: bool) -> Self {
        self.needs_cleanup = enabled;
        self
    }

    /// Responds with the given capabilities.
    pub async fn send(self) -> Result<(), Error> {
        let msg = Message::new(
            MessageType::Capabilities,
            vec![
                ("Version", &self.version),
                (
                    "Send-Config",
                    if self.send_config { "true" } else { "false" },
                ),
                (
                    "Single-Instance",
                    if self.single_instance {
                        "true"
                    } else {
                        "false"
                    },
                ),
                ("Pipeline", if self.pipeline { "true" } else { "false" }),
                ("Local-Only", if self.local_only { "true" } else { "false" }),
                ("Removable", if self.removable { "true" } else { "false" }),
                (
                    "Needs-Cleanup",
                    if self.needs_cleanup { "true" } else { "false" },
                ),
            ],
        )
        .to_string();

        log::debug!("sending capabilities response");
        log::trace!("sending capabilities response: {}", msg.trim());
        self.this
            .output_stream
            .as_mut()
            .write_all(msg.as_bytes())
            .await?;

        log::debug!("flushing output stream");
        self.this.output_stream.as_mut().flush().await?;

        log::debug!("capabilities response sent successfully");
        Ok(())
    }
}

/// Configuration request, containing all configuration options.
///
/// This request has no response.
pub struct ConfigRequest {
    options: HashMap<String, String>,
}

impl ConfigRequest {
    fn from(message: Message) -> Self {
        let options = message
            .headers
            .into_iter()
            .filter_map(|(key, value)| {
                if key == "Config-Item" {
                    let (key, value) = value.split_once('=')?;
                    Some((key.to_string(), value.to_string()))
                } else {
                    None
                }
            })
            .collect::<HashMap<String, String>>();

        ConfigRequest { options }
    }

    /// Gets the configuration options.
    pub fn options(&self) -> &HashMap<String, String> {
        &self.options
    }
}

/// A request against the URI.
pub struct UriRequest<'a> {
    this: &'a mut AptStream,
    uri: String,
    repo_uri: String,
    filename: String,
}

impl<'a> UriRequest<'a> {
    /// Attempts to parse the URI request from the message.
    async fn from(this: &'a mut AptStream, message: Message) -> Result<UriRequest<'a>, Error> {
        let Ok(uri) = message.header("URI") else {
            let msg = Message::uri_failure("", "Missing URI header").to_string();
            this.output_stream.write_all(msg.as_bytes()).await?;
            this.output_stream.as_mut().flush().await?;
            return Err(Error::HeaderNotFound("URI".to_string()));
        };

        let Ok(filename) = message.header("Filename") else {
            let msg = Message::uri_failure(uri, "Missing Filename header").to_string();
            this.output_stream.write_all(msg.as_bytes()).await?;
            this.output_stream.as_mut().flush().await?;
            return Err(Error::HeaderNotFound("Filename".to_string()));
        };

        let Ok(target_uri) = message.header("Target-Site") else {
            let msg = Message::uri_failure(uri, "Missing Target-Site header").to_string();
            this.output_stream.write_all(msg.as_bytes()).await?;
            this.output_stream.as_mut().flush().await?;
            return Err(Error::HeaderNotFound("Target-Site".to_string()));
        };

        Ok(UriRequest {
            this,
            uri: uri.to_string(),
            repo_uri: target_uri.to_string(),
            filename: filename.to_string(),
        })
    }

    /// Gets the URI for the request.
    pub fn uri(&self) -> &str {
        &self.uri
    }

    /// The filename to save the acquired URI to.
    pub fn filename(&self) -> &str {
        &self.filename
    }

    /// The originally configured URI for the repository.
    ///
    /// This is different from `uri()` in that it is the original
    /// URI configured in the APT sources, whereas `uri()` has the
    /// full URI that is being requested.
    pub fn repo_uri(&self) -> &str {
        &self.repo_uri
    }

    /// Sends a status (update message) to the APT client,
    /// indicating progress on the request.
    ///
    /// This does not consume the request, so multiple status
    /// updates can be sent for a single request.
    pub async fn send_status(&mut self, status: &str) -> Result<(), Error> {
        let msg = Message::status(status).to_string();
        self.this
            .output_stream
            .as_mut()
            .write_all(msg.as_bytes())
            .await?;
        self.this.output_stream.as_mut().flush().await?;
        Ok(())
    }

    /// Responds with a failure for the URI request.
    ///
    /// This immediately cancels the fetch and notifies APT of the failure.
    pub async fn fail(self, reason: &str) -> Result<(), Error> {
        let msg = Message::uri_failure(&self.uri, reason).to_string();
        self.this
            .output_stream
            .as_mut()
            .write_all(msg.as_bytes())
            .await?;
        self.this.output_stream.as_mut().flush().await?;
        Ok(())
    }

    /// Begins the response for the URI request.
    ///
    /// This tells APT that the initial metadata fetch was successful,
    /// and provides the size and last-modified time of the resource.
    pub async fn start(
        self,
        size_in_bytes: u64,
        last_modified: &str,
    ) -> Result<UriResponse<'a>, Error> {
        let msg = Message::uri_start(&self.uri, size_in_bytes, last_modified).to_string();
        self.this
            .output_stream
            .as_mut()
            .write_all(msg.as_bytes())
            .await?;
        self.this.output_stream.as_mut().flush().await?;
        Ok(UriResponse {
            this: self.this,
            uri: self.uri,
            filename: self.filename,
        })
    }
}

/// A builder for URI responses.
pub struct UriResponse<'a> {
    this: &'a mut AptStream,
    uri: String,
    filename: String,
}

impl UriResponse<'_> {
    /// Tells the APT client that the URI request has completed successfully.
    pub async fn complete(self) -> Result<(), Error> {
        let msg = Message::uri_success(&self.uri, &self.filename).to_string();
        self.this
            .output_stream
            .as_mut()
            .write_all(msg.as_bytes())
            .await?;
        self.this.output_stream.as_mut().flush().await?;
        Ok(())
    }

    /// Tells the APT client that the URI request has failed.
    pub async fn fail(self, reason: &str) -> Result<(), Error> {
        let msg = Message::uri_failure(&self.uri, reason).to_string();
        self.this
            .output_stream
            .as_mut()
            .write_all(msg.as_bytes())
            .await?;
        self.this.output_stream.as_mut().flush().await?;
        Ok(())
    }

    /// Creates a writer for writing the URI data to the target file.
    ///
    /// Creates the file with user read/write permissions only.
    pub async fn writer(&self) -> Result<impl tokio::io::AsyncWrite + 'static, Error> {
        Ok(tokio::fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .mode(0o600)
            .open(self.filename.clone())
            .await?)
    }
}

trait TryFromMessage<'a>: Sized {
    type Error;

    async fn try_from_message(msg: (&'a mut AptStream, Message)) -> Result<Self, Self::Error>;
}

impl<'a> TryFromMessage<'a> for AptRequest<'a> {
    type Error = Error;

    async fn try_from_message(
        (this, message): (&'a mut AptStream, Message),
    ) -> Result<Self, Self::Error> {
        match message.message_type {
            MessageType::Configuration => {
                let config_request = ConfigRequest::from(message);
                Ok(AptRequest::Configuration(config_request))
            }
            MessageType::URIAcquire => {
                let uri_request = UriRequest::from(this, message).await?;
                Ok(AptRequest::Uri(uri_request))
            }
            other => Err(Error::UnexpectedMessageType(other, message)),
        }
    }
}