cargo_doc_ngrok/
main.rs

1use std::{
2    io,
3    path::PathBuf,
4    process::Stdio,
5    sync::Arc,
6};
7
8use axum::BoxError;
9use clap::{
10    Args,
11    Parser,
12    Subcommand,
13};
14use futures::TryStreamExt;
15use hyper::service::service_fn;
16use hyper_util::{
17    rt::TokioExecutor,
18    server,
19};
20use ngrok::prelude::*;
21use watchexec::{
22    action::{
23        Action,
24        Outcome,
25    },
26    command::Command,
27    config::{
28        InitConfig,
29        RuntimeConfig,
30    },
31    error::CriticalError,
32    handler::PrintDebug,
33    signal::source::MainSignal,
34    Watchexec,
35};
36
37#[derive(Parser, Debug)]
38struct Cargo {
39    #[command(subcommand)]
40    cmd: Cmd,
41}
42
43#[derive(Debug, Subcommand)]
44enum Cmd {
45    DocNgrok(DocNgrok),
46}
47
48#[derive(Debug, Args)]
49struct DocNgrok {
50    #[arg(short)]
51    package: Option<String>,
52
53    #[arg(long, short)]
54    domain: Option<String>,
55
56    #[arg(long, short)]
57    watch: bool,
58
59    #[arg(last = true)]
60    doc_args: Vec<String>,
61}
62
63#[tokio::main]
64async fn main() -> Result<(), BoxError> {
65    let Cmd::DocNgrok(args) = Cargo::parse().cmd;
66
67    std::process::Command::new("cargo")
68        .arg("doc")
69        .args(args.doc_args.iter())
70        .stderr(Stdio::inherit())
71        .stdout(Stdio::inherit())
72        .spawn()?
73        .wait()?;
74
75    let meta = cargo_metadata::MetadataCommand::new().exec()?;
76
77    let default_package = args
78        .package
79        .or(meta.root_package().map(|p| p.name.clone()))
80        .ok_or("No default package found. You must provide one with -p")?;
81    let root_dir = meta.workspace_root;
82    let target_dir = meta.target_directory;
83    let doc_dir = target_dir.join("doc");
84
85    let sess = ngrok::Session::builder()
86        .authtoken_from_env()
87        .connect()
88        .await?;
89
90    let mut listen_cfg = sess.http_endpoint();
91    if let Some(domain) = args.domain {
92        listen_cfg.domain(domain);
93    }
94
95    let mut listener = listen_cfg.listen().await?;
96
97    let service = service_fn(move |req| {
98        let stat = hyper_staticfile::Static::new(&doc_dir);
99        stat.serve(req)
100    });
101
102    println!(
103        "serving docs on: {}/{}/",
104        listener.url(),
105        default_package.replace('-', "_")
106    );
107
108    let server = async move {
109        let (dropref, waiter) = awaitdrop::awaitdrop();
110
111        // Continuously accept new connections.
112        while let Some(conn) = listener.try_next().await? {
113            let service = service.clone();
114            let dropref = dropref.clone();
115            // Spawn a task to handle the connection. That way we can multiple connections
116            // concurrently.
117            tokio::spawn(async move {
118                if let Err(err) = server::conn::auto::Builder::new(TokioExecutor::new())
119                    .serve_connection(conn, service)
120                    .await
121                {
122                    eprintln!("failed to serve connection: {err:#}");
123                }
124                drop(dropref);
125            });
126        }
127
128        // Wait until all children have finished, not just the listener.
129        drop(dropref);
130        waiter.await;
131
132        Ok::<(), BoxError>(())
133    };
134
135    if args.watch {
136        let we = make_watcher(args.doc_args, root_dir, target_dir)?;
137
138        we.main().await??;
139    } else {
140        server.await?;
141    }
142
143    Ok(())
144}
145
146fn make_watcher(
147    args: Vec<String>,
148    root_dir: impl Into<PathBuf>,
149    target_dir: impl Into<PathBuf>,
150) -> Result<Arc<Watchexec>, Box<CriticalError>> {
151    let target_dir = target_dir.into();
152    let root_dir = root_dir.into();
153    let mut init = InitConfig::default();
154    init.on_error(PrintDebug(std::io::stderr()));
155
156    let mut runtime = RuntimeConfig::default();
157    runtime.pathset([root_dir]);
158    runtime.command(Command::Exec {
159        prog: "cargo".into(),
160        args: [String::from("doc")].into_iter().chain(args).collect(),
161    });
162    runtime.on_action({
163        move |action: Action| {
164            let target_dir = target_dir.clone();
165            async move {
166                let sigs = action
167                    .events
168                    .iter()
169                    .flat_map(|event| event.signals())
170                    .collect::<Vec<_>>();
171                if sigs.iter().any(|sig| sig == &MainSignal::Interrupt) {
172                    action.outcome(Outcome::Exit);
173                } else if action
174                    .events
175                    .iter()
176                    .any(|e| e.paths().any(|(p, _)| !p.starts_with(&target_dir)))
177                {
178                    action.outcome(Outcome::if_running(
179                        Outcome::both(Outcome::Stop, Outcome::Start),
180                        Outcome::Start,
181                    ));
182                }
183
184                Result::<_, io::Error>::Ok(())
185            }
186        }
187    });
188    Watchexec::new(init, runtime).map_err(Box::new)
189}