Skip to main content

cargo_doc_ngrok/
main.rs

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