Refactored to be more KISS, now only 79 lines of rust code, version 2.0.0

Cargo.lock

@@ -2,4 +2,4 @@
 # It is not intended for manual editing.
 [[package]]
 name = "kiss-ntpd"
-version = "1.0.0"
+version = "2.0.0"

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "kiss-ntpd"
-version = "1.0.0"
+version = "2.0.0"
 authors = ["moparisthebest <admin@moparisthebest.com>"]
 edition = "2018"
 

README.md

@@ -16,16 +16,11 @@ ntp client keep that clock in sync.
 ##### Usage
 
 ```
-$ kiss-ntpd -h
+usage: kiss-ntpd [options...] [bind_addresses...]
-usage: kiss-ntpd [options...]
- -b, --bind                      address to bind to, default '0.0.0.0:123'
  -h, --help                      print this usage text
  -V, -v, --version               Show version number then quit
- -d, --debug                     Print packets sent and recieved, very verbose
 
- Environment variable support:
+ If no bind_addresses supplied, defaults to 0.0.0.0:123
- You if environmental variable KISS_NTPD_BIND is set, it is used in place of --bind
- Also KISS_NTPD_DEBUG=true can be used in place of --debug
 ```
 
 There is an example systemd unit in `systemd/kiss-ntpd.service` which runs it with minimal permissions

src/main.rs

@@ -17,330 +17,107 @@ You should have received a copy of the GNU Affero General Public License
 along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
 
-use std::convert::TryInto;
+use std::io::{Error, ErrorKind, Result};
-use std::env;
+use std::net::UdpSocket;
-use std::io;
-use std::io::{Error, ErrorKind};
-use std::net::{SocketAddr, UdpSocket};
-use std::str::FromStr;
 use std::time::SystemTime;
 
-#[derive(Debug, Copy, Clone)]
+fn ts_now() -> [u8; 8] {
-struct NtpTimestamp {
+    let now = SystemTime::now();
-    ts: u64,
+    let dur = now.duration_since(std::time::UNIX_EPOCH).unwrap(); // this should be unable to happen
+    let secs = dur.as_secs() + 2208988800; // 1900 epoch
+    let nanos = dur.subsec_nanos();
+
+    ((secs << 32) + (nanos as f64 * 4.294967296) as u64).to_be_bytes()
 }
 
-impl NtpTimestamp {
+struct NtpServer {
-    fn now() -> NtpTimestamp {
+    socket: UdpSocket,
-        let now = SystemTime::now();
+    buf: [u8; 48],
-        let dur = now.duration_since(std::time::UNIX_EPOCH).unwrap(); // this should be unable to happen
+}
-        let secs = dur.as_secs() + 2208988800; // 1900 epoch
-        let nanos = dur.subsec_nanos();
 
-        NtpTimestamp {
+impl NtpServer {
-            ts: (secs << 32) + (nanos as f64 * 4.294967296) as u64,
+    fn new(local_addr: String) -> NtpServer {
+        NtpServer {
+            socket: UdpSocket::bind(local_addr).expect("could not bind to socket"),
+            buf: [0u8; 48],
         }
     }
 
-    fn zero() -> NtpTimestamp {
+    fn respond(&mut self) -> Result<usize> {
-        NtpTimestamp { ts: 0 }
+        let (len, remote_addr) = self.socket.recv_from(&mut self.buf)?;
-    }
-
-    fn read(buf: &[u8]) -> NtpTimestamp {
-        // this unwrap can never fail because we always send in exactly 8 bytes
-        NtpTimestamp {
-            ts: u64::from_be_bytes(buf.try_into().unwrap()),
-        }
-    }
-
-    fn write(&self, buf: &mut [u8]) {
-        buf.copy_from_slice(&self.ts.to_be_bytes());
-    }
-}
-
-impl PartialEq for NtpTimestamp {
-    fn eq(&self, other: &NtpTimestamp) -> bool {
-        self.ts == other.ts
-    }
-}
-
-#[derive(Debug, Copy, Clone)]
-struct NtpFracValue {
-    val: u32,
-}
-
-impl NtpFracValue {
-    fn read(buf: &[u8]) -> NtpFracValue {
-        // this unwrap can never fail because we always send in exactly 4 bytes
-        NtpFracValue {
-            val: u32::from_be_bytes(buf.try_into().unwrap()),
-        }
-    }
-
-    fn write(&self, buf: &mut [u8]) {
-        buf.copy_from_slice(&self.val.to_be_bytes());
-    }
-
-    fn zero() -> NtpFracValue {
-        NtpFracValue { val: 0 }
-    }
-}
-
-#[derive(Debug)]
-struct NtpPacket {
-    remote_addr: SocketAddr,
-    local_ts: NtpTimestamp,
-
-    leap: u8,
-    version: u8,
-    mode: u8,
-    stratum: u8,
-    poll: i8,
-    precision: i8,
-    delay: NtpFracValue,
-    dispersion: NtpFracValue,
-    ref_id: u32,
-    ref_ts: NtpTimestamp,
-    orig_ts: NtpTimestamp,
-    rx_ts: NtpTimestamp,
-    tx_ts: NtpTimestamp,
-}
-
-impl NtpPacket {
-    fn receive(socket: &UdpSocket) -> io::Result<NtpPacket> {
-        let mut buf = [0; 1024];
-
-        let (len, addr) = socket.recv_from(&mut buf)?;
-
-        let local_ts = NtpTimestamp::now();
 
         if len < 48 {
             return Err(Error::new(ErrorKind::UnexpectedEof, "Packet too short"));
         }
 
-        let leap = buf[0] >> 6;
+        let version = (self.buf[0] >> 3) & 0x7;
-        let version = (buf[0] >> 3) & 0x7;
+        let mode = self.buf[0] & 0x7;
-        let mode = buf[0] & 0x7;
 
         if version < 1 || version > 4 {
             return Err(Error::new(ErrorKind::Other, "Unsupported version"));
         }
 
-        Ok(NtpPacket {
+        if !(mode == 1 || mode == 3 || (mode == 0 && version == 1 && remote_addr.port() != 123)) {
-            remote_addr: addr,
+            return Err(Error::new(ErrorKind::Other, "Not a valid NTP request"));
-            local_ts: local_ts,
-            leap: leap,
-            version: version,
-            mode: mode,
-            stratum: buf[1],
-            poll: buf[2] as i8,
-            precision: buf[3] as i8,
-            delay: NtpFracValue::read(&buf[4..8]),
-            dispersion: NtpFracValue::read(&buf[8..12]),
-            // this unwrap can never fail because we always send in exactly 4 bytes
-            ref_id: u32::from_be_bytes((&buf[12..16]).try_into().unwrap()),
-            ref_ts: NtpTimestamp::read(&buf[16..24]),
-            orig_ts: NtpTimestamp::read(&buf[24..32]),
-            rx_ts: NtpTimestamp::read(&buf[32..40]),
-            tx_ts: NtpTimestamp::read(&buf[40..48]),
-        })
-    }
-
-    fn send(&self, socket: &UdpSocket) -> io::Result<usize> {
-        let mut buf = [0; 48];
-
-        buf[0] = self.leap << 6 | self.version << 3 | self.mode;
-        buf[1] = self.stratum;
-        buf[2] = self.poll as u8;
-        buf[3] = self.precision as u8;
-        self.delay.write(&mut buf[4..8]);
-        self.dispersion.write(&mut buf[8..12]);
-        &mut buf[12..16].copy_from_slice(&self.ref_id.to_be_bytes());
-        self.ref_ts.write(&mut buf[16..24]);
-        self.orig_ts.write(&mut buf[24..32]);
-        self.rx_ts.write(&mut buf[32..40]);
-        self.tx_ts.write(&mut buf[40..48]);
-
-        socket.send_to(&buf, self.remote_addr)
-    }
-
-    fn is_request(&self) -> bool {
-        self.mode == 1 || self.mode == 3 || (self.mode == 0 && self.version == 1 && self.remote_addr.port() != 123)
-    }
-
-    fn make_response(&self) -> Option<NtpPacket> {
-        if !self.is_request() {
-            return None;
         }
 
-        Some(NtpPacket {
+        let ts = ts_now();
-            remote_addr: self.remote_addr,
-            local_ts: NtpTimestamp::zero(),
-            leap: 0,
-            version: self.version,
-            mode: if self.mode == 1 { 2 } else { 4 },
-            stratum: 8,
-            poll: self.poll,
-            precision: 0,
-            delay: NtpFracValue::zero(),
-            dispersion: NtpFracValue::zero(),
-            ref_id: 0,
-            ref_ts: NtpTimestamp::now(),
-            orig_ts: self.tx_ts,
-            rx_ts: self.local_ts,
-            tx_ts: NtpTimestamp::now(),
-        })
-    }
-}
 
-struct NtpServer {
+        // first 0u8 here is leap
-    socket: UdpSocket,
+        self.buf[0] = 0u8 << 6 | version << 3 | if mode == 1 { 2 } else { 4 };
-    debug: bool,
+        self.buf[1] = 8; // stratum
-}
+        // poll is at index 2 in both request and response, so do nothing
+        self.buf[3] = 0; // precision
+        // delay (4 bytes), dispersion (4 bytes), ref_id (4 bytes), but we don't really care about these
+        // we could set them to 0 but should we bother?
+        //&mut self.buf[4..16].copy_from_slice(&[0u8; 12]);
+        &mut self.buf[16..24].copy_from_slice(&ts); // ref_ts
+        // orig_ts needs moved from 40..48 in request to 24..32 in response
+        let (dst, src) = self.buf.split_at_mut(40);
+        dst[24..32].copy_from_slice(&src[..8]); // orig_ts
+        &mut self.buf[32..40].copy_from_slice(&ts); // rx_ts
+        &mut self.buf[40..48].copy_from_slice(&ts); // tx_ts
 
-impl NtpServer {
+        self.socket.send_to(&self.buf, remote_addr)
-    fn new(local_addr: String, debug: bool) -> NtpServer {
-        NtpServer {
-            socket: UdpSocket::bind(local_addr).expect("could not bind to socket"),
-            debug: debug,
-        }
     }
 
-    fn process_requests(debug: bool, socket: UdpSocket) {
+    fn run(mut self) {
-        println!("Server thread started");
-
         loop {
-            match NtpPacket::receive(&socket) {
+            if let Err(e) = self.respond() {
-                Ok(request) => {
+                eprintln!("error: {}", e);
-                    if debug {
-                        println!("received {:?}", request);
-                    }
-
-                    match request.make_response() {
-                        Some(response) => match response.send(&socket) {
-                            Ok(_) => {
-                                if debug {
-                                    println!("sent {:?}", response);
-                                }
-                            }
-                            Err(e) => println!("failed to send packet to {}: {}", response.remote_addr, e),
-                        },
-                        None => {}
-                    }
-                }
-                Err(e) => {
-                    println!("failed to receive packet: {}", e);
-                }
             }
         }
     }
-
-    fn run(self) {
-        NtpServer::process_requests(self.debug, self.socket);
-    }
-}
-
-fn arg_to_env(arg: &str) -> Option<String> {
-    if !arg.starts_with("--") {
-        return None;
-    }
-    let env = "KISS_NTPD_".to_owned();
-    let mut env = env + &arg.trim_matches('-').replace("-", "_");
-    env.make_ascii_uppercase();
-    Some(env)
-}
-
-fn env_for_arg(arg: &str) -> Option<String> {
-    arg_to_env(arg).and_then(|key| std::env::var(key).ok())
-}
-
-pub struct Args<'a> {
-    args: &'a Vec<String>,
-}
-
-impl<'a> Args<'a> {
-    pub fn new(args: &'a Vec<String>) -> Args {
-        Args { args }
-    }
-    pub fn flag(&self, flag: &'a str) -> bool {
-        if self.args.contains(&flag.to_owned()) {
-            return true;
-        }
-        // because env we want slightly special handling of empty/0/false
-        match env_for_arg(flag) {
-            Some(env) => &env != "" && &env != "0" && &env != "false",
-            None => false,
-        }
-    }
-    pub fn get_option(&self, flags: &[&'a str]) -> Option<String> {
-        for flag in flags.iter() {
-            let mut found = false;
-            for arg in self.args.iter() {
-                if found {
-                    return Some(arg.to_owned());
-                }
-                if arg == flag {
-                    found = true;
-                }
-            }
-        }
-        // no matching arguments are found, so check env variables as a fallback
-        for flag in flags.iter() {
-            let env = env_for_arg(flag);
-            if env.is_some() {
-                return env;
-            }
-        }
-        return None;
-    }
-    pub fn get_str(&self, flags: &[&'a str], def: &'a str) -> String {
-        match self.get_option(flags) {
-            Some(ret) => ret,
-            None => def.to_owned(),
-        }
-    }
-    pub fn get<T: FromStr>(&self, flags: &[&'a str], def: T) -> T {
-        match self.get_option(flags) {
-            Some(ret) => match ret.parse::<T>() {
-                Ok(ret) => ret,
-                Err(_) => def, // or panic
-            },
-            None => def,
-        }
-    }
 }
 
 fn main() {
-    let raw_args = env::args().collect();
+    let mut args = std::env::args().skip(1);
-    let args = Args::new(&raw_args);
-
-    if args.flag("-V") || args.flag("-v") || args.flag("--version") {
-        println!("kiss-ntpd {} ", env!("CARGO_PKG_VERSION"));
-        return;
-    }
 
     let default_udp_host = "0.0.0.0:123";
 
-    let bind = args.get_str(&["-b", "--bind"], default_udp_host).to_owned();
+    let first_arg = args.next().map_or(default_udp_host.to_owned(), |a| a.to_owned());
 
-    if args.flag("-h") || args.flag("--help") {
+    if first_arg == "-h" || first_arg == "--help" {
         println!(
-            r#"usage: kiss-ntpd [options...]
+            r#"usage: kiss-ntpd [options...] [bind_addresses...]
- -b, --bind                      address to bind to, default '{}'
  -h, --help                      print this usage text
  -V, -v, --version               Show version number then quit
- -d, --debug                     Print packets sent and recieved, very verbose
 
- Environment variable support:
+ If no bind_addresses supplied, defaults to {}
- You if environmental variable KISS_NTPD_BIND is set, it is used in place of --bind
- Also KISS_NTPD_DEBUG=true can be used in place of --debug
         "#,
             default_udp_host
         );
         return;
+    } else if first_arg == "-V" || first_arg == "-v" || first_arg == "--version" {
+        println!("kiss-ntpd {} ", env!("CARGO_PKG_VERSION"));
+        return;
     }
 
-    let server = NtpServer::new(bind, args.flag("-d") || args.flag("--debug"));
+    for bind_address in args {
+        std::thread::spawn(|| {
+            NtpServer::new(bind_address).run();
+        });
+    }
 
-    server.run();
+    NtpServer::new(first_arg).run();
 }