Skip to main content

dre_protocol/
host.rs

1//! Core's side of the protocol: spawn a plugin, negotiate a version, send requests.
2//!
3//! Frames are read on a background thread so core can time out a silent plugin during the
4//! handshake and never blocks forever on a plugin that has died. stderr is drained on another
5//! thread, forwarded to a log sink and kept (last lines) for error messages.
6
7use std::collections::VecDeque;
8use std::io::{BufRead, BufReader, BufWriter, Write};
9use std::path::{Path, PathBuf};
10use std::process::{Child, ChildStdin, Command, ExitStatus, Stdio};
11use std::sync::mpsc::{Receiver, RecvTimeoutError, channel};
12use std::sync::{Arc, Mutex, OnceLock};
13use std::time::Duration;
14
15use arrow::array::RecordBatch;
16use arrow::datatypes::SchemaRef;
17use serde_json::{Map, Value};
18
19use crate::frame::{self, Frame, FrameError};
20use crate::msg::{ConnectionField, DeliveryFile, Request, Response, ResultSetMeta};
21use crate::options::OptionField;
22use crate::{Kind, MAX_VERSION, MIN_VERSION, PluginId, parse_executable_name};
23
24/// Receives each stderr line a plugin writes.
25pub type LogSink = Arc<dyn Fn(&str, &str) + Send + Sync>;
26
27/// A log sink that prefixes each plugin line with the plugin's file name, on core's stderr.
28pub fn stderr_log() -> LogSink {
29    Arc::new(|plugin, line| eprintln!("[{plugin}] {line}"))
30}
31
32pub const DEFAULT_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30);
33
34/// Set to `1` in a plugin's environment when a person is at core's terminal, else `0`.
35pub const INTERACTIVE_ENV: &str = "DRE_INTERACTIVE";
36const STDERR_TAIL: usize = 20;
37
38#[derive(Debug, Clone, PartialEq)]
39pub struct PluginInfo {
40    pub protocol_version: u32,
41    pub kind: Kind,
42    pub name: String,
43    pub version: String,
44    pub capabilities: Vec<String>,
45    /// Every plugin the executable serves (at least the one it's serving now).
46    pub provides: Vec<PluginId>,
47}
48
49/// A plugin's `describe` reply.
50#[derive(Debug, Clone, PartialEq, Default)]
51pub struct Description {
52    pub connection_fields: Vec<ConnectionField>,
53    pub option_fields: Vec<OptionField>,
54    /// A source's identifier quote character.
55    pub identifier_quote: Option<String>,
56}
57
58#[derive(Debug)]
59pub enum HostError {
60    Spawn {
61        plugin: String,
62        error: std::io::Error,
63    },
64    Crashed {
65        plugin: String,
66        status: Option<ExitStatus>,
67        stderr: Vec<String>,
68    },
69    Malformed {
70        plugin: String,
71        message: String,
72    },
73    Incompatible {
74        plugin: String,
75        core: (u32, u32),
76        plugin_range: (u32, u32),
77    },
78    Timeout {
79        plugin: String,
80        waiting_for: &'static str,
81    },
82    Unexpected {
83        plugin: String,
84        expected: &'static str,
85        got: String,
86    },
87    /// The plugin reported an error for a request.
88    Plugin {
89        plugin: String,
90        message: String,
91    },
92    Arrow {
93        plugin: String,
94        message: String,
95    },
96}
97
98impl std::fmt::Display for HostError {
99    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        match self {
101            HostError::Spawn { plugin, error } => write!(f, "can't start plugin `{plugin}`: {error}"),
102            HostError::Crashed {
103                plugin,
104                status,
105                stderr,
106            } => {
107                match status {
108                    Some(s) => write!(f, "plugin `{plugin}` exited unexpectedly ({s})")?,
109                    None => write!(f, "plugin `{plugin}` closed its output unexpectedly")?,
110                }
111                if !stderr.is_empty() {
112                    write!(f, "; its last log lines:\n  {}", stderr.join("\n  "))?;
113                }
114                Ok(())
115            }
116            HostError::Malformed { plugin, message } => {
117                write!(f, "plugin `{plugin}` broke the protocol: {message}")
118            }
119            HostError::Incompatible {
120                plugin,
121                core,
122                plugin_range,
123            } => write!(
124                f,
125                "plugin `{plugin}` speaks protocol versions {}..={}, but this DRE core speaks {}..={}; update {}",
126                plugin_range.0,
127                plugin_range.1,
128                core.0,
129                core.1,
130                if plugin_range.1 < core.0 {
131                    "the plugin"
132                } else {
133                    "DRE"
134                }
135            ),
136            HostError::Timeout { plugin, waiting_for } => {
137                write!(
138                    f,
139                    "plugin `{plugin}` didn't answer in time (waiting for {waiting_for})"
140                )
141            }
142            HostError::Unexpected {
143                plugin,
144                expected,
145                got,
146            } => {
147                write!(f, "plugin `{plugin}` sent {got} where {expected} was expected")
148            }
149            HostError::Plugin { message, .. } => f.write_str(message),
150            HostError::Arrow { plugin, message } => {
151                write!(f, "plugin `{plugin}` sent invalid Arrow data: {message}")
152            }
153        }
154    }
155}
156
157impl std::error::Error for HostError {}
158
159pub type Result<T> = std::result::Result<T, HostError>;
160
161/// A message from the plugin.
162pub enum Incoming {
163    Json(Response),
164    Arrow(Vec<u8>),
165}
166
167impl std::fmt::Debug for Incoming {
168    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
169        match self {
170            Incoming::Json(r) => write!(f, "Json({r:?})"),
171            Incoming::Arrow(b) => write!(f, "Arrow frame ({} bytes)", b.len()),
172        }
173    }
174}
175
176/// The outcome of a source `execute`.
177#[derive(Debug)]
178pub enum Execution {
179    NoResult { rows_affected: Option<u64> },
180    Result { schema: SchemaRef, rows: u64 },
181}
182
183pub struct PluginProcess {
184    label: String,
185    path: PathBuf,
186    child: Child,
187    stdin: Option<BufWriter<ChildStdin>>,
188    rx: Receiver<std::result::Result<Frame, FrameError>>,
189    stderr: Arc<Mutex<VecDeque<String>>>,
190    /// Forwards stderr; joined on close/drop so no log line is lost when the plugin exits.
191    stderr_thread: Option<std::thread::JoinHandle<()>>,
192    info: Option<PluginInfo>,
193    /// A reply that arrived while core was still sending (a format failing mid-stream).
194    early: Option<std::result::Result<Frame, FrameError>>,
195    /// The plugin asked for in the handshake. By default the one a `dre-<kind>-<name>` file name
196    /// names; none for a package executable, which then serves its first plugin.
197    serve: Option<PluginId>,
198}
199
200impl PluginProcess {
201    /// Spawn the plugin at `path` and complete the handshake.
202    pub fn start(path: &Path, log: LogSink) -> Result<PluginProcess> {
203        Self::start_in(path, log, None)
204    }
205
206    /// Like `start`, running the plugin in `cwd` (core uses the project directory).
207    pub fn start_in(path: &Path, log: LogSink, cwd: Option<&Path>) -> Result<PluginProcess> {
208        Self::start_for(path, None, log, cwd)
209    }
210
211    /// Like `start_in`, asking the executable for `plugin` (one of a package's plugins).
212    pub fn start_for(
213        path: &Path,
214        plugin: Option<&PluginId>,
215        log: LogSink,
216        cwd: Option<&Path>,
217    ) -> Result<PluginProcess> {
218        let timeout = std::env::var("DRE_PLUGIN_HANDSHAKE_TIMEOUT_MS")
219            .ok()
220            .and_then(|v| v.parse().ok())
221            .map(Duration::from_millis)
222            .unwrap_or(DEFAULT_HANDSHAKE_TIMEOUT);
223        let mut p = Self::spawn_in(path, log, &[], cwd)?;
224        if let Some(id) = plugin {
225            p.serve = Some(id.clone());
226        }
227        p.handshake((MIN_VERSION, MAX_VERSION), timeout)?;
228        Ok(p)
229    }
230
231    /// Spawn and handshake offering an explicit version range (the conformance suite uses this).
232    pub fn start_with(
233        path: &Path,
234        log: LogSink,
235        versions: (u32, u32),
236        timeout: Duration,
237    ) -> Result<PluginProcess> {
238        let mut p = Self::spawn(path, log)?;
239        p.handshake(versions, timeout)?;
240        Ok(p)
241    }
242
243    /// Spawn without a handshake (for tests that probe raw protocol behaviour).
244    pub fn spawn(path: &Path, log: LogSink) -> Result<PluginProcess> {
245        Self::spawn_env(path, log, &[])
246    }
247
248    /// Spawn with extra environment variables, without a handshake.
249    pub fn spawn_env(path: &Path, log: LogSink, env: &[(&str, &str)]) -> Result<PluginProcess> {
250        Self::spawn_in(path, log, env, None)
251    }
252
253    /// Spawn with extra environment variables and a working directory, without a handshake.
254    pub fn spawn_in(
255        path: &Path,
256        log: LogSink,
257        env: &[(&str, &str)],
258        cwd: Option<&Path>,
259    ) -> Result<PluginProcess> {
260        let label = path
261            .file_name()
262            .map(|f| f.to_string_lossy().to_string())
263            .unwrap_or_default();
264        let mut cmd = Command::new(path);
265        if let Some(dir) = cwd {
266            cmd.current_dir(dir);
267        }
268        // Plugins can't see the terminal (their stdio is piped), so tell them whether a person is
269        // there, e.g. to finish a browser sign-in. An existing DRE_INTERACTIVE is left alone.
270        if std::env::var_os(INTERACTIVE_ENV).is_none() {
271            use std::io::IsTerminal;
272            let there = std::io::stderr().is_terminal() && std::io::stdin().is_terminal();
273            cmd.env(INTERACTIVE_ENV, if there { "1" } else { "0" });
274        }
275        let mut child = cmd
276            .envs(env.iter().copied())
277            .stdin(Stdio::piped())
278            .stdout(Stdio::piped())
279            .stderr(Stdio::piped())
280            .spawn()
281            .map_err(|error| HostError::Spawn {
282                plugin: label.clone(),
283                error,
284            })?;
285        let stdout = child.stdout.take().unwrap();
286        let stderr_pipe = child.stderr.take().unwrap();
287        let stdin = child.stdin.take().map(BufWriter::new);
288
289        let (tx, rx) = channel();
290        std::thread::spawn(move || {
291            let mut r = BufReader::new(stdout);
292            loop {
293                let f = frame::read_frame(&mut r);
294                let stop = f.is_err();
295                if tx.send(f).is_err() || stop {
296                    break;
297                }
298            }
299        });
300        let stderr = Arc::new(Mutex::new(VecDeque::new()));
301        let tail = stderr.clone();
302        let who = label.clone();
303        let stderr_thread = std::thread::spawn(move || {
304            for line in BufReader::new(stderr_pipe).lines() {
305                let Ok(line) = line else { break };
306                log(&who, &line);
307                let mut t = tail.lock().unwrap();
308                if t.len() == STDERR_TAIL {
309                    t.pop_front();
310                }
311                t.push_back(line);
312            }
313        });
314        let serve = parse_executable_name(&label).map(|(k, n)| PluginId::new(k, n));
315        Ok(PluginProcess {
316            label,
317            path: path.to_path_buf(),
318            child,
319            stdin,
320            rx,
321            stderr,
322            stderr_thread: Some(stderr_thread),
323            info: None,
324            early: None,
325            serve,
326        })
327    }
328
329    /// Name the plugin to ask for in the handshake (`None`: a package's first).
330    pub fn ask_for(&mut self, plugin: Option<PluginId>) {
331        self.serve = plugin;
332    }
333
334    /// Complete the handshake on a process from `spawn`/`spawn_env`.
335    pub fn handshake(&mut self, (min, max): (u32, u32), timeout: Duration) -> Result<()> {
336        self.send(&Request::Hello {
337            min_version: min,
338            max_version: max,
339            core_version: core_version(),
340            plugin: self.serve.clone(),
341        })?;
342        match self.recv(Some(timeout), "the hello reply")? {
343            Incoming::Json(Response::Hello {
344                protocol_version,
345                kind,
346                name,
347                version,
348                capabilities,
349                mut provides,
350            }) => {
351                if protocol_version < min || protocol_version > max {
352                    return Err(HostError::Incompatible {
353                        plugin: self.label.clone(),
354                        core: (min, max),
355                        plugin_range: (protocol_version, protocol_version),
356                    });
357                }
358                if let Some(want) = &self.serve
359                    && (want.kind != kind || want.name != name)
360                {
361                    return Err(HostError::Plugin {
362                        plugin: self.label.clone(),
363                        message: format!(
364                            "`{}` was asked for the {want} plugin but serves {kind}/{name}",
365                            self.label
366                        ),
367                    });
368                }
369                if provides.is_empty() {
370                    provides.push(PluginId::new(kind, name.clone()));
371                }
372                self.info = Some(PluginInfo {
373                    protocol_version,
374                    kind,
375                    name,
376                    version,
377                    capabilities,
378                    provides,
379                });
380                Ok(())
381            }
382            Incoming::Json(Response::VersionMismatch {
383                min_version,
384                max_version,
385            }) => Err(HostError::Incompatible {
386                plugin: self.label.clone(),
387                core: (min, max),
388                plugin_range: (min_version, max_version),
389            }),
390            Incoming::Json(Response::Error { message }) => Err(HostError::Plugin {
391                plugin: self.label.clone(),
392                message,
393            }),
394            other => Err(self.unexpected("a hello reply", &other)),
395        }
396    }
397
398    pub fn label(&self) -> &str {
399        &self.label
400    }
401
402    pub fn path(&self) -> &Path {
403        &self.path
404    }
405
406    pub fn info(&self) -> &PluginInfo {
407        self.info.as_ref().expect("handshake completed")
408    }
409
410    pub fn has(&self, capability: &str) -> bool {
411        self.info
412            .as_ref()
413            .is_some_and(|i| i.capabilities.iter().any(|c| c == capability))
414    }
415
416    fn unexpected(&self, expected: &'static str, got: &Incoming) -> HostError {
417        let got = match got {
418            Incoming::Json(r) => format!("{:?}", r)
419                .split([' ', '{'])
420                .next()
421                .unwrap_or("?")
422                .to_string(),
423            Incoming::Arrow(_) => "Arrow data".into(),
424        };
425        HostError::Unexpected {
426            plugin: self.label.clone(),
427            expected,
428            got: format!("`{got}`"),
429        }
430    }
431
432    fn crashed(&mut self) -> HostError {
433        // Give the process a moment to finish exiting so its status and last stderr are known.
434        let mut status = None;
435        for _ in 0..50 {
436            if let Ok(Some(s)) = self.child.try_wait() {
437                status = Some(s);
438                break;
439            }
440            std::thread::sleep(Duration::from_millis(20));
441        }
442        std::thread::sleep(Duration::from_millis(50));
443        let stderr = self.stderr.lock().unwrap().iter().cloned().collect();
444        HostError::Crashed {
445            plugin: self.label.clone(),
446            status,
447            stderr,
448        }
449    }
450
451    pub fn send(&mut self, req: &Request) -> Result<()> {
452        let ok = match self.stdin.as_mut() {
453            Some(w) => frame::write_json(w, req).is_ok(),
454            None => false,
455        };
456        if ok { Ok(()) } else { Err(self.write_failed()) }
457    }
458
459    pub fn send_batch(&mut self, batch: &RecordBatch) -> Result<()> {
460        // A plugin that fails part-way replies before reading the rest: stop sending and
461        // surface its error rather than streaming everything into a plugin that gave up.
462        if self.early.is_none()
463            && let Ok(f) = self.rx.try_recv()
464        {
465            self.early = Some(f);
466        }
467        if let Some(Ok(Frame::Json(v))) = &self.early
468            && let Ok(Response::Error { message }) = serde_json::from_value::<Response>(v.clone())
469        {
470            return Err(HostError::Plugin {
471                plugin: self.label.clone(),
472                message,
473            });
474        }
475        let ipc = frame::encode_batch(batch).map_err(|e| HostError::Arrow {
476            plugin: self.label.clone(),
477            message: e.to_string(),
478        })?;
479        let ok = match self.stdin.as_mut() {
480            Some(w) => frame::write_arrow(w, &ipc).is_ok(),
481            None => false,
482        };
483        if ok { Ok(()) } else { Err(self.write_failed()) }
484    }
485
486    /// Writing failed: the plugin has gone away. Prefer its own error message if it sent one.
487    fn write_failed(&mut self) -> HostError {
488        match self.rx.recv_timeout(Duration::from_secs(2)) {
489            Ok(Ok(Frame::Json(v))) => match serde_json::from_value::<Response>(v) {
490                Ok(Response::Error { message }) => HostError::Plugin {
491                    plugin: self.label.clone(),
492                    message,
493                },
494                _ => self.crashed(),
495            },
496            _ => self.crashed(),
497        }
498    }
499
500    /// Receive the next message. `timeout` of `None` waits as long as the plugin is alive.
501    pub fn recv(&mut self, timeout: Option<Duration>, waiting_for: &'static str) -> Result<Incoming> {
502        let got = match timeout {
503            _ if self.early.is_some() => self.early.take(),
504            Some(t) => match self.rx.recv_timeout(t) {
505                Ok(f) => Some(f),
506                Err(RecvTimeoutError::Timeout) => {
507                    return Err(HostError::Timeout {
508                        plugin: self.label.clone(),
509                        waiting_for,
510                    });
511                }
512                Err(RecvTimeoutError::Disconnected) => None,
513            },
514            None => self.rx.recv().ok(),
515        };
516        match got {
517            Some(Ok(Frame::Json(v))) => match serde_json::from_value::<Response>(v.clone()) {
518                Ok(r) => Ok(Incoming::Json(r)),
519                Err(e) => Err(HostError::Malformed {
520                    plugin: self.label.clone(),
521                    message: format!("unknown message {v}: {e}"),
522                }),
523            },
524            Some(Ok(Frame::Arrow(b))) => Ok(Incoming::Arrow(b)),
525            Some(Err(FrameError::Malformed(m))) => Err(HostError::Malformed {
526                plugin: self.label.clone(),
527                message: m,
528            }),
529            Some(Err(_)) | None => Err(self.crashed()),
530        }
531    }
532
533    /// Receive a JSON response, turning `error` into `HostError::Plugin`.
534    pub fn recv_json(&mut self, waiting_for: &'static str) -> Result<Response> {
535        match self.recv(None, waiting_for)? {
536            Incoming::Json(Response::Error { message }) => Err(HostError::Plugin {
537                plugin: self.label.clone(),
538                message,
539            }),
540            Incoming::Json(r) => Ok(r),
541            other => Err(self.unexpected(waiting_for, &other)),
542        }
543    }
544
545    /// Send a request and expect `ok`.
546    fn call_ok(&mut self, req: &Request, what: &'static str) -> Result<()> {
547        self.send(req)?;
548        match self.recv_json(what)? {
549            Response::Ok {} => Ok(()),
550            other => Err(self.unexpected(what, &Incoming::Json(other))),
551        }
552    }
553
554    pub fn describe(&mut self) -> Result<Vec<ConnectionField>> {
555        Ok(self.describe_all()?.0)
556    }
557
558    /// The connection fields and the options the plugin declares.
559    pub fn describe_all(&mut self) -> Result<(Vec<ConnectionField>, Vec<OptionField>)> {
560        let d = self.description()?;
561        Ok((d.connection_fields, d.option_fields))
562    }
563
564    /// The whole `describe` reply.
565    pub fn description(&mut self) -> Result<Description> {
566        self.send(&Request::Describe {})?;
567        match self.recv_json("a describe reply")? {
568            Response::Describe {
569                connection_fields,
570                option_fields,
571                identifier_quote,
572            } => Ok(Description {
573                connection_fields,
574                option_fields,
575                identifier_quote,
576            }),
577            other => Err(self.unexpected("a describe reply", &Incoming::Json(other))),
578        }
579    }
580
581    /// Check a config block of options (needs `validate`); returns every problem found.
582    pub fn validate(&mut self, options: Map<String, Value>) -> Result<Vec<String>> {
583        self.send(&Request::Validate { options })?;
584        match self.recv_json("a validate reply")? {
585            Response::Validated { errors } => Ok(errors),
586            other => Err(self.unexpected("a validate reply", &Incoming::Json(other))),
587        }
588    }
589
590    pub fn open(&mut self, connection: Map<String, Value>, read_only: bool) -> Result<()> {
591        self.call_ok(
592            &Request::Open {
593                connection,
594                read_only,
595            },
596            "an open reply",
597        )
598    }
599
600    pub fn check(&mut self, sql: &str) -> Result<()> {
601        self.call_ok(&Request::Check { sql: sql.to_string() }, "a check reply")
602    }
603
604    /// Run one statement, handing every batch to `on_batch` as it arrives.
605    pub fn execute(
606        &mut self,
607        sql: &str,
608        row_limit: Option<u64>,
609        mut on_batch: impl FnMut(&SchemaRef, RecordBatch) -> std::result::Result<(), String>,
610    ) -> Result<Execution> {
611        self.send(&Request::Execute {
612            sql: sql.to_string(),
613            row_limit,
614        })?;
615        match self.recv_json("an execute reply")? {
616            Response::NoResult { rows_affected } => return Ok(Execution::NoResult { rows_affected }),
617            Response::Result { .. } => {}
618            other => return Err(self.unexpected("an execute reply", &Incoming::Json(other))),
619        }
620        let mut schema: Option<SchemaRef> = None;
621        let mut rows = 0u64;
622        let mut sink_error = None;
623        loop {
624            match self.recv(None, "result data")? {
625                Incoming::Arrow(ipc) => {
626                    let (s, batches) = frame::decode_batches(&ipc).map_err(|e| HostError::Arrow {
627                        plugin: self.label.clone(),
628                        message: e.to_string(),
629                    })?;
630                    let schema = schema.get_or_insert(s);
631                    for b in batches {
632                        rows += b.num_rows() as u64;
633                        if sink_error.is_none()
634                            && let Err(e) = on_batch(schema, b)
635                        {
636                            // Keep draining so the session stays usable; report after.
637                            sink_error = Some(e);
638                        }
639                    }
640                }
641                Incoming::Json(Response::ResultEnd { .. }) => break,
642                Incoming::Json(Response::Error { message }) => {
643                    return Err(HostError::Plugin {
644                        plugin: self.label.clone(),
645                        message,
646                    });
647                }
648                other => return Err(self.unexpected("result data", &other)),
649            }
650        }
651        if let Some(e) = sink_error {
652            return Err(HostError::Plugin {
653                plugin: self.label.clone(),
654                message: e,
655            });
656        }
657        let schema = schema.ok_or_else(|| HostError::Malformed {
658            plugin: self.label.clone(),
659            message: "a result ended without any Arrow frame carrying its schema".into(),
660        })?;
661        Ok(Execution::Result { schema, rows })
662    }
663
664    /// Start a format `write`; follow with `write_result_set` per result set, then `write_finish`.
665    pub fn write_begin(
666        &mut self,
667        path: &str,
668        format: &str,
669        options: Map<String, Value>,
670        result_sets: Vec<ResultSetMeta>,
671        template: Option<Value>,
672    ) -> Result<()> {
673        self.send(&Request::Write {
674            path: path.to_string(),
675            format: format.to_string(),
676            options,
677            result_sets,
678            template,
679        })
680    }
681
682    /// Stream one result set: at least one batch (carrying the schema), then the end marker.
683    pub fn write_result_set(
684        &mut self,
685        schema: &SchemaRef,
686        batches: impl IntoIterator<Item = RecordBatch>,
687    ) -> Result<()> {
688        let mut any = false;
689        for b in batches {
690            any = true;
691            self.send_batch(&b)?;
692        }
693        if !any {
694            self.send_batch(&RecordBatch::new_empty(schema.clone()))?;
695        }
696        self.send(&Request::ResultSetEnd {})
697    }
698
699    /// Load rows into a temporary table on the source's session. Returns the relation to use in
700    /// SQL, and the plugin's warning, if any. Only for plugins advertising `load`.
701    pub fn load(
702        &mut self,
703        name: &str,
704        schema: &SchemaRef,
705        batches: impl IntoIterator<Item = RecordBatch>,
706    ) -> Result<(String, Option<String>)> {
707        self.send(&Request::Load {
708            name: name.to_string(),
709        })?;
710        self.write_result_set(schema, batches)?;
711        match self.recv_json("a loaded reply")? {
712            Response::Loaded {
713                relation, warning, ..
714            } => Ok((relation, warning)),
715            other => Err(self.unexpected("a loaded reply", &Incoming::Json(other))),
716        }
717    }
718
719    pub fn write_finish(&mut self) -> Result<Vec<String>> {
720        Ok(self.write_finish_with_warnings()?.0)
721    }
722
723    /// `write_finish`, also returning the format's warnings.
724    pub fn write_finish_with_warnings(&mut self) -> Result<(Vec<String>, Vec<String>)> {
725        self.send(&Request::Finish {})?;
726        match self.recv_json("a written reply")? {
727            Response::Written { files, warnings } => Ok((files, warnings)),
728            other => Err(self.unexpected("a written reply", &Incoming::Json(other))),
729        }
730    }
731
732    /// Deliver one file with no plugin options.
733    pub fn deliver(
734        &mut self,
735        local_path: &str,
736        remote_path: Option<&str>,
737        connection: Map<String, Value>,
738    ) -> Result<String> {
739        let file = DeliveryFile {
740            local_path: local_path.to_string(),
741            remote_path: remote_path.map(str::to_string),
742        };
743        self.deliver_files(&[file], connection, Map::new())
744    }
745
746    /// Deliver `files` in one request with the destination entry's `options`. More than one file
747    /// needs a plugin advertising `multi_file`; call once per file otherwise.
748    pub fn deliver_files(
749        &mut self,
750        files: &[DeliveryFile],
751        connection: Map<String, Value>,
752        options: Map<String, Value>,
753    ) -> Result<String> {
754        let req = match files {
755            [] => {
756                return Err(HostError::Malformed {
757                    plugin: self.label.clone(),
758                    message: "nothing to deliver".into(),
759                });
760            }
761            [one] => Request::Deliver {
762                local_path: Some(one.local_path.clone()),
763                remote_path: one.remote_path.clone(),
764                files: Vec::new(),
765                connection,
766                options,
767            },
768            many => {
769                if !self.has(crate::CAP_MULTI_FILE) {
770                    return Err(HostError::Plugin {
771                        plugin: self.label.clone(),
772                        message: format!(
773                            "{} files in one delivery, but the plugin doesn't advertise `multi_file`",
774                            many.len()
775                        ),
776                    });
777                }
778                Request::Deliver {
779                    local_path: None,
780                    remote_path: None,
781                    files: many.to_vec(),
782                    connection,
783                    options,
784                }
785            }
786        };
787        self.send(&req)?;
788        match self.recv_json("a delivered reply")? {
789            Response::Delivered { location } => Ok(location),
790            other => Err(self.unexpected("a delivered reply", &Incoming::Json(other))),
791        }
792    }
793
794    /// Ask the plugin to exit, and wait for it.
795    pub fn close(mut self) -> Result<()> {
796        let _ = self.send(&Request::Close {});
797        let _ = self.recv(Some(Duration::from_secs(5)), "the close reply");
798        self.stdin.take();
799        for _ in 0..250 {
800            if let Ok(Some(_)) = self.child.try_wait() {
801                return Ok(());
802            }
803            std::thread::sleep(Duration::from_millis(20));
804        }
805        let _ = self.child.kill();
806        Ok(())
807    }
808
809    /// Raw access for protocol tests: write arbitrary bytes to the plugin.
810    pub fn write_raw(&mut self, bytes: &[u8]) -> Result<()> {
811        let ok = match self.stdin.as_mut() {
812            Some(w) => w.write_all(bytes).and_then(|_| w.flush()).is_ok(),
813            None => false,
814        };
815        if ok { Ok(()) } else { Err(self.write_failed()) }
816    }
817
818    /// Wait up to `timeout` for the process to exit.
819    pub fn wait_exit(&mut self, timeout: Duration) -> Option<ExitStatus> {
820        self.stdin.take();
821        let deadline = std::time::Instant::now() + timeout;
822        while std::time::Instant::now() < deadline {
823            if let Ok(Some(s)) = self.child.try_wait() {
824                return Some(s);
825            }
826            std::thread::sleep(Duration::from_millis(20));
827        }
828        None
829    }
830}
831
832impl PluginProcess {
833    /// Wait (briefly) for the stderr forwarder to reach end of stream.
834    fn drain_stderr(&mut self) {
835        if let Some(t) = self.stderr_thread.take() {
836            let deadline = std::time::Instant::now() + Duration::from_secs(2);
837            while !t.is_finished() && std::time::Instant::now() < deadline {
838                std::thread::sleep(Duration::from_millis(5));
839            }
840            if t.is_finished() {
841                let _ = t.join();
842            }
843        }
844    }
845}
846
847impl Drop for PluginProcess {
848    fn drop(&mut self) {
849        self.stdin.take();
850        if let Ok(None) = self.child.try_wait() {
851            // Give it a moment to exit on its own (stdin is closed), then insist.
852            for _ in 0..25 {
853                if let Ok(Some(_)) = self.child.try_wait() {
854                    break;
855                }
856                std::thread::sleep(Duration::from_millis(10));
857            }
858            let _ = self.child.kill();
859            let _ = self.child.wait();
860        }
861        self.drain_stderr();
862    }
863}
864
865static CORE_VERSION: OnceLock<String> = OnceLock::new();
866
867/// Set the DRE version sent to plugins in the handshake (`core_version`). This crate is versioned
868/// on its own, so the host names the version of DRE it's part of; unset, it's "unreleased".
869pub fn set_core_version(version: &str) {
870    let _ = CORE_VERSION.set(version.to_string());
871}
872
873fn core_version() -> String {
874    CORE_VERSION.get().cloned().unwrap_or_else(|| "unreleased".into())
875}