Skip to main content

ciplink_server/
server.rs

1//! The HTTP side (WebSocket `/ws`, `GET /health`, `/stats`, `/tags`) and the server that ties the hub and the PLC
2//! workers together.
3
4use std::collections::HashMap;
5use std::future::Future;
6use std::net::SocketAddr;
7use std::sync::Arc;
8use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
9
10use axum::Router;
11use axum::extract::ws::{CloseFrame, Message, WebSocket, WebSocketUpgrade};
12use axum::extract::{Query, State};
13use axum::http::{HeaderMap, Method, StatusCode, header};
14use axum::response::{IntoResponse, Response};
15use axum::routing::get;
16use futures_util::{SinkExt, StreamExt};
17use indexmap::IndexMap;
18use serde_json::{Map, Value as Json, json};
19use tokio::net::TcpListener;
20use tokio::sync::{mpsc, oneshot, watch};
21use tower_http::cors::{Any, CorsLayer};
22
23use crate::config::{Config, py_repr};
24use crate::hub::{FEATURES, Hub, Msg, Outbox, PROTOCOL, Shared};
25use crate::plc::Factory;
26use crate::rules::Rules;
27use crate::worker::{PlcWorker, WorkerOptions};
28
29#[derive(Clone)]
30struct AppState {
31    shared: Arc<Shared>,
32    tx: mpsc::UnboundedSender<Msg>,
33    next_id: Arc<AtomicU64>,
34    closing: watch::Receiver<bool>,
35}
36
37/// The shared state of a server on this config: its rules and a (not yet started) worker per PLC, whose events go
38/// to `tx`.
39pub(crate) fn build(
40    config: Config,
41    factory: Factory,
42    tx: &mpsc::UnboundedSender<Msg>,
43) -> Result<Arc<Shared>, String> {
44    let config = Arc::new(config);
45    let rules = Arc::new(Rules::new(&config)?);
46    let options = WorkerOptions {
47        cycle: config.cycle_ms as f64 / 1000.0,
48        tier2_period: config.tier2_period_s,
49        grace_s: config.grace_s,
50        planner: config.planner,
51        whole_max_bytes: config.whole_max_bytes,
52    };
53    let workers: IndexMap<String, Arc<PlcWorker>> = config
54        .plcs
55        .iter()
56        .map(|(name, cfg)| {
57            let events = tx.clone();
58            let on_event = Arc::new(move |event| {
59                let _ = events.send(Msg::Worker(event));
60            });
61            let worker = PlcWorker::new(
62                cfg.clone(),
63                Arc::clone(&factory),
64                options.clone(),
65                Arc::clone(&rules),
66                on_event,
67            );
68            (name.clone(), Arc::new(worker))
69        })
70        .collect();
71    Ok(Arc::new(Shared {
72        config,
73        rules,
74        workers,
75        sessions: AtomicUsize::new(0),
76        writes: AtomicU64::new(0),
77        write_rejects: AtomicU64::new(0),
78        momentary_timeouts: AtomicU64::new(0),
79    }))
80}
81
82/// Serves on `listener` until `shutdown` completes, then closes the WebSockets, releases every pressed button,
83/// stops the PLC workers (their queued writes still go out) and flushes the audit log.
84pub async fn serve(
85    config: Config,
86    factory: Factory,
87    listener: TcpListener,
88    shutdown: impl Future<Output = ()> + Send + 'static,
89) -> Result<(), String> {
90    let (tx, rx) = mpsc::unbounded_channel();
91    let shared = build(config, factory, &tx)?;
92    serve_built(shared, tx, rx, listener, shutdown).await
93}
94
95async fn serve_built(
96    shared: Arc<Shared>,
97    tx: mpsc::UnboundedSender<Msg>,
98    rx: mpsc::UnboundedReceiver<Msg>,
99    listener: TcpListener,
100    shutdown: impl Future<Output = ()> + Send + 'static,
101) -> Result<(), String> {
102    let hub = Hub::new(Arc::clone(&shared), tx.clone());
103    let hub_task = tokio::spawn(hub.run(rx));
104
105    let (closing_tx, closing) = watch::channel(false);
106    let state = AppState {
107        shared,
108        tx: tx.clone(),
109        next_id: Arc::new(AtomicU64::new(1)),
110        closing,
111    };
112    let cors = CorsLayer::new()
113        .allow_origin(Any)
114        .allow_methods([Method::GET])
115        .allow_headers([header::AUTHORIZATION]);
116    let app = Router::new()
117        .route("/health", get(health))
118        .route("/stats", get(stats))
119        .route("/tags", get(tags))
120        .route("/ws", get(ws))
121        .layer(cors)
122        .with_state(state);
123    let result = axum::serve(listener, app)
124        .with_graceful_shutdown(async move {
125            shutdown.await;
126            let _ = closing_tx.send(true); // the WebSockets close, so the graceful shutdown can finish
127        })
128        .await
129        .map_err(|e| e.to_string());
130    let (done_tx, done_rx) = oneshot::channel();
131    if tx.send(Msg::Stop(done_tx)).is_ok() {
132        let _ = done_rx.await;
133    }
134    let _ = hub_task.await;
135    result
136}
137
138/// Binds the config's host and port and serves until Ctrl+C.
139pub async fn run(config: Config, factory: Factory) -> Result<(), String> {
140    let listener = TcpListener::bind((config.host.as_str(), config.port))
141        .await
142        .map_err(|e| format!("cannot listen on {}:{}: {e}", config.host, config.port))?;
143    log::info!(
144        "serving on http://{}",
145        listener
146            .local_addr()
147            .map_or_else(|_| "?".into(), |a| a.to_string())
148    );
149    serve(config, factory, listener, async {
150        let _ = tokio::signal::ctrl_c().await;
151        log::info!("stopping");
152    })
153    .await
154}
155
156/// A server running on its own threads, for embedding (bindings, tests).
157pub struct ServerHandle {
158    addr: SocketAddr,
159    shared: Arc<Shared>,
160    stop: Option<oneshot::Sender<()>>,
161    thread: Option<std::thread::JoinHandle<Result<(), String>>>,
162}
163
164impl ServerHandle {
165    /// Starts a server with its own runtime; returns once it listens. Port 0 picks a free port.
166    pub fn start(config: Config, factory: Factory) -> Result<Self, String> {
167        let (ready_tx, ready_rx) = std::sync::mpsc::channel();
168        let (stop_tx, stop_rx) = oneshot::channel::<()>();
169        let (tx, rx) = mpsc::unbounded_channel();
170        let shared = build(config, factory, &tx)?;
171        let server = Arc::clone(&shared);
172        let thread = std::thread::Builder::new()
173            .name("ciplink-server".into())
174            .spawn(move || {
175                let runtime = tokio::runtime::Builder::new_multi_thread()
176                    .enable_all()
177                    .build()
178                    .map_err(|e| e.to_string())?;
179                runtime.block_on(async move {
180                    let (host, port) = (server.config.host.clone(), server.config.port);
181                    let listener = match TcpListener::bind((host.as_str(), port)).await {
182                        Ok(l) => l,
183                        Err(e) => {
184                            let error = format!("cannot listen on {host}:{port}: {e}");
185                            let _ = ready_tx.send(Err(error.clone()));
186                            return Err(error);
187                        }
188                    };
189                    let _ = ready_tx.send(listener.local_addr().map_err(|e| e.to_string()));
190                    serve_built(server, tx, rx, listener, async {
191                        let _ = stop_rx.await;
192                    })
193                    .await
194                })
195            })
196            .map_err(|e| e.to_string())?;
197        let addr = ready_rx
198            .recv()
199            .map_err(|_| "the server thread ended".to_owned())??;
200        Ok(ServerHandle {
201            addr,
202            shared,
203            stop: Some(stop_tx),
204            thread: Some(thread),
205        })
206    }
207
208    /// The server's shared state (config, rules, workers, counters).
209    #[doc(hidden)]
210    pub fn shared(&self) -> &Arc<Shared> {
211        &self.shared
212    }
213
214    /// The address the server listens on.
215    pub fn addr(&self) -> SocketAddr {
216        self.addr
217    }
218
219    /// Stops the server (releasing pressed buttons, flushing the audit log) and waits for it.
220    pub fn stop(mut self) -> Result<(), String> {
221        self.shutdown()
222    }
223
224    fn shutdown(&mut self) -> Result<(), String> {
225        if let Some(stop) = self.stop.take() {
226            let _ = stop.send(());
227        }
228        match self.thread.take() {
229            Some(thread) => thread
230                .join()
231                .map_err(|_| "the server thread panicked".to_owned())?,
232            None => Ok(()),
233        }
234    }
235}
236
237impl Drop for ServerHandle {
238    fn drop(&mut self) {
239        let _ = self.shutdown();
240    }
241}
242
243// ------------------------------------------------------------------------------------------------ HTTP
244
245fn detail(status: StatusCode, message: &str) -> Response {
246    (status, axum::Json(json!({"detail": message}))).into_response()
247}
248
249async fn health(State(state): State<AppState>) -> axum::Json<Json> {
250    let plcs: Map<String, Json> = state
251        .shared
252        .workers
253        .iter()
254        .map(|(name, w)| {
255            let stats = w.stats();
256            (
257                name.clone(),
258                json!({
259                    "state": w.state(),
260                    "last_cycle_ms": stats.last_cycle_ms,
261                    "last_ok_ts": stats.last_ok_ts,
262                    "reconnects": stats.reconnects,
263                    "reloads": stats.reloads,
264                }),
265            )
266        })
267        .collect();
268    axum::Json(json!({"plcs": plcs}))
269}
270
271async fn stats(State(state): State<AppState>) -> axum::Json<Json> {
272    let shared = &state.shared;
273    let plcs: Map<String, Json> = shared
274        .workers
275        .iter()
276        .map(|(name, w)| {
277            let stats = w.stats();
278            let mut entry = Map::new();
279            entry.insert("reads_per_cycle".into(), json!(stats.reads_per_cycle));
280            entry.insert("tags_per_cycle".into(), json!(stats.tags_per_cycle));
281            entry.insert("last_cycle_ms".into(), json!(stats.last_cycle_ms));
282            entry.insert("cycle_ms".into(), json!(shared.config.cycle_ms));
283            entry.insert("overruns".into(), json!(stats.overruns));
284            entry.insert("subscriptions".into(), json!(w.subscriptions()));
285            entry.extend(w.tier_counts());
286            entry.insert("group_fallbacks".into(), json!(stats.group_fallbacks));
287            entry.insert("groups".into(), Json::Object(w.group_stats()));
288            (name.clone(), Json::Object(entry))
289        })
290        .collect();
291    axum::Json(json!({
292        "sessions": shared.sessions.load(Ordering::Relaxed),
293        "writes": shared.writes.load(Ordering::Relaxed),
294        "write_rejects": shared.write_rejects.load(Ordering::Relaxed),
295        "momentary_timeouts": shared.momentary_timeouts.load(Ordering::Relaxed),
296        "plcs": plcs,
297    }))
298}
299
300async fn tags(
301    State(state): State<AppState>,
302    headers: HeaderMap,
303    Query(query): Query<HashMap<String, String>>,
304) -> Response {
305    let shared = &state.shared;
306    let secret = headers
307        .get(header::AUTHORIZATION)
308        .and_then(|v| v.to_str().ok())
309        .and_then(|v| v.strip_prefix("Bearer "));
310    if shared
311        .authenticate(secret.map(|s| Json::String(s.to_owned())).as_ref())
312        .is_none()
313    {
314        return detail(StatusCode::UNAUTHORIZED, "invalid token");
315    }
316    let plc = match query.get("plc") {
317        Some(plc) => plc.clone(),
318        None if shared.workers.len() == 1 => {
319            shared.workers.keys().next().cloned().unwrap_or_default()
320        }
321        None => return detail(StatusCode::NOT_FOUND, "unknown PLC None"),
322    };
323    let (Some(worker), Some(cfg)) = (shared.workers.get(&plc), shared.config.plcs.get(&plc)) else {
324        return detail(
325            StatusCode::NOT_FOUND,
326            &format!("unknown PLC {}", py_repr(&plc)),
327        );
328    };
329    let structure = worker.structure();
330    let rows = |tables: &[Map<String, Json>]| {
331        Json::Array(tables.iter().cloned().map(Json::Object).collect())
332    };
333    axum::Json(json!({
334        "plc": plc,
335        "state": worker.state(),
336        "version": worker.version(),
337        "tags": structure.tag_list,
338        "types": structure.types,
339        "rules": rows(&cfg.rules),
340        "groups": rows(&cfg.groups),
341        "decorators": shared.config.decorators,
342        "udt": cfg.udt,
343        "tag_settings": cfg.tags,
344    }))
345    .into_response()
346}
347
348// ------------------------------------------------------------------------------------------------ WebSocket
349
350/// Completes once the server is closing.
351async fn closed(mut closing: watch::Receiver<bool>) {
352    while !*closing.borrow_and_update() {
353        if closing.changed().await.is_err() {
354            return;
355        }
356    }
357}
358
359async fn ws(State(state): State<AppState>, upgrade: WebSocketUpgrade) -> Response {
360    upgrade.on_upgrade(move |socket| session(socket, state))
361}
362
363/// One client: `hello` first (else closed with 4401), then messages to the hub and the outbox to the client.
364async fn session(socket: WebSocket, state: AppState) {
365    let (mut sink, mut stream) = socket.split();
366    let hello = loop {
367        match stream.next().await {
368            Some(Ok(Message::Text(text))) => break serde_json::from_str::<Json>(&text).ok(),
369            Some(Ok(Message::Ping(_) | Message::Pong(_))) => {}
370            Some(Ok(Message::Binary(_))) => break None,
371            _ => return,
372        }
373    };
374    let token = hello
375        .filter(|h| h.get("op").and_then(Json::as_str) == Some("hello"))
376        .filter(|h| h.get("v").and_then(Json::as_u64) == Some(PROTOCOL))
377        .and_then(|h| state.shared.authenticate(h.get("token")));
378    let Some(token) = token else {
379        let frame = CloseFrame {
380            code: 4401,
381            reason: "".into(),
382        };
383        let _ = sink.send(Message::Close(Some(frame))).await;
384        return;
385    };
386    let id = state.next_id.fetch_add(1, Ordering::Relaxed);
387    let outbox = Arc::new(Outbox::default());
388    outbox.send(json!({
389        "op": "welcome",
390        "v": PROTOCOL,
391        "plcs": state.shared.states(),
392        "write": !token.write.is_empty(),
393        "features": FEATURES,
394    }));
395    if state
396        .tx
397        .send(Msg::Open(id, token, Arc::clone(&outbox)))
398        .is_err()
399    {
400        return;
401    }
402    let mut closing = Box::pin(closed(state.closing.clone()));
403    let sender = tokio::spawn(async move {
404        loop {
405            tokio::select! {
406                messages = outbox.next() => {
407                    for message in messages {
408                        if sink.send(Message::Text(message.to_string().into())).await.is_err() {
409                            return;
410                        }
411                    }
412                }
413                () = &mut closing => {
414                    let frame = CloseFrame { code: 1001, reason: "server stopping".into() };
415                    let _ = sink.send(Message::Close(Some(frame))).await;
416                    return;
417                }
418            }
419        }
420    });
421    let mut closing = Box::pin(closed(state.closing.clone()));
422    loop {
423        tokio::select! {
424            message = stream.next() => match message {
425                Some(Ok(Message::Text(text))) => {
426                    let _ = state.tx.send(Msg::Text(id, text.to_string()));
427                }
428                Some(Ok(Message::Close(_)) | Err(_)) | None => break,
429                Some(Ok(_)) => {}
430            },
431            () = &mut closing => break,
432        }
433    }
434    let _ = state.tx.send(Msg::Close(id));
435    if *state.closing.borrow() {
436        let _ = sender.await; // it sends the close frame
437    } else {
438        sender.abort();
439    }
440}