1use std::collections::HashMap;
16use std::io::{BufReader, Read, Write};
17use std::net::{Ipv4Addr, SocketAddr, TcpListener, TcpStream};
18use std::path::PathBuf;
19use std::sync::atomic::{AtomicBool, Ordering};
20use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender, channel};
21use std::sync::{Arc, Mutex, PoisonError};
22use std::time::{Duration, Instant};
23
24use crate::ControlError;
25use crate::key::{KeyFile, key_path, random_key, write_key_file};
26use crate::lines::{Auth, DEFAULT_MAX_LINE, Source, serve_lines};
27use crate::protocol::{Reply, Request};
28
29#[derive(Debug, Clone)]
31pub struct Config {
32 pub port: u16,
34 pub key_dir: PathBuf,
36 pub max_line: usize,
38 pub max_connections: usize,
40 pub first_request: Duration,
42 pub reply_timeout: Duration,
44}
45
46impl Config {
47 pub fn new(port: u16, key_dir: impl Into<PathBuf>) -> Self {
48 Self {
49 port,
50 key_dir: key_dir.into(),
51 max_line: DEFAULT_MAX_LINE,
52 max_connections: 16,
53 first_request: Duration::from_secs(30),
54 reply_timeout: Duration::from_secs(60),
55 }
56 }
57}
58
59pub type Wake = Arc<dyn Fn() + Send + Sync>;
61
62pub struct Server {
65 port: u16,
66 key: Arc<str>,
67 key_file: KeyFile,
68 stop: Arc<AtomicBool>,
69 conns: Arc<Mutex<Conns>>,
70}
71
72#[derive(Default)]
74struct Conns {
75 next: u64,
76 open: HashMap<u64, TcpStream>,
77}
78
79impl std::fmt::Debug for Server {
80 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81 f.debug_struct("Server").field("port", &self.port).field("key_file", &self.key_file).finish_non_exhaustive()
82 }
83}
84
85impl Server {
86 pub fn start(config: Config, wake: impl Fn() + Send + Sync + 'static) -> Result<(Server, Receiver<Request>), ControlError> {
90 let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, config.port))
91 .map_err(|e| ControlError::Io(format!("cannot bind 127.0.0.1:{}: {e}", config.port)))?;
92 let port = listener.local_addr().map_err(|e| ControlError::Io(e.to_string()))?.port();
93 let key: Arc<str> = random_key().map_err(ControlError::Key)?.into();
94 let path = key_path(&config.key_dir, port);
95 write_key_file(&path, &key).map_err(ControlError::Key)?;
96 let key_file = KeyFile::new(path);
97 log::info!("control server listening on 127.0.0.1:{port}, key in {}", key_file.path().display());
98 let (tx, rx) = channel();
99 let stop = Arc::new(AtomicBool::new(false));
100 let ctx = Ctx { tx, wake: Arc::new(wake), key: key.clone(), config: Arc::new(config) };
101 let conns = Arc::new(Mutex::new(Conns::default()));
102 let (stop2, conns2) = (stop.clone(), conns.clone());
103 std::thread::Builder::new()
104 .name("control-accept".into())
105 .spawn(move || accept_loop(listener, ctx, &stop2, &conns2))
106 .map_err(|e| ControlError::Io(format!("cannot start the control thread: {e}")))?;
107 Ok((Server { port, key, key_file, stop, conns }, rx))
108 }
109
110 pub fn port(&self) -> u16 {
111 self.port
112 }
113
114 pub fn key(&self) -> &str {
116 &self.key
117 }
118
119 pub fn key_file(&self) -> &std::path::Path {
120 self.key_file.path()
121 }
122}
123
124impl Drop for Server {
125 fn drop(&mut self) {
126 self.stop.store(true, Ordering::SeqCst);
127 for s in self.conns.lock().unwrap_or_else(PoisonError::into_inner).open.values() {
128 let _ = s.shutdown(std::net::Shutdown::Both);
129 }
130 let _ = TcpStream::connect_timeout(&SocketAddr::from((Ipv4Addr::LOCALHOST, self.port)), Duration::from_millis(200));
132 }
133}
134
135#[derive(Clone)]
136struct Ctx {
137 tx: Sender<Request>,
138 wake: Wake,
139 key: Arc<str>,
140 config: Arc<Config>,
141}
142
143fn accept_loop(listener: TcpListener, ctx: Ctx, stop: &AtomicBool, conns: &Arc<Mutex<Conns>>) {
144 for stream in listener.incoming() {
145 if stop.load(Ordering::SeqCst) {
146 return;
147 }
148 let Ok(mut stream) = stream else { continue };
149 let slot = {
150 let mut c = conns.lock().unwrap_or_else(PoisonError::into_inner);
151 match stream.try_clone() {
152 Ok(handle) if c.open.len() < ctx.config.max_connections => {
153 let id = c.next;
154 c.next = c.next.wrapping_add(1);
155 c.open.insert(id, handle);
156 Some(Slot(id, Arc::clone(conns)))
157 }
158 _ => None,
159 }
160 };
161 let Some(slot) = slot else {
164 let _ = stream.set_write_timeout(Some(Duration::from_millis(200)));
165 let _ = stream.write_all(format!("{}\n", Reply::err("too many control connections").to_json(None)).as_bytes());
166 continue;
167 };
168 let ctx = ctx.clone();
169 let spawned = std::thread::Builder::new().name("control-conn".into()).spawn(move || {
170 let _slot = slot;
171 serve(stream, &ctx);
172 });
173 if let Err(e) = spawned {
174 log::error!("control server could not start a connection thread: {e}");
175 }
176 }
177}
178
179struct Slot(u64, Arc<Mutex<Conns>>);
181
182impl Drop for Slot {
183 fn drop(&mut self) {
184 self.1.lock().unwrap_or_else(PoisonError::into_inner).open.remove(&self.0);
185 }
186}
187
188struct Timed {
191 s: TcpStream,
192 deadline: Option<Instant>,
193}
194
195impl Read for Timed {
196 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
197 if let Some(d) = self.deadline {
198 let left = d.checked_duration_since(Instant::now()).filter(|l| !l.is_zero()).ok_or(std::io::ErrorKind::TimedOut)?;
199 self.s.set_read_timeout(Some(left))?;
200 }
201 self.s.read(buf)
202 }
203}
204
205impl Source for Timed {
206 fn set_deadline(&mut self, d: Option<Instant>) {
207 self.deadline = d;
208 if d.is_none() {
209 let _ = self.s.set_read_timeout(None);
210 }
211 }
212}
213
214fn serve(stream: TcpStream, ctx: &Ctx) {
215 let Ok(read) = stream.try_clone() else { return };
216 let _ = stream.set_nodelay(true);
217 let _ = stream.set_write_timeout(Some(Duration::from_secs(30)));
219 let deadline = Instant::now().checked_add(ctx.config.first_request);
220 let mut reader = BufReader::new(Timed { s: read, deadline });
221 serve_lines(&mut reader, stream, ctx.config.max_line, Auth::Key(&ctx.key), |method, params| {
222 let (req, rrx) = Request::new(method, params);
223 ctx.tx.send(req).ok()?;
224 (ctx.wake)();
225 Some(match rrx.recv_timeout(ctx.config.reply_timeout) {
226 Ok(r) => r,
227 Err(RecvTimeoutError::Timeout) => Reply::err(format!("timeout: the app did not answer within {} s", ctx.config.reply_timeout.as_secs())),
228 Err(RecvTimeoutError::Disconnected) => Reply::err("the app dropped the request"),
229 })
230 });
231}