KeyBoxen 0.1.0

// Copyright (C) 2022 KeyBoxen Authors
// SPDX-License-Identifier: GPL-3.0-or-later

//! KeyBoxen to Pinentry communication (Assuan requests)
//!
//! This is run in a separate task. This module deals with sending the
//! specific commands to the pinentry process via stdin in assuan
//! format.

use super::{Command, Result};
use tokio::{
    io::{AsyncWriteExt, BufWriter},
    process::ChildStdin,
    sync::mpsc::Receiver,
};

/// Pinentry transmit loop
///
/// This function is executed as a separate async task to listen for
/// commands from the [controller][super::controller] task and forward it to pinentry
/// subprocess in assuan format. This is the only interface for the
/// transmit loop.
///
/// # Arguments
/// * `stream`  : The stdin of the pinentry process
/// * `channel` : The receiver side of the channel that is used by the
/// controller to pass commands to the loop
pub(super) async fn transmit<'a>(stream: ChildStdin, mut channel: Receiver<Command>) -> Result<()> {
    let mut writer = BufWriter::new(stream);
    loop {
        match channel.recv().await {
            // Return when channel tx is closed
            None => {
                println!("🛈 Pinentry: Command channel closed");
                break;
            }
            Some(cmd) => {
                let lines: Vec<Vec<u8>> = cmd.into();
                for mut line in lines {
                    line.push(b'\n');
                    writer.write(&line).await?;
                    writer.flush().await?;
                }
            }
        }
    }
    Ok(())
}

/// Automatic conversion of Command to 2D byte vector
///
/// This trait implementation abstracts the conversion of commands to
/// 2D byte vector away from the [transmit] function.
impl From<Command> for Vec<Vec<u8>> {
    fn from(cmd: Command) -> Self {
        match cmd {
            Command::Title(s) => CmdType::Text("SETTITLE", s),
            Command::Description(s) => CmdType::Text("SETDESC", s),
            Command::Prompt(s) => CmdType::Text("SETPROMPT", s),
            Command::OKButton(s) => CmdType::Text("SETOK", s),
            Command::CancelButton(s) => CmdType::Text("SETCANCEL", s),
            Command::NotOKButton(s) => CmdType::Text("SETNOTOK", s),
            Command::TimeOut(t) => CmdType::Number("SETTIMEOUT", t),
            Command::Data(data) => CmdType::Data("D", data),
            Command::End => CmdType::Plain("END"),
            Command::Cancel => CmdType::Plain("CAN"),
            Command::Password => CmdType::Plain("GETPIN"),
            Command::Confirm => CmdType::Plain("CONFIRM"),
            Command::Notify => CmdType::Plain("MESSAGE"),
            Command::Exit => CmdType::Plain("BYE"),
        }
        .into()
    }
}

/// Reduced set of command types
///
/// This enum reduces the commands into a set of broadly similar
/// classes based on the payload type. Each type can then be
/// serialized automatically.
enum CmdType {
    /// Format: Plain(cmd)
    Plain(&'static str),
    /// Format: Text(cmd, text)
    Text(&'static str, String),
    /// Format: Number(cmd, number)
    Number(&'static str, i64),
    /// Format: Data(cmd, text)
    Data(&'static str, Vec<u8>),
}

/// Automatic conversion of CmdType to 2D byte vector
///
/// The conversion of a [Command] to a 2D byte vector consists of two
/// steps:   
/// 1. [Command] to [CmdType] mapping (orthogonalization)
/// 2. CmdType to 2D byte vector
///
/// This implementation abstracts the second step mentioned
/// above.
impl From<CmdType> for Vec<Vec<u8>> {
    fn from(data: CmdType) -> Self {
        match data {
            CmdType::Text(cmd, text) => text.export(cmd),
            CmdType::Data(cmd, data) => data.segment(cmd),
            CmdType::Plain(cmd) => ().export(cmd),
            CmdType::Number(cmd, num) => num.export(cmd),
        }
    }
}

/// Assuan escape for any binary vector
///
/// This trait enables the use of method syntax for escaping binary
/// data as required by Assuan protocol.
trait AssuanEscape {
    fn escape(self) -> Vec<u8>;
}

impl<'a> AssuanEscape for &'a [u8] {
    fn escape(self) -> Vec<u8> {
        let mut dest = Vec::with_capacity(self.len());
        for byte in self {
            match byte {
                b'%' => dest.extend_from_slice(b"%25"),
                b'\r' => dest.extend_from_slice(b"%0D"),
                b'\n' => dest.extend_from_slice(b"%0A"),
                _ => dest.push(*byte),
            }
        }
        dest
    }
}

/// Assuan segmentation for binary data
///
/// This trait is used for large data that may not fit in one Assuan
/// command. The data is segmented to a maximum size of 1000,
/// including the leading command and trailing linefeed.
trait AssuanSegmented {
    fn segment(self, cmd: &'static str) -> Vec<Vec<u8>>;
}

impl<T: Into<Vec<u8>>> AssuanSegmented for T {
    fn segment(self, cmd: &'static str) -> Vec<Vec<u8>> {
        const MAX_CHUNK_SIZE: usize = 999;
        let cmd = format!("{cmd} ").into_bytes();
        let sub_chunk_size = MAX_CHUNK_SIZE - cmd.len();
        self.into()
            .escape()
            .chunks(sub_chunk_size)
            .map(|chunk| cmd.iter().chain(chunk).map(|x| *x).collect())
            .collect()
    }
}

/// Trait for converting single line commands to byte vector
///
/// This trait is used for those command lines that don't need
/// segmentation. It converts the data into 2D byte array.
trait AssuanSingleLine {
    fn export(&self, cmd: &'static str) -> Vec<Vec<u8>>;
}

impl AssuanSingleLine for () {
    fn export(&self, cmd: &'static str) -> Vec<Vec<u8>> {
        vec![format!("{cmd}").into_bytes()]
    }
}

impl AssuanSingleLine for i64 {
    fn export(&self, cmd: &'static str) -> Vec<Vec<u8>> {
        vec![format!("{cmd} {self}").into_bytes()]
    }
}

impl AssuanSingleLine for String {
    fn export(&self, cmd: &'static str) -> Vec<Vec<u8>> {
        let mut line = format!("{cmd} ").into_bytes();
        line.extend_from_slice(&self.as_bytes().escape());
        vec![line]
    }
}