Implement health check and automated rollback decision logic

1 files changed, 238 insertions, 1 deletions

diff --git a/src/main.rs b/src/main.rs
index fcea38d..a1fe3d0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -18,6 +18,243 @@
 //   If DROPOUT for 1h? => REBOOT
 //   Wait 5m => MONITOR
 
+use std::fs;
+use std::io::{self, Read, Write};
+use std::net::{Shutdown, SocketAddr, TcpListener, TcpStream};
+use std::time::{Duration, Instant};
+
+const UPDATE_FILE: &str = "/data/update";
+const LISTEN_SOCKET: &str = "[::]:12808";
+const MAGIC: [u8; 4] = [0x32, 0x7f, 0xfe, 0x4c];
+const RESP_OK: [u8; 5] = [0x32, 0x7f, 0xfe, 0x4c, 0x00];
+const RESP_NORMAL: [u8; 5] = [0x32, 0x7f, 0xfe, 0x4c, 0x01];
+const CHECK_INTERVAL: Duration = Duration::from_secs(30);
+const CHECK_TIMEOUT: Duration = Duration::from_secs(300);
+const MONITOR_INTERVAL: Duration = Duration::from_secs(300);
+const MONITOR_TIMEOUT: Duration = Duration::from_secs(3600);
+const TCP_TIMEOUT: Duration = Duration::from_secs(8);
+const POLL_INTERVAL: Duration = Duration::from_millis(500);
+const PING_V4: &str = "ipv4.google.com:80";
+const PING_V6: &str = "ipv6.google.com:80";
+
 fn main() {
-    println!("Hello, world!");
+    println!("[info] init");
+
+    match run() {
+        Ok(_) => eprintln!("[warn] logic terminated unexpectedly"),
+        Err(e) => eprintln!("[warn] {}", e),
+    }
+}
+
+fn run() -> io::Result<()> {
+    let ln = TcpListener::bind(LISTEN_SOCKET)?;
+    ln.set_nonblocking(true)?;
+
+    if fs::exists(UPDATE_FILE)? {
+        eprintln!("[info] update detected");
+        check_rollback(&ln)?;
+        eprintln!("[info] no rollback needed");
+    } else {
+        println!("[info] no update, skipping rollback check");
+    }
+
+    monitor(&ln)
+}
+
+fn check_rollback(ln: &TcpListener) -> io::Result<()> {
+    let mut outbound_healthy_v4 = false;
+    let mut outbound_healthy_v6 = false;
+    let mut inbound_healthy = false;
+
+    let t_start = Instant::now();
+    let mut t = t_start;
+    loop {
+        match ln.accept() {
+            Ok((conn, raddr)) => handle_conn_check(conn, raddr, &mut inbound_healthy)?,
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
+            Err(e) => return Err(e),
+        }
+
+        let now = Instant::now();
+        if now.duration_since(t) >= CHECK_INTERVAL {
+            check_connectivity(&mut outbound_healthy_v4, &mut outbound_healthy_v6)?;
+            t = now;
+        }
+
+        if outbound_healthy_v4 && outbound_healthy_v6 && inbound_healthy {
+            break;
+        }
+
+        if now.duration_since(t_start) >= CHECK_TIMEOUT {
+            return rollback();
+        }
+
+        std::thread::sleep(POLL_INTERVAL);
+    }
+
+    Ok(())
+}
+
+fn monitor(ln: &TcpListener) -> io::Result<()> {
+    let mut t_healthy = Instant::now();
+    let mut t = t_healthy;
+    loop {
+        match ln.accept() {
+            Ok((conn, raddr)) => handle_conn_monitor(conn, raddr)?,
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
+            Err(e) => return Err(e),
+        }
+
+        let mut outbound_healthy_v4 = false;
+        let mut outbound_healthy_v6 = false;
+
+        let now = Instant::now();
+        if now.duration_since(t) >= MONITOR_INTERVAL {
+            check_connectivity(&mut outbound_healthy_v4, &mut outbound_healthy_v6)?;
+            t = now;
+        }
+
+        if outbound_healthy_v4 && outbound_healthy_v6 {
+            t_healthy = now;
+            continue;
+        }
+
+        if now.duration_since(t_healthy) >= MONITOR_TIMEOUT {
+            return reboot();
+        }
+
+        std::thread::sleep(POLL_INTERVAL);
+    }
+}
+
+fn handle_conn_check(
+    mut conn: TcpStream,
+    raddr: SocketAddr,
+    inbound_healthy: &mut bool,
+) -> io::Result<()> {
+    let mut buf = [0; 4];
+    loop {
+        match conn.read_exact(&mut buf) {
+            Ok(_) => break,
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
+            Err(e) => return Err(e),
+        }
+    }
+
+    if buf != MAGIC {
+        eprintln!("[warn] handle {raddr}: bad magic {buf:?}");
+        return Ok(());
+    }
+
+    *inbound_healthy = true;
+    conn.write_all(&RESP_OK)?;
+    conn.shutdown(Shutdown::Both)?;
+
+    Ok(())
+}
+
+fn handle_conn_monitor(mut conn: TcpStream, raddr: SocketAddr) -> io::Result<()> {
+    let mut buf = [0; 4];
+    loop {
+        match conn.read_exact(&mut buf) {
+            Ok(_) => break,
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
+            Err(e) => return Err(e),
+        }
+    }
+
+    if buf != MAGIC {
+        eprintln!("[warn] handle {raddr}: bad magic {buf:?}");
+        return Ok(());
+    }
+
+    conn.write_all(&RESP_NORMAL)?;
+    conn.shutdown(Shutdown::Both)?;
+
+    Ok(())
+}
+
+fn check_connectivity(
+    outbound_healthy_v4: &mut bool,
+    outbound_healthy_v6: &mut bool,
+) -> io::Result<()> {
+    let mut buf = [0; 1024];
+
+    let conn4 =
+        match TcpStream::connect_timeout(&PING_V4.parse().expect("PING_V4 invalid"), TCP_TIMEOUT) {
+            Ok(conn) => Some(conn),
+            Err(e) => {
+                eprintln!("[warn] IPv4: connect: {}", e);
+                *outbound_healthy_v4 = false;
+                None
+            }
+        };
+
+    if let Some(mut conn4) = conn4 {
+        conn4.write_all(b"GET / HTTP/1.1")?;
+
+        loop {
+            match conn4.read(&mut buf) {
+                Ok(0) => {
+                    eprintln!("[warn] IPv4: connection closed");
+                    *outbound_healthy_v4 = false;
+                }
+                Ok(_) => {
+                    *outbound_healthy_v4 = true;
+                    break;
+                }
+                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
+                Err(e) => {
+                    eprintln!("[warn] IPv4: read: {}", e);
+                    *outbound_healthy_v4 = false;
+                }
+            }
+        }
+
+        conn4.shutdown(Shutdown::Both)?;
+    }
+
+    let conn6 =
+        match TcpStream::connect_timeout(&PING_V6.parse().expect("PING_V6 invalid"), TCP_TIMEOUT) {
+            Ok(conn) => Some(conn),
+            Err(e) => {
+                eprintln!("[warn] IPv6: connect: {}", e);
+                *outbound_healthy_v6 = false;
+                None
+            }
+        };
+
+    if let Some(mut conn6) = conn6 {
+        conn6.write_all(b"GET / HTTP/1.1")?;
+
+        loop {
+            match conn6.read(&mut buf) {
+                Ok(0) => {
+                    eprintln!("[warn] IPv6: connection closed");
+                    *outbound_healthy_v6 = false;
+                }
+                Ok(_) => {
+                    *outbound_healthy_v6 = true;
+                    break;
+                }
+                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
+                Err(e) => {
+                    eprintln!("[warn] IPv6: read: {}", e);
+                    *outbound_healthy_v6 = false;
+                }
+            }
+        }
+
+        conn6.shutdown(Shutdown::Both)?;
+    }
+
+    Ok(())
+}
+
+fn rollback() -> io::Result<()> {
+    todo!()
+}
+
+fn reboot() -> io::Result<()> {
+    todo!()
 }