1use std::path::PathBuf;
4use std::sync::{Arc, LazyLock};
5use std::time::Duration;
6
7use axum::extract::State;
8use axum::http::{header, HeaderMap, StatusCode};
9use axum::response::{IntoResponse, Response};
10use axum::routing::{get, post};
11use axum::{Json, Router};
12use subtle::ConstantTimeEq;
13
14use super::{Command, Reply};
15use crate::store::TaskStore;
16
17static CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
18 reqwest::Client::builder()
19 .timeout(Duration::from_secs(30))
20 .build()
21 .unwrap_or_default()
22});
23
24#[derive(Debug, Clone)]
26pub struct Endpoint {
27 pub url: String,
28 pub token: String,
29}
30
31fn port_path() -> anyhow::Result<PathBuf> {
32 if cfg!(test) {
33 Ok(std::env::temp_dir()
34 .join("orch-test-state")
35 .join("cluster.port"))
36 } else {
37 Ok(crate::home::state_dir()?.join("cluster.port"))
38 }
39}
40
41#[derive(Clone)]
42struct App {
43 store: TaskStore,
44 token: Arc<String>,
45 host: Arc<String>,
46}
47
48pub async fn serve(store: TaskStore) -> anyhow::Result<Endpoint> {
50 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
51 let port = listener.local_addr()?.port();
52 let token = format!(
53 "{}{}",
54 uuid::Uuid::new_v4().simple(),
55 uuid::Uuid::new_v4().simple()
56 );
57 write_private(
58 &port_path()?,
59 &format!("{port} {token} {}", std::process::id()),
60 )?;
61 let app = Router::new()
62 .route("/write", post(write))
63 .route("/progress", post(progress))
64 .route("/nodes", get(nodes))
65 .with_state(App {
66 store,
67 token: Arc::new(token.clone()),
68 host: Arc::new(format!("127.0.0.1:{port}")),
69 });
70 tokio::spawn(async move {
71 if let Err(e) = axum::serve(listener, app).await {
72 tracing::warn!(err = %e, "cluster write endpoint stopped");
73 }
74 });
75 tracing::info!(port, "cluster write endpoint started on 127.0.0.1");
76 Ok(Endpoint {
77 url: format!("http://127.0.0.1:{port}/write"),
78 token,
79 })
80}
81
82fn write_private(path: &std::path::Path, body: &str) -> anyhow::Result<()> {
84 use std::io::Write;
85 use std::os::unix::fs::OpenOptionsExt;
86 if let Some(parent) = path.parent() {
87 std::fs::create_dir_all(parent)?;
88 }
89
90 let tmp = path.with_extension(format!("tmp.{}", std::process::id()));
91 let mut file = std::fs::OpenOptions::new()
92 .write(true)
93 .create(true)
94 .truncate(true)
95 .mode(0o600)
96 .open(&tmp)?;
97 file.write_all(body.as_bytes())?;
98 std::fs::rename(&tmp, path)?;
99 Ok(())
100}
101
102fn authorized(app: &App, headers: &HeaderMap) -> bool {
103 let host_ok = headers
105 .get(header::HOST)
106 .is_some_and(|h| h.as_bytes() == app.host.as_bytes());
107 let given = headers
108 .get(header::AUTHORIZATION)
109 .and_then(|v| v.to_str().ok())
110 .and_then(|v| v.strip_prefix("Bearer "))
111 .unwrap_or("");
112 let token_ok: bool = given.as_bytes().ct_eq(app.token.as_bytes()).into();
113 host_ok && token_ok
114}
115
116async fn write(State(app): State<App>, headers: HeaderMap, Json(cmd): Json<Command>) -> Response {
117 if !authorized(&app, &headers) {
118 return StatusCode::UNAUTHORIZED.into_response();
119 }
120
121 if let Err(e) = cmd.check_replicable() {
123 return Json(Reply::Err(format!("{e:#}"))).into_response();
124 }
125
126 Json(match app.store.write(cmd).await {
127 Ok(reply) => reply,
128 Err(e) => Reply::Err(format!("{e:#}")),
129 })
130 .into_response()
131}
132
133async fn progress(
134 State(app): State<App>,
135 headers: HeaderMap,
136 Json(report): Json<super::progress::Report>,
137) -> Response {
138 if !authorized(&app, &headers) {
139 return StatusCode::UNAUTHORIZED.into_response();
140 }
141
142 let node = crate::cluster::discovery::status()
143 .await
144 .map(|t| t.me.name)
145 .unwrap_or_else(|_| "local".to_string());
146 match super::progress::handle(&app.store, &node, report).await {
147 Ok(record) => Json(record).into_response(),
148 Err(e) => (StatusCode::BAD_REQUEST, format!("{e:#}")).into_response(),
149 }
150}
151
152async fn nodes(State(app): State<App>, headers: HeaderMap) -> Response {
153 if !authorized(&app, &headers) {
154 return StatusCode::UNAUTHORIZED.into_response();
155 }
156
157 let members = app.store.node().map(super::members).unwrap_or_default();
158 Json(members).into_response()
159}
160
161pub async fn fetch_nodes(endpoint: &Endpoint) -> anyhow::Result<Vec<super::Member>> {
162 let url = endpoint.url.replace("/write", "/nodes");
163 let res = CLIENT.get(url).bearer_auth(&endpoint.token).send().await?;
164 Ok(res.error_for_status()?.json().await?)
165}
166
167pub async fn send_progress(
169 endpoint: &Endpoint,
170 report: &super::progress::Report,
171) -> anyhow::Result<super::progress::Record> {
172 let url = endpoint.url.replace("/write", "/progress");
173 let res = CLIENT
174 .post(url)
175 .bearer_auth(&endpoint.token)
176 .json(report)
177 .send()
178 .await?;
179 let status = res.status();
180 if !status.is_success() {
181 anyhow::bail!(
182 "orch service refused the report ({status}): {}",
183 res.text().await?
184 );
185 }
186
187 Ok(res.json().await?)
188}
189
190pub async fn cleanup() {
191 if let Ok(path) = port_path() {
192 let _ = tokio::fs::remove_file(path).await;
193 }
194}
195
196pub async fn daemon_endpoint() -> Option<Endpoint> {
198 let body = tokio::fs::read_to_string(port_path().ok()?).await.ok()?;
199 let mut parts = body.split_whitespace();
200 let port: u16 = parts.next()?.parse().ok()?;
201 let token = parts.next()?.to_string();
202 let pid: u32 = parts.next()?.parse().ok()?;
203 let alive = tokio::task::spawn_blocking(move || crate::home::pid_is_orch_serve(pid))
204 .await
205 .ok()?;
206 alive.then(|| Endpoint {
207 url: format!("http://127.0.0.1:{port}/write"),
208 token,
209 })
210}
211
212pub(crate) async fn send(endpoint: &Endpoint, cmd: &Command) -> reqwest::Result<Reply> {
213 CLIENT
214 .post(&endpoint.url)
215 .bearer_auth(&endpoint.token)
216 .json(cmd)
217 .send()
218 .await?
219 .error_for_status()?
220 .json()
221 .await
222}