1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
// Copyright (C) 2023 Andreas Hartmann <hartan@7x.de>
// GNU General Public License v3.0+ (https://www.gnu.org/licenses/gpl-3.0.txt)
// SPDX-License-Identifier: GPL-3.0-or-later

//! # Custom provider
//!
//! Adds integration to execute arbitrary commands as extra command providers. The main purpose of
//! this mechanism is to allow user-written scripts in any language to be used to enhance `cnf`s
//! default experience.
//!
//! Communication between `cnf` and custom providers takes place by passing messages via
//! stdin/stdout. The messages are JSON-formatted and must be terminated with a `BEL` char (ASCII
//! 7, 0x07), followed by a `\n` (newline) char (see [`MESSAGE_TERMINATOR`]). The search term is
//! passed as the first and only CLI argument to the custom provider.
//!
//! Child processes can run without restriction, there are no timeouts or similar measures. Cleanly
//! exiting from the custom provider executable can happen in one of two ways:
//!
//! - By sending a `CustomToCnf::Results` message with valid results to display
//! - By sending a `CustomToCnf::Error` message with an error description
//!
//! Any other type of exiting (regular exit, exit by signal) is recognized and will cause the
//! provider to report an appropriate error message.
//!
//!
//! ## Configuring custom providers
//!
//! Custom providers are registered through the application config file. On Linux systems, it is
//! found under `$XDG_CONFIG_DIR/cnf` (usually `~/.config/cnf`). Registering a provider looks like
//! this:
//!
//! ```yml
//! custom_providers:
//!     # A pretty name, displayed in the application
//!   - name: "cnf_fd (Bash)"
//!     # Main command to execute
//!     command: "/home/andi/Downloads/cnf_fd.sh"
//!     # Any additional arguments you may need
//!     args: []
//! ```
//!
//! You can configure an arbitrary amount of providers, just copy the snippet above and add more
//! list entries!
use crate::provider::prelude::*;
use async_process::{Command, Stdio};
use async_std::{
    io::{ReadExt, WriteExt},
    stream::StreamExt,
};
use logerr::LoggableError;
use serde_derive::{Deserialize, Serialize};

/// Error variants for custom providers.
///
/// These are returned to the application as `ProviderError::ApplicationError` variant and can be
/// accessed with `anyhow`s `downcast_*` functions.
#[derive(Debug, ThisError)]
#[non_exhaustive]
pub enum Error {
    #[error("failed to capture stdin of spawned child process")]
    NoStdin,

    #[error("failed to capture stdout of spawned child process")]
    NoStdout,

    #[error("failed to parse command output into string")]
    InvalidUtf8(#[from] std::string::FromUtf8Error),

    #[error("failed to write message to child process")]
    BrokenStdin(#[from] BrokenStdinError),

    #[error("failed to deserialize message from custom provider")]
    Deserialize(#[from] serde_json::Error),

    #[error("failed to read message from child process")]
    BrokenStdout(#[from] BrokenStdoutError),

    /// Miscellaneous error from the external provider executable
    #[error("{0}")]
    Child(String),
}

impl From<Error> for ProviderError {
    fn from(value: Error) -> Self {
        Self::ApplicationError(anyhow::Error::new(value))
    }
}

#[derive(Debug, ThisError)]
#[error(transparent)]
pub struct BrokenStdinError(#[from] std::io::Error);

#[derive(Debug, ThisError)]
#[error(transparent)]
pub struct BrokenStdoutError(#[from] std::io::Error);

/// Terminating byte sequence for messages passed between CNF and custom plugin.
///
/// We use this two-byte sequence for the following reasons:
///
/// 1. This sequence seems to be reasonably unlikely in regular shell output
/// 2. The terminating `\n` makes sure that messages can be received even by languages which have
///    no trivial way to read raw (unbuffered) stdin
/// 3. We can distinguish between a newline as part of the payload and the message termination
///    (under the assumption that the message payload doesn't hold the exact terminating sequence).
pub const MESSAGE_TERMINATOR: [u8; 2] = [0x07, b'\n'];

#[derive(Default, Debug, PartialEq, Clone, Deserialize, Serialize)]
pub struct Custom {
    /// A human-readable name/short identifier
    pub name: String,
    /// The main command to execute
    pub command: String,
    /// Additional arguments to provide
    pub args: Vec<String>,
}

impl fmt::Display for Custom {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "custom ({})", self.name)
    }
}

/// Messages from `cnf` to custom provider
#[derive(Debug, Serialize)]
#[serde(rename_all = "kebab-case")]
enum CnfToCustom {
    CommandResponse {
        stdout: String,
        stderr: String,
        exit_code: i32,
    },
}

/// Messages from custom provider to `cnf`
#[derive(Debug, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum CustomToCnf {
    Execute(CommandLine),

    Results(Vec<Candidate>),

    Error(String),
}

#[async_trait]
impl IsProvider for Custom {
    async fn search_internal(
        &self,
        command: &str,
        target_env: Arc<Environment>,
    ) -> ProviderResult<Vec<Candidate>> {
        let mut result: Vec<Candidate> = vec![];

        let mut child = Command::new(&self.command)
            .args(&self.args)
            .arg(command)
            .kill_on_drop(true)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null())
            .spawn()
            .map_err(|e| match e.kind() {
                std::io::ErrorKind::NotFound => {
                    ProviderError::Requirements(self.command.to_string())
                }
                _ => ProviderError::ApplicationError(anyhow::Error::new(e)),
            })?;

        let mut stdin = child.stdin.take().ok_or(Error::NoStdin)?;
        let mut stdout = child.stdout.take().ok_or(Error::NoStdout)?.bytes();
        let mut message: Vec<u8> = vec![];

        // TODO(hartan): Streams are free to return `None` at some point and continue afterwards
        // regardless. I'm not exactly sure how this stream behaves...
        while let Some(byte) = stdout.next().await {
            let byte = byte
                .map_err(BrokenStdoutError)
                .map_err(Error::BrokenStdout)?;
            if message
                .last()
                .map(|c| *c != MESSAGE_TERMINATOR[0])
                .unwrap_or(true)
                || (byte != MESSAGE_TERMINATOR[1])
            {
                message.push(byte);
                continue;
            };
            // Message terminated correctly
            message.truncate(message.len() - (MESSAGE_TERMINATOR.len() - 1));

            match serde_json::from_slice::<CustomToCnf>(&message)
                .map_err(Error::Deserialize)
                .with_context(|| format!("error communicating with '{}'", self.name))
                .to_log()?
            {
                CustomToCnf::Execute(commandline) => {
                    let output = target_env.output_of(commandline).await;

                    let message = match output {
                        Ok(stdout) => CnfToCustom::CommandResponse {
                            stdout,
                            stderr: "".to_string(),
                            exit_code: 0,
                        },
                        Err(ExecutionError::NonZero { output, .. }) => {
                            CnfToCustom::CommandResponse {
                                stdout: String::from_utf8(output.stdout)
                                    .map_err(Error::InvalidUtf8)?,
                                stderr: String::from_utf8(output.stderr)
                                    .map_err(Error::InvalidUtf8)?,
                                exit_code: output.status.code().unwrap_or(256),
                            }
                        }
                        _ => return output.map(|_| vec![]).map_err(ProviderError::from),
                    };
                    let mut response = serde_json::to_vec(&message)
                        .with_context(|| format!("failed to send response to provider {}", self))
                        .map_err(ProviderError::ApplicationError)?;
                    response.push(MESSAGE_TERMINATOR[0]);
                    response.push(MESSAGE_TERMINATOR[1]);
                    stdin
                        .write_all(&response)
                        .await
                        .map_err(BrokenStdinError)
                        .map_err(Error::BrokenStdin)?;
                }
                CustomToCnf::Results(results) => {
                    result = results;
                    break;
                }
                CustomToCnf::Error(error) => {
                    return Err(Error::Child(error).into());
                }
            }

            message.clear();
        }

        let _ = child
            .status()
            .await
            .with_context(|| format!("child process of '{}' terminated unexpectedly", self))
            .to_log();

        Ok(result)
    }
}