1use std::error;
5use std::fmt::Write as _;
6use std::sync::{Arc, LazyLock};
7use std::time::Duration;
8
9use axum::extract::State;
10use axum::http::{HeaderMap, StatusCode};
11use axum::response::{IntoResponse, Response};
12use axum::routing::{get, post};
13use axum::{Json, Router};
14use hmac::Hmac;
15use hmac::Mac;
16use openraft::BasicNode;
17use openraft::error::ClientWriteError;
18use openraft::error::{
19 InstallSnapshotError, NetworkError, RPCError, RaftError, RemoteError, Unreachable,
20};
21use openraft::network::{RPCOption, RaftNetwork, RaftNetworkFactory};
22use openraft::raft::{
23 AppendEntriesRequest, AppendEntriesResponse, InstallSnapshotRequest, InstallSnapshotResponse,
24 VoteRequest, VoteResponse,
25};
26use reqwest::header::HeaderValue;
27use serde::Serialize;
28use serde::de::DeserializeOwned;
29use subtle::ConstantTimeEq;
30
31use super::{NodeId, Raft, Reply, Request, TypeConfig};
32
33const SECRET_HEADER: &str = "x-orch-cluster-secret";
34const NONCE_HEADER: &str = "x-orch-nonce";
35const PROOF_HEADER: &str = "x-orch-proof";
36
37static CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
38 reqwest::Client::builder()
39 .timeout(Duration::from_secs(30))
40 .build()
41 .unwrap_or_default()
42});
43
44#[derive(Clone)]
45struct App {
46 raft: Raft,
47 secret: Arc<str>,
48 me: Arc<Hello>,
49 addr: Arc<str>,
51}
52
53#[derive(Debug, Clone, PartialEq, Serialize, serde::Deserialize)]
55pub struct Hello {
56 pub name: String,
57 pub id: NodeId,
58 pub witness: bool,
59 #[serde(default)]
60 pub leader: Option<NodeId>,
61 #[serde(default)]
62 pub term: u64,
63 #[serde(default)]
65 pub proof: String,
66}
67
68fn mac(secret: &str, parts: &[&str]) -> String {
70 let mut m =
71 Hmac::<sha2::Sha256>::new_from_slice(secret.as_bytes()).expect("HMAC takes any key");
72 for part in parts {
73 m.update(part.as_bytes());
74 m.update(b"\0");
75 }
76
77 let bytes = m.finalize().into_bytes();
78 let mut out = String::with_capacity(bytes.len() * 2);
79 for b in bytes {
80 let _ = write!(out, "{b:02x}");
81 }
82 out
83}
84
85#[derive(Debug, Clone, Serialize, serde::Deserialize)]
87pub struct Join {
88 pub id: NodeId,
89 pub addr: String,
90}
91
92fn authorized(app: &App, headers: &HeaderMap) -> bool {
93 let given = headers
94 .get(SECRET_HEADER)
95 .map(HeaderValue::as_bytes)
96 .unwrap_or_default();
97 given.ct_eq(app.secret.as_bytes()).into()
98}
99
100pub(crate) fn router(raft: Raft, secret: Arc<str>, me: Hello, addr: String) -> Router {
101 Router::new()
102 .route("/raft/append", post(append))
103 .route("/raft/snapshot", post(snapshot))
104 .route("/raft/vote", post(vote))
105 .route("/cluster/write", post(write))
106 .route("/cluster/hello", get(hello))
107 .route("/cluster/join", post(join))
108 .with_state(App {
109 raft,
110 secret,
111 me: Arc::new(me),
112 addr: Arc::from(addr),
113 })
114}
115
116async fn hello(State(a): State<App>, headers: HeaderMap) -> Response {
118 let header = |name: &str| {
119 headers
120 .get(name)
121 .and_then(|v| v.to_str().ok())
122 .unwrap_or_default()
123 .to_string()
124 };
125 let nonce = header(NONCE_HEADER);
126 let expected = mac(&a.secret, &["probe", &nonce, &a.addr]);
127 let proven: bool = header(PROOF_HEADER)
128 .as_bytes()
129 .ct_eq(expected.as_bytes())
130 .into();
131 if nonce.len() < 16 || !proven {
132 return StatusCode::UNAUTHORIZED.into_response();
133 }
134
135 let m = a.raft.metrics().borrow().clone();
136 Json(Hello {
137 leader: m.current_leader,
138 term: m.current_term,
139 proof: mac(&a.secret, &["hello", &nonce, &a.me.name, &a.addr]),
140 ..(*a.me).clone()
141 })
142 .into_response()
143}
144
145async fn join(State(a): State<App>, headers: HeaderMap, Json(req): Json<Join>) -> Response {
147 if !authorized(&a, &headers) {
148 return StatusCode::UNAUTHORIZED.into_response();
149 }
150
151 let node = BasicNode::new(&req.addr);
152 match a.raft.add_learner(req.id, node, false).await {
153 Ok(_) => StatusCode::OK.into_response(),
154 Err(RaftError::APIError(ClientWriteError::ForwardToLeader(f))) => match f.leader_node {
155 Some(leader) => match send_join(&leader.addr, &a.secret, &req).await {
156 Ok(()) => StatusCode::OK.into_response(),
157 Err(e) => (StatusCode::BAD_GATEWAY, e.to_string()).into_response(),
158 },
159 None => (StatusCode::SERVICE_UNAVAILABLE, "no leader yet").into_response(),
160 },
161 Err(e) => (StatusCode::SERVICE_UNAVAILABLE, e.to_string()).into_response(),
162 }
163}
164
165pub async fn send_join(addr: &str, secret: &str, req: &Join) -> anyhow::Result<()> {
166 let res = CLIENT
167 .post(format!("http://{addr}/cluster/join"))
168 .header(SECRET_HEADER, secret)
169 .json(req)
170 .send()
171 .await?;
172 anyhow::ensure!(
173 res.status().is_success(),
174 "join refused by {addr} ({}): {}",
175 res.status(),
176 res.text().await?
177 );
178 Ok(())
179}
180
181pub async fn probe(addr: &str, secret: &str, expect_name: Option<&str>) -> Option<Hello> {
185 let nonce = uuid::Uuid::new_v4().simple().to_string();
186 let res = CLIENT
187 .get(format!("http://{addr}/cluster/hello"))
188 .header(NONCE_HEADER, &nonce)
189 .header(PROOF_HEADER, mac(secret, &["probe", &nonce, addr]))
190 .timeout(Duration::from_secs(1))
191 .send()
192 .await
193 .ok()?;
194 let hello: Hello = res.error_for_status().ok()?.json().await.ok()?;
195 if expect_name.is_some_and(|name| name != hello.name) {
196 return None;
197 }
198
199 let expected = mac(secret, &["hello", &nonce, &hello.name, addr]);
200 let proven: bool = hello.proof.as_bytes().ct_eq(expected.as_bytes()).into();
201 proven.then_some(hello)
202}
203
204async fn append(
205 State(a): State<App>,
206 headers: HeaderMap,
207 Json(req): Json<AppendEntriesRequest<TypeConfig>>,
208) -> Response {
209 if !authorized(&a, &headers) {
210 return StatusCode::UNAUTHORIZED.into_response();
211 }
212
213 Json(a.raft.append_entries(req).await).into_response()
214}
215
216async fn snapshot(
217 State(a): State<App>,
218 headers: HeaderMap,
219 Json(req): Json<InstallSnapshotRequest<TypeConfig>>,
220) -> Response {
221 if !authorized(&a, &headers) {
222 return StatusCode::UNAUTHORIZED.into_response();
223 }
224
225 Json(a.raft.install_snapshot(req).await).into_response()
226}
227
228async fn vote(
229 State(a): State<App>,
230 headers: HeaderMap,
231 Json(req): Json<VoteRequest<NodeId>>,
232) -> Response {
233 if !authorized(&a, &headers) {
234 return StatusCode::UNAUTHORIZED.into_response();
235 }
236
237 Json(a.raft.vote(req).await).into_response()
238}
239
240async fn write(State(a): State<App>, headers: HeaderMap, Json(req): Json<Request>) -> Response {
242 if !authorized(&a, &headers) {
243 return StatusCode::UNAUTHORIZED.into_response();
244 }
245
246 if let Err(e) = req.cmd.check_replicable() {
247 return (StatusCode::BAD_REQUEST, e.to_string()).into_response();
248 }
249
250 match a.raft.client_write(req).await {
251 Ok(r) => Json(r.data).into_response(),
252 Err(e @ RaftError::APIError(ClientWriteError::ForwardToLeader(_))) => {
254 (StatusCode::MISDIRECTED_REQUEST, e.to_string()).into_response()
255 }
256 Err(e) => (StatusCode::SERVICE_UNAVAILABLE, e.to_string()).into_response(),
257 }
258}
259
260#[derive(Debug)]
262pub(crate) enum ForwardError {
263 NotApplied(String),
265 Unknown(anyhow::Error),
267}
268
269pub(crate) async fn forward(
270 addr: &str,
271 secret: &str,
272 req: &Request,
273) -> Result<Reply, ForwardError> {
274 let res = CLIENT
275 .post(format!("http://{addr}/cluster/write"))
276 .header(SECRET_HEADER, secret)
277 .json(req)
278 .send()
279 .await
280 .map_err(|e| {
281 if e.is_connect() {
282 ForwardError::NotApplied(e.to_string())
283 } else {
284 ForwardError::Unknown(e.into())
285 }
286 })?;
287 let status = res.status();
288 if status == StatusCode::MISDIRECTED_REQUEST {
289 let body = res.text().await.unwrap_or_default();
290 return Err(ForwardError::NotApplied(format!(
291 "{addr} is not the leader: {body}"
292 )));
293 }
294
295 if !status.is_success() {
296 let body = res.text().await.unwrap_or_default();
297 return Err(ForwardError::Unknown(anyhow::anyhow!(
298 "leader at {addr} refused the write ({status}): {body}"
299 )));
300 }
301
302 res.json()
303 .await
304 .map_err(|e| ForwardError::Unknown(e.into()))
305}
306
307pub(crate) struct HttpNetwork {
308 pub secret: Arc<str>,
309}
310
311impl RaftNetworkFactory<TypeConfig> for HttpNetwork {
312 type Network = Peer;
313
314 async fn new_client(&mut self, target: NodeId, node: &BasicNode) -> Peer {
315 Peer {
316 target,
317 addr: node.addr.clone(),
318 secret: Arc::clone(&self.secret),
319 }
320 }
321}
322
323pub(crate) struct Peer {
324 target: NodeId,
325 addr: String,
326 secret: Arc<str>,
327}
328
329impl Peer {
330 async fn send<Req, Resp, E>(
331 &self,
332 path: &str,
333 req: Req,
334 option: &RPCOption,
335 ) -> Result<Resp, RPCError<NodeId, BasicNode, RaftError<NodeId, E>>>
336 where
337 Req: Serialize,
338 Resp: DeserializeOwned,
339 E: error::Error + DeserializeOwned,
340 {
341 let res = CLIENT
342 .post(format!("http://{}/raft/{path}", self.addr))
343 .header(SECRET_HEADER, &*self.secret)
344 .timeout(option.hard_ttl())
345 .json(&req)
346 .send()
347 .await
348 .map_err(|e| {
349 if e.is_connect() {
350 RPCError::Unreachable(Unreachable::new(&e))
351 } else {
352 RPCError::Network(NetworkError::new(&e))
353 }
354 })?;
355 let res = res
356 .error_for_status()
357 .map_err(|e| RPCError::Network(NetworkError::new(&e)))?;
358 let body: Result<Resp, RaftError<NodeId, E>> = res
359 .json()
360 .await
361 .map_err(|e| RPCError::Network(NetworkError::new(&e)))?;
362 body.map_err(|e| RPCError::RemoteError(RemoteError::new(self.target, e)))
363 }
364}
365
366impl RaftNetwork<TypeConfig> for Peer {
367 async fn append_entries(
368 &mut self,
369 req: AppendEntriesRequest<TypeConfig>,
370 option: RPCOption,
371 ) -> Result<AppendEntriesResponse<NodeId>, RPCError<NodeId, BasicNode, RaftError<NodeId>>> {
372 self.send("append", req, &option).await
373 }
374
375 async fn install_snapshot(
376 &mut self,
377 req: InstallSnapshotRequest<TypeConfig>,
378 option: RPCOption,
379 ) -> Result<
380 InstallSnapshotResponse<NodeId>,
381 RPCError<NodeId, BasicNode, RaftError<NodeId, InstallSnapshotError>>,
382 > {
383 self.send("snapshot", req, &option).await
384 }
385
386 async fn vote(
387 &mut self,
388 req: VoteRequest<NodeId>,
389 option: RPCOption,
390 ) -> Result<VoteResponse<NodeId>, RPCError<NodeId, BasicNode, RaftError<NodeId>>> {
391 self.send("vote", req, &option).await
392 }
393}