// Copyright 2020-2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

#![deny(missing_docs)]
#![allow(clippy::upper_case_acronyms)]
//! NitroSecurityModule IO
//! # Overview
//! This module contains the structure definitions that allows data interchange between
//! a NitroSecureModule and the client using it. It uses CBOR to encode the data to allow
//! easy IPC between components.

// BTreeMap preserves ordering, which makes the tests easier to write
use std::collections::{BTreeMap, BTreeSet};
use std::io::Error as IoError;
use std::result;

use serde::{Deserialize, Serialize};
use serde_bytes::ByteBuf;
use serde_cbor::error::Error as CborError;
use serde_cbor::{from_slice, to_vec};

#[derive(Debug)]
/// Possible error types return from this library.
pub enum Error {
    /// An IO error of type `std::io::Error`
    Io(IoError),
    /// A CBOR ser/de error of type `serde_cbor::error::Error`.
    Cbor(CborError),
}

/// Result type return nsm-io::Error on failure.
pub type Result<T> = result::Result<T, Error>;

impl From<IoError> for Error {
    fn from(error: IoError) -> Self {
        Error::Io(error)
    }
}

impl From<CborError> for Error {
    fn from(error: CborError) -> Self {
        Error::Cbor(error)
    }
}

/// List of error codes that the NSM module can return as part of a Response
#[repr(C)]
#[derive(Debug, Serialize, Deserialize)]
pub enum ErrorCode {
    /// No errors
    Success,

    /// Input argument(s) invalid
    InvalidArgument,

    /// PlatformConfigurationRegister index out of bounds
    InvalidIndex,

    /// The received response does not correspond to the earlier request
    InvalidResponse,

    /// PlatformConfigurationRegister is in read-only mode and the operation
    /// attempted to modify it
    ReadOnlyIndex,

    /// Given request cannot be fulfilled due to missing capabilities
    InvalidOperation,

    /// Operation succeeded but provided output buffer is too small
    BufferTooSmall,

    /// The user-provided input is too large
    InputTooLarge,

    /// NitroSecureModule cannot fulfill request due to internal errors
    InternalError,
}

/// Operations that a NitroSecureModule should implement. Assumes 64K registers will be enough for everyone.
#[derive(Debug, Serialize, Deserialize)]
#[non_exhaustive]
pub enum Request {
    /// Read data from PlatformConfigurationRegister at `index`
    DescribePCR {
        /// index of the PCR to describe
        index: u16,
    },

    /// Extend PlatformConfigurationRegister at `index` with `data`
    ExtendPCR {
        /// index the PCR to extend
        index: u16,

        #[serde(with = "serde_bytes")]
        /// data to extend it with
        data: Vec<u8>,
    },

    /// Lock PlatformConfigurationRegister at `index` from further modifications
    LockPCR {
        /// index to lock
        index: u16,
    },

    /// Lock PlatformConfigurationRegisters at indexes `[0, range)` from further modifications
    LockPCRs {
        /// number of PCRs to lock, starting from index 0
        range: u16,
    },

    /// Return capabilities and version of the connected NitroSecureModule. Clients are recommended to decode
    /// major_version and minor_version first, and use an appropriate structure to hold this data, or fail
    /// if the version is not supported.
    DescribeNSM,

    /// Requests the NSM to create an AttestationDoc and sign it with it's private key to ensure
    /// authenticity.
    Attestation {
        /// Includes additional user data in the AttestationDoc.
        user_data: Option<ByteBuf>,

        /// Includes an additional nonce in the AttestationDoc.
        nonce: Option<ByteBuf>,

        /// Includes a user provided public key in the AttestationDoc.
        public_key: Option<ByteBuf>,
    },

    /// Requests entropy from the NSM side.
    GetRandom,
}

/// Responses received from a NitroSecureModule as a result of a Request
#[derive(Debug, Serialize, Deserialize)]
#[non_exhaustive]
pub enum Response {
    /// returns the current PlatformConfigurationRegister state
    DescribePCR {
        /// true if the PCR is read-only, false otherwise
        lock: bool,
        #[serde(with = "serde_bytes")]
        /// the current value of the PCR
        data: Vec<u8>,
    },

    /// returned if PlatformConfigurationRegister has been successfully extended
    ExtendPCR {
        #[serde(with = "serde_bytes")]
        /// The new value of the PCR after extending the data into the register.
        data: Vec<u8>,
    },

    /// returned if PlatformConfigurationRegister has been successfully locked
    LockPCR,

    /// returned if PlatformConfigurationRegisters have been successfully locked
    LockPCRs,

    /// returns the runtime configuration of the NitroSecureModule
    DescribeNSM {
        /// Breaking API changes are denoted by `major_version`
        version_major: u16,
        /// Minor API changes are denoted by `minor_version`. Minor versions should be backwards compatible.
        version_minor: u16,
        /// Patch version. These are security and stability updates and do not affect API.
        version_patch: u16,
        /// `module_id` is an identifier for a singular NitroSecureModule
        module_id: String,
        /// The maximum number of PCRs exposed by the NitroSecureModule.
        max_pcrs: u16,
        /// The PCRs that are read-only.
        locked_pcrs: BTreeSet<u16>,
        /// The digest of the PCR Bank
        digest: Digest,
    },

    /// A response to an Attestation Request containing the CBOR-encoded AttestationDoc and the
    /// signature generated from the doc by the NitroSecureModule
    Attestation {
        /// A signed COSE structure containing a CBOR-encoded AttestationDocument as the payload.
        #[serde(with = "serde_bytes")]
        document: Vec<u8>,
    },

    /// A response containing a number of bytes of entropy.
    GetRandom {
        #[serde(with = "serde_bytes")]
        /// The random bytes.
        random: Vec<u8>,
    },

    /// An error has occured, and the NitroSecureModule could not successfully complete the operation
    Error(ErrorCode),
}

/// The digest implementation used by a NitroSecureModule
#[repr(C)]
#[derive(Debug, Serialize, Deserialize, Copy, Clone, PartialEq)]
pub enum Digest {
    /// SHA256
    SHA256,
    /// SHA384
    SHA384,
    /// SHA512
    SHA512,
}

/// An attestation response.  This is also used for sealing
/// data.
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct AttestationDoc {
    /// Issuing NSM ID
    pub module_id: String,

    /// The digest function used for calculating the register values
    /// Can be: "SHA256" | "SHA512"
    pub digest: Digest,

    /// UTC time when document was created expressed as milliseconds since Unix Epoch
    pub timestamp: u64,

    /// Map of all locked PCRs at the moment the attestation document was generated
    pub pcrs: BTreeMap<usize, ByteBuf>,

    /// The infrastucture certificate used to sign the document, DER encoded
    pub certificate: ByteBuf,
    /// Issuing CA bundle for infrastructure certificate
    pub cabundle: Vec<ByteBuf>,

    /// An optional DER-encoded key the attestation consumer can use to encrypt data with
    pub public_key: Option<ByteBuf>,

    /// Additional signed user data, as defined by protocol.
    pub user_data: Option<ByteBuf>,

    /// An optional cryptographic nonce provided by the attestation consumer as a proof of
    /// authenticity.
    pub nonce: Option<ByteBuf>,
}

impl AttestationDoc {
    /// Creates a new AttestationDoc.
    ///
    /// # Arguments
    ///
    /// * module_id: a String representing the name of the NitroSecureModule
    /// * digest: nsm_io::Digest that describes what the PlatformConfigurationRegisters
    ///           contain
    /// * pcrs: BTreeMap containing the index to PCR value
    /// * certificate: the serialized certificate that will be used to sign this AttestationDoc
    /// * cabundle: the serialized set of certificates up to the root of trust certificate that
    ///             emitted `certificate`
    /// * user_data: optional user definted data included in the AttestationDoc
    /// * nonce: optional cryptographic nonce that will be included in the AttestationDoc
    /// * public_key: optional DER-encoded public key that will be included in the AttestationDoc
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        module_id: String,
        digest: Digest,
        timestamp: u64,
        pcrs: BTreeMap<usize, Vec<u8>>,
        certificate: Vec<u8>,
        cabundle: Vec<Vec<u8>>,
        user_data: Option<Vec<u8>>,
        nonce: Option<Vec<u8>>,
        public_key: Option<Vec<u8>>,
    ) -> Self {
        let mut pcrs_serialized = BTreeMap::new();

        for (i, pcr) in pcrs.into_iter() {
            let pcr = ByteBuf::from(pcr);
            pcrs_serialized.insert(i, pcr);
        }

        let cabundle_serialized = cabundle.into_iter().map(ByteBuf::from).collect();

        AttestationDoc {
            module_id,
            digest,
            timestamp,
            pcrs: pcrs_serialized,
            cabundle: cabundle_serialized,
            certificate: ByteBuf::from(certificate),
            user_data: user_data.map(ByteBuf::from),
            nonce: nonce.map(ByteBuf::from),
            public_key: public_key.map(ByteBuf::from),
        }
    }

    /// Helper function that converts an AttestationDoc structure to its CBOR representation
    pub fn to_binary(&self) -> Vec<u8> {
        // This should not fail
        to_vec(self).unwrap()
    }

    /// Helper function that parses a CBOR representation of an AttestationDoc and creates the
    /// structure from it, if possible.
    pub fn from_binary(bin: &[u8]) -> Result<Self> {
        from_slice(bin).map_err(Error::Cbor)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_attestationdoc_binary_encode() {
        let mut pcrs = BTreeMap::new();
        pcrs.insert(1, vec![1, 2, 3]);
        pcrs.insert(2, vec![4, 5, 6]);
        pcrs.insert(3, vec![7, 8, 9]);

        let doc1 = AttestationDoc::new(
            "abcd".to_string(),
            Digest::SHA256,
            1234,
            pcrs,
            vec![42; 10],
            vec![],
            Some(vec![255; 10]),
            None,
            None,
        );
        let bin1 = doc1.to_binary();
        let doc2 = AttestationDoc::from_binary(&bin1).unwrap();
        let bin2 = doc2.to_binary();
        assert_eq!(doc1, doc2);
        assert_eq!(bin1, bin2);
    }
}


// Copyright 2020-2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! ***NitroSecureModule driver communication support***
//! # Overview
//! This module implements support functions for communicating with the NSM
//! driver by encoding requests to / decoding responses from a C-compatible
//! message structure which is shared with the driver via `ioctl()`.
//! In general, a message contains:
//! 1. A *request* content structure, holding CBOR-encoded user input data.
//! 2. A *response* content structure, with an initial memory capacity provided by
//! the user, which then gets populated with information from the NSM driver and
//! then decoded from CBOR.

use crate::api::{ErrorCode, Request, Response};
use libc::ioctl;
use log::{debug, error};
use nix::errno::Errno;
use nix::request_code_readwrite;
use nix::unistd::close;

use std::fs::OpenOptions;
use std::io::{IoSlice, IoSliceMut};
use std::mem;
use std::os::unix::io::{IntoRawFd, RawFd};

const DEV_FILE: &str = "/dev/nsm";
const NSM_IOCTL_MAGIC: u8 = 0x0A;
const NSM_REQUEST_MAX_SIZE: usize = 0x1000;
const NSM_RESPONSE_MAX_SIZE: usize = 0x3000;

/// NSM message structure to be used with `ioctl()`.
#[repr(C)]
struct NsmMessage<'a> {
    /// User-provided data for the request
    pub request: IoSlice<'a>,
    /// Response data provided by the NSM pipeline
    pub response: IoSliceMut<'a>,
}

/// Encode an NSM `Request` value into a vector.
/// *Argument 1 (input)*: The NSM request.
/// *Returns*: The vector containing the CBOR encoding.
fn nsm_encode_request_to_cbor(request: Request) -> Vec<u8> {
    serde_cbor::to_vec(&request).unwrap()
}

/// Decode an NSM `Response` value from a raw memory buffer.
/// *Argument 1 (input)*: The `iovec` holding the memory buffer.
/// *Returns*: The decoded NSM response.
fn nsm_decode_response_from_cbor(response_data: &IoSliceMut<'_>) -> Response {
    match serde_cbor::from_slice(response_data) {
        Ok(response) => response,
        Err(_) => Response::Error(ErrorCode::InternalError),
    }
}

/// Do an `ioctl()` of a given type for a given message.
/// *Argument 1 (input)*: The descriptor to the device file.
/// *Argument 2 (input/output)*: The message to be sent and updated via `ioctl()`.
/// *Returns*: The status of the operation.
fn nsm_ioctl(fd: i32, message: &mut NsmMessage) -> Option<Errno> {
    let status = unsafe {
        ioctl(
            fd,
            request_code_readwrite!(NSM_IOCTL_MAGIC, 0, mem::size_of::<NsmMessage>()),
            message,
        )
    };
    let errno = Errno::last();

    match status {
        // If ioctl() succeeded, the status is the message's response code
        0 => None,

        // If ioctl() failed, the error is given by errno
        _ => Some(errno),
    }
}

/// Create a message with input data and output capacity from a given
/// request, then send it to the NSM driver via `ioctl()` and wait
/// for the driver's response.
/// *Argument 1 (input)*: The descriptor to the NSM device file.
/// *Argument 2 (input)*: The NSM request.
/// *Returns*: The corresponding NSM response from the driver.
pub fn nsm_process_request(fd: i32, request: Request) -> Response {
    let cbor_request = nsm_encode_request_to_cbor(request);

    // Check if the request is too large
    if cbor_request.len() > NSM_REQUEST_MAX_SIZE {
        return Response::Error(ErrorCode::InputTooLarge);
    }

    let mut cbor_response: [u8; NSM_RESPONSE_MAX_SIZE] = [0; NSM_RESPONSE_MAX_SIZE];
    let mut message = NsmMessage {
        request: IoSlice::new(&cbor_request),
        response: IoSliceMut::new(&mut cbor_response),
    };
    let status = nsm_ioctl(fd, &mut message);

    match status {
        None => nsm_decode_response_from_cbor(&message.response),
        Some(errno) => match errno {
            Errno::EMSGSIZE => Response::Error(ErrorCode::InputTooLarge),
            _ => Response::Error(ErrorCode::InternalError),
        },
    }
}

/// NSM library initialization function.
/// *Returns*: A descriptor for the opened device file.
pub fn nsm_init() -> i32 {
    let mut open_options = OpenOptions::new();
    let open_dev = open_options.read(true).write(true).open(DEV_FILE);

    match open_dev {
        Ok(open_dev) => {
            debug!("Device file '{}' opened successfully.", DEV_FILE);
            open_dev.into_raw_fd()
        }
        Err(e) => {
            error!("Device file '{}' failed to open: {}", DEV_FILE, e);
            -1
        }
    }
}

/// NSM library exit function.
/// *Argument 1 (input)*: The descriptor for the opened device file, as
/// obtained from `nsm_init()`.
pub fn nsm_exit(fd: i32) {
    let result = close(fd as RawFd);
    match result {
        Ok(()) => debug!("File of descriptor {} closed successfully.", fd),
        Err(e) => error!("File of descriptor {} failed to close: {}", fd, e),
    }
}


use nsm_io::Request;
use serde_bytes::ByteBuf;

fn main() {
    let nsm_fd = nsm_driver::nsm_init();

    let public_key = ByteBuf::from("my super secret key");
    let hello = ByteBuf::from("hello, world!");

    let request = Request::Attestation {
        public_key: Some(public_key),
        user_data: Some(hello),
        nonce: None,
    };

    let response = nsm_driver::nsm_process_request(nsm_fd, request);
    println!("{:?}", response);

    nsm_driver::nsm_exit(nsm_fd);
}


use clap::ArgMatches;

#[derive(Debug, Clone)]
pub struct ServerArgs {
    pub port: u32,
}

impl ServerArgs {
    pub fn new_with(args: &ArgMatches) -> Result<Self, String> {
        Ok(ServerArgs {
            port: parse_port(args)?,
        })
    }
}

#[derive(Debug, Clone)]
pub struct ClientArgs {
    pub cid: u32,
    pub port: u32,
}

impl ClientArgs {
    pub fn new_with(args: &ArgMatches) -> Result<Self, String> {
        Ok(ClientArgs {
            cid: parse_cid_client(args)?,
            port: parse_port(args)?,
        })
    }
}

fn parse_cid_client(args: &ArgMatches) -> Result<u32, String> {
    let cid = args.value_of("cid").ok_or("Could not find cid argument")?;
    cid.parse()
        .map_err(|_err| "cid is not a number".to_string())
}

fn parse_port(args: &ArgMatches) -> Result<u32, String> {
    let port = args
        .value_of("port")
        .ok_or("Could not find port argument")?;
    port.parse()
        .map_err(|_err| "port is not a number".to_string())
}


pub mod command_parser;
pub mod protocol_helpers;
pub mod utils;

use command_parser::{ClientArgs, ServerArgs};
use protocol_helpers::{recv_loop, recv_u64, send_loop, send_u64};

use nix::sys::socket::listen as listen_vsock;
use nix::sys::socket::{accept, bind, connect, shutdown, socket};
use nix::sys::socket::{AddressFamily, Shutdown, SockAddr, SockFlag, SockType};
use nix::unistd::close;
use std::convert::TryInto;
use std::os::unix::io::{AsRawFd, RawFd};

const VMADDR_CID_ANY: u32 = 0xFFFFFFFF;
const BUF_MAX_LEN: usize = 8192;
// Maximum number of outstanding connections in the socket's
// listen queue
const BACKLOG: usize = 128;
// Maximum number of connection attempts
const MAX_CONNECTION_ATTEMPTS: usize = 5;

struct VsockSocket {
    socket_fd: RawFd,
}

impl VsockSocket {
    fn new(socket_fd: RawFd) -> Self {
        VsockSocket { socket_fd }
    }
}

impl Drop for VsockSocket {
    fn drop(&mut self) {
        shutdown(self.socket_fd, Shutdown::Both)
            .unwrap_or_else(|e| eprintln!("Failed to shut socket down: {:?}", e));
        close(self.socket_fd).unwrap_or_else(|e| eprintln!("Failed to close socket: {:?}", e));
    }
}

impl AsRawFd for VsockSocket {
    fn as_raw_fd(&self) -> RawFd {
        self.socket_fd
    }
}

/// Initiate a connection on an AF_VSOCK socket
fn vsock_connect(cid: u32, port: u32) -> Result<VsockSocket, String> {
    let sockaddr = SockAddr::new_vsock(cid, port);
    let mut err_msg = String::new();

    for i in 0..MAX_CONNECTION_ATTEMPTS {
        let vsocket = VsockSocket::new(
            socket(
                AddressFamily::Vsock,
                SockType::Stream,
                SockFlag::empty(),
                None,
            )
            .map_err(|err| format!("Failed to create the socket: {:?}", err))?,
        );
        match connect(vsocket.as_raw_fd(), &sockaddr) {
            Ok(_) => return Ok(vsocket),
            Err(e) => err_msg = format!("Failed to connect: {}", e),
        }

        // Exponentially backoff before retrying to connect to the socket
        std::thread::sleep(std::time::Duration::from_secs(1 << i));
    }

    Err(err_msg)
}

/// Send 'Hello, world!' to the server
pub fn client(args: ClientArgs) -> Result<(), String> {
    let vsocket = vsock_connect(args.cid, args.port)?;
    let fd = vsocket.as_raw_fd();

    // TODO: Replace this with your client code
    let data = "Hello, world!".to_string();
    let buf = data.as_bytes();
    let len: u64 = buf.len().try_into().map_err(|err| format!("{:?}", err))?;
    send_u64(fd, len)?;
    send_loop(fd, buf, len)?;

    Ok(())
}

/// Accept connections on a certain port and print
/// the received data
pub fn server(args: ServerArgs) -> Result<(), String> {
    let socket_fd = socket(
        AddressFamily::Vsock,
        SockType::Stream,
        SockFlag::empty(),
        None,
    )
    .map_err(|err| format!("Create socket failed: {:?}", err))?;

    let sockaddr = SockAddr::new_vsock(VMADDR_CID_ANY, args.port);

    bind(socket_fd, &sockaddr).map_err(|err| format!("Bind failed: {:?}", err))?;

    listen_vsock(socket_fd, BACKLOG).map_err(|err| format!("Listen failed: {:?}", err))?;

    loop {
        let fd = accept(socket_fd).map_err(|err| format!("Accept failed: {:?}", err))?;

        // TODO: Replace this with your server code
        let len = recv_u64(fd)?;
        let mut buf = [0u8; BUF_MAX_LEN];
        recv_loop(fd, &mut buf, len)?;
        println!(
            "{}",
            String::from_utf8(buf.to_vec())
                .map_err(|err| format!("The received bytes are not UTF-8: {:?}", err))?
        );
    }
}


use clap::{App, AppSettings, Arg, SubCommand};

use vsock_sample::command_parser::{ClientArgs, ServerArgs};
use vsock_sample::create_app;
use vsock_sample::{client, server};

fn main() {
    let app = create_app!();
    let args = app.get_matches();

    match args.subcommand() {
        Some(("server", args)) => {
            let server_args = ServerArgs::new_with(args).unwrap();
            server(server_args).unwrap();
        }
        Some(("client", args)) => {
            let client_args = ClientArgs::new_with(args).unwrap();
            client(client_args).unwrap();
        }
        Some(_) | None => ()
    }
}


use byteorder::{ByteOrder, LittleEndian};
use nix::errno::Errno::EINTR;
use nix::sys::socket::MsgFlags;
use nix::sys::socket::{recv, send};
use std::convert::TryInto;
use std::mem::size_of;
use std::os::unix::io::RawFd;

pub fn send_u64(fd: RawFd, val: u64) -> Result<(), String> {
    let mut buf = [0u8; size_of::<u64>()];
    LittleEndian::write_u64(&mut buf, val);
    send_loop(fd, &buf, size_of::<u64>().try_into().unwrap())?;
    Ok(())
}

pub fn recv_u64(fd: RawFd) -> Result<u64, String> {
    let mut buf = [0u8; size_of::<u64>()];
    recv_loop(fd, &mut buf, size_of::<u64>().try_into().unwrap())?;
    let val = LittleEndian::read_u64(&buf);
    Ok(val)
}

/// Send `len` bytes from `buf` to a connection-oriented socket
pub fn send_loop(fd: RawFd, buf: &[u8], len: u64) -> Result<(), String> {
    let len: usize = len.try_into().map_err(|err| format!("{:?}", err))?;
    let mut send_bytes = 0;

    while send_bytes < len {
        let size = match send(fd, &buf[send_bytes..len], MsgFlags::empty()) {
            Ok(size) => size,
            Err(nix::Error::Sys(EINTR)) => 0,
            Err(err) => return Err(format!("{:?}", err)),
        };
        send_bytes += size;
    }

    Ok(())
}

/// Receive `len` bytes from a connection-orriented socket
pub fn recv_loop(fd: RawFd, buf: &mut [u8], len: u64) -> Result<(), String> {
    let len: usize = len.try_into().map_err(|err| format!("{:?}", err))?;
    let mut recv_bytes = 0;

    while recv_bytes < len {
        let size = match recv(fd, &mut buf[recv_bytes..len], MsgFlags::empty()) {
            Ok(size) => size,
            Err(nix::Error::Sys(EINTR)) => 0,
            Err(err) => return Err(format!("{:?}", err)),
        };
        recv_bytes += size;
    }

    Ok(())
}


use log::error;

pub trait ExitGracefully<T, E> {
    fn ok_or_exit(self, message: &str) -> T;
}

impl<T, E: std::fmt::Debug> ExitGracefully<T, E> for Result<T, E> {
    fn ok_or_exit(self, message: &str) -> T {
        match self {
            Ok(val) => val,
            Err(err) => {
                error!("{:?}: {}", err, message);
                std::process::exit(1);
            }
        }
    }
}

#[macro_export]
macro_rules! create_app {
    () => {
        App::new("Vsock Sample")
            .about("Hello world example for vsock server and client communication.")
            .setting(AppSettings::ArgRequiredElseHelp)
            .version(env!("CARGO_PKG_VERSION"))
            .subcommand(
                SubCommand::with_name("server")
                    .about("Listen on a given port.")
                    .arg(
                        Arg::with_name("port")
                            .long("port")
                            .help("port")
                            .takes_value(true)
                            .required(true),
                    ),
            )
            .subcommand(
                SubCommand::with_name("client")
                    .about("Connect to a given cid and port.")
                    .arg(
                        Arg::with_name("port")
                            .long("port")
                            .help("port")
                            .takes_value(true)
                            .required(true),
                    )
                    .arg(
                        Arg::with_name("cid")
                            .long("cid")
                            .help("cid")
                            .takes_value(true)
                            .required(true),
                    ),
            )
    };
}


use lambda_runtime::{run, service_fn, Error, LambdaEvent};

use serde::{Deserialize, Serialize};

/// This is a made-up example. Requests come into the runtime as unicode
/// strings in json format, which can map to any structure that implements `serde::Deserialize`
/// The runtime pays no attention to the contents of the request payload.
#[derive(Deserialize)]
struct Request {
}

/// This is a made-up example of what a response structure may look like.
/// There is no restriction on what it can be. The runtime requires responses
/// to be serialized into json. The runtime pays no attention
/// to the contents of the response payload.
#[derive(Serialize)]
struct Response {
    req_id: String,
    msg: String,
}

/// This is the main body for the function.
/// Write your code inside it.
/// There are some code example in the following URLs:
/// - https://github.com/awslabs/aws-lambda-rust-runtime/tree/main/examples
/// - https://github.com/aws-samples/serverless-rust-demo/
async fn function_handler(event: LambdaEvent<Request>) -> Result<Response, Error> {
    // Prepare the response
    let resp = Response {
        req_id: "63".to_string(),
        msg: "Hello FunctionA".to_string(),
    };

    // Return `Response` (it will be serialized to JSON automatically by the runtime)
    Ok(resp)
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        // disable printing the name of the module in every log line.
        .with_target(false)
        // disabling time is handy because CloudWatch will add the ingestion time.
        .without_time()
        .init();

    run(service_fn(function_handler)).await
}


use lambda_runtime::{run, service_fn, Error, LambdaEvent};

use serde::{Deserialize, Serialize};

/// This is a made-up example. Requests come into the runtime as unicode
/// strings in json format, which can map to any structure that implements `serde::Deserialize`
/// The runtime pays no attention to the contents of the request payload.
#[derive(Deserialize)]
struct Request {
}

/// This is a made-up example of what a response structure may look like.
/// There is no restriction on what it can be. The runtime requires responses
/// to be serialized into json. The runtime pays no attention
/// to the contents of the response payload.
#[derive(Serialize)]
struct Response {
    req_id: String,
    msg: String,
}

/// This is the main body for the function.
/// Write your code inside it.
/// There are some code example in the following URLs:
/// - https://github.com/awslabs/aws-lambda-rust-runtime/tree/main/examples
/// - https://github.com/aws-samples/serverless-rust-demo/
async fn function_handler(event: LambdaEvent<Request>) -> Result<Response, Error> {
    // Prepare the response
    let resp = Response {
        req_id: "99".to_string(),
        msg: "Hello FunctionB".to_string(),
    };

    // Return `Response` (it will be serialized to JSON automatically by the runtime)
    Ok(resp)
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        // disable printing the name of the module in every log line.
        .with_target(false)
        // disabling time is handy because CloudWatch will add the ingestion time.
        .without_time()
        .init();

    run(service_fn(function_handler)).await
}


use lambda_runtime::{run, service_fn, Error, LambdaEvent};

use serde::{Deserialize, Serialize};

/// This is a made-up example. Requests come into the runtime as unicode
/// strings in json format, which can map to any structure that implements `serde::Deserialize`
/// The runtime pays no attention to the contents of the request payload.
#[derive(Deserialize)]
struct Request {
}

/// This is a made-up example of what a response structure may look like.
/// There is no restriction on what it can be. The runtime requires responses
/// to be serialized into json. The runtime pays no attention
/// to the contents of the response payload.
#[derive(Serialize)]
struct Response {
    req_id: String,
    msg: String,
}

/// This is the main body for the function.
/// Write your code inside it.
/// There are some code example in the following URLs:
/// - https://github.com/awslabs/aws-lambda-rust-runtime/tree/main/examples
/// - https://github.com/aws-samples/serverless-rust-demo/
async fn function_handler(event: LambdaEvent<Request>) -> Result<Response, Error> {
    // Prepare the response
    let resp = Response {
        req_id: "34".to_string(),
        msg: "Hello World".to_string(),
    };

    // Return `Response` (it will be serialized to JSON automatically by the runtime)
    Ok(resp)
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        // disable printing the name of the module in every log line.
        .with_target(false)
        // disabling time is handy because CloudWatch will add the ingestion time.
        .without_time()
        .init();

    run(service_fn(function_handler)).await
}


use lambda_runtime::{run, service_fn, Error, LambdaEvent};

use serde::{Deserialize, Serialize};

/// This is a made-up example. Requests come into the runtime as unicode
/// strings in json format, which can map to any structure that implements `serde::Deserialize`
/// The runtime pays no attention to the contents of the request payload.
#[derive(Deserialize)]
struct Request {
}

/// This is a made-up example of what a response structure may look like.
/// There is no restriction on what it can be. The runtime requires responses
/// to be serialized into json. The runtime pays no attention
/// to the contents of the response payload.
#[derive(Serialize)]
struct Response {
    statusCode: i32,
    body: String,
}

/// This is the main body for the function.
/// Write your code inside it.
/// There are some code example in the following URLs:
/// - https://github.com/awslabs/aws-lambda-rust-runtime/tree/main/examples
/// - https://github.com/aws-samples/serverless-rust-demo/
async fn function_handler(event: LambdaEvent<Request>) -> Result<Response, Error> {
    // Prepare the response
    let resp = Response {
        statusCode: 200,
        body: "Hello World!".to_string(),
    };

    // Return `Response` (it will be serialized to JSON automatically by the runtime)
    Ok(resp)
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        // disable printing the name of the module in every log line.
        .with_target(false)
        // disabling time is handy because CloudWatch will add the ingestion time.
        .without_time()
        .init();

    run(service_fn(function_handler)).await
}


use aws_sdk_dynamodb::{model::AttributeValue, Client};
use lambda_http::{service_fn, Body, Error, Request, RequestExt, Response};
use std::env;

/// Main function
#[tokio::main]
async fn main() -> Result<(), Error> {
    // Initialize the AWS SDK for Rust
    let config = aws_config::load_from_env().await;
    let table_name = env::var("TABLE_NAME").expect("TABLE_NAME must be set");
    let dynamodb_client = Client::new(&config);

    // Register the Lambda handler
    //
    // We use a closure to pass the `dynamodb_client` and `table_name` as arguments
    // to the handler function.
    lambda_http::run(service_fn(|request: Request| {
        put_item(&dynamodb_client, &table_name, request)
    }))
    .await?;

    Ok(())
}

/// Put Item Lambda function
///
/// This function will run for every invoke of the Lambda function.
async fn put_item(
    client: &Client,
    table_name: &str,
    request: Request,
) -> Result<Response<Body>, Error> {
    // Extract path parameter from request
    let path_parameters = request.path_parameters();
    let id = match path_parameters.first("id") {
        Some(id) => id,
        None => return Ok(Response::builder().status(400).body("id is required".into())?),
    };

    // Extract body from request
    let body = match request.body() {
        Body::Empty => "".to_string(),
        Body::Text(body) => body.clone(),
        Body::Binary(body) => String::from_utf8_lossy(body).to_string(),
    };

    // Put the item in the DynamoDB table
    let res = client
        .put_item()
        .table_name(table_name)
        .item("id", AttributeValue::S(id.to_string()))
        .item("payload", AttributeValue::S(body))
        .send()
        .await;

    // Return a response to the end-user
    match res {
        Ok(_) => Ok(Response::builder().status(200).body("item saved".into())?),
        Err(_) => Ok(Response::builder().status(500).body("internal error".into())?),
    }
}

/// Unit tests
///
/// These tests are run using the `cargo test` command.
#[cfg(test)]
mod tests {
    use super::*;
    use aws_sdk_dynamodb::{Client, Config, Credentials, Region};
    use aws_smithy_client::{erase::DynConnector, test_connection::TestConnection};
    use aws_smithy_http::body::SdkBody;
    use std::collections::HashMap;

    // Helper function to create a mock AWS configuration
    async fn get_mock_config(conn: &TestConnection<SdkBody>) -> Config {
        let cfg = aws_config::from_env()
            .region(Region::new("eu-west-1"))
            .http_connector(DynConnector::new(conn.clone()))
            .credentials_provider(Credentials::new(
                "access_key",
                "privatekey",
                None,
                None,
                "dummy",
            ))
            .load()
            .await;

        Config::new(&cfg)
    }

    /// Helper function to generate a sample DynamoDB request
    fn get_request_builder() -> http::request::Builder {
        http::Request::builder()
            .header("content-type", "application/x-amz-json-1.0")
            .uri(http::uri::Uri::from_static(
                "https://dynamodb.eu-west-1.amazonaws.com/",
            ))
    }

    #[tokio::test]
    async fn test_put_item() {
        // Mock DynamoDB client
        //
        // `TestConnection` takes a vector of requests and responses, allowing us to
        // simulate the behaviour of the DynamoDB API endpoint. Since we are only
        // making a single request in this test, we only need to provide a single
        // entry in the vector.
        let conn = TestConnection::new(vec![(
            get_request_builder()
                .header("x-amz-target", "DynamoDB_20120810.PutItem")
                .body(SdkBody::from(
                    r#"{"TableName":"test","Item":{"id":{"S":"1"},"payload":{"S":"test1"}}}"#,
                ))
                .unwrap(),
            http::Response::builder()
                .status(200)
                .body(SdkBody::from(
                    r#"{"Attributes": {"id": {"S": "1"}, "payload": {"S": "test1"}}}"#,
                ))
                .unwrap(),
        )]);
        let client = Client::from_conf(get_mock_config(&conn).await);

        let table_name = "test_table";

        // Mock API Gateway request
        let mut path_parameters = HashMap::new();
        path_parameters.insert("id".to_string(), vec!["1".to_string()]);

        let request = http::Request::builder()
            .method("PUT")
            .uri("/1")
            .body(Body::Text("test1".to_string()))
            .unwrap()
            .with_path_parameters(path_parameters);

        // Send mock request to Lambda handler function
        let response = put_item(&client, table_name, request)
            .await
            .unwrap();

        // Assert that the response is correct
        assert_eq!(response.status(), 200);
        assert_eq!(response.body(), &Body::Text("item saved".to_string()));
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: GPL-2.0-only

use chrono::Utc;
use chrony_candm::Client;
use chrony_candm::{reply::ReplyBody, request::RequestBody};
use std::{
    net::{Ipv6Addr, SocketAddr, SocketAddrV6},
    str::FromStr,
};
use tokio::runtime::Handle;

#[tokio::main(flavor = "current_thread")]
async fn main() {
    let request_body = RequestBody::Tracking;
    let client = Client::spawn(&Handle::current(), Default::default());
    let server_addr = SocketAddr::V6(SocketAddrV6::new(
        Ipv6Addr::from_str("::1").unwrap(),
        323,
        0,
        0,
    ));
    let result = client.query(request_body, server_addr).await;

    match result {
        Err(e) => println!("{:?}", e),
        Ok(reply) => {
            println!("Status: {:?}", reply.status);
            if let ReplyBody::Tracking(body) = reply.body {
                println!("Reference ID: {}", body.ref_id);
                println!("Source IP: {}", body.ip_addr);
                println!("Stratum: {}", body.stratum);
                println!(
                    "Ref time: {}",
                    chrono::DateTime::<Utc>::from(body.ref_time)
                        .to_rfc3339_opts(chrono::SecondsFormat::Nanos, true)
                );
                let current_correction: f64 = body.current_correction.into();
                println!(
                    "System time: {:.9} seconds {} of NTP time",
                    current_correction.abs(),
                    if current_correction.is_sign_negative() {
                        "fast"
                    } else {
                        "slow"
                    }
                );
                println!("Last offset: {:+.9} seconds", body.last_offset);
                println!("RMS offset: {:.9} seconds", body.rms_offset);
                let freq_ppm: f64 = body.freq_ppm.into();
                println!(
                    "Frequency: {:.3} ppm {}",
                    freq_ppm.abs(),
                    if freq_ppm.is_sign_negative() {
                        "slow"
                    } else {
                        "fast"
                    }
                );
                println!("Residual freq: {:+.3} ppm", body.resid_freq_ppm);
                println!("Skew: {:.3} ppm", body.skew_ppm);
                println!("Root delay: {:.9} seconds", body.root_delay);
                println!("Root dispersion: {:.9} seconds", body.root_dispersion);
                println!("Update interval: {:.1} seconds", body.last_update_interval);
                println!(
                    "Leap status: {}",
                    match body.leap_status {
                        0 => "Normal",
                        1 => "Insert",
                        2 => "Delete",
                        3 => "Unsynchronized",
                        _ => "Invalid",
                    }
                )
            }
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: GPL-2.0-only

use std::collections::{HashMap, VecDeque};
use std::ffi::OsStr;
use std::fs::Permissions;
use std::future::Future;
use std::hash::{Hash, Hasher};
use std::net::{Ipv6Addr, SocketAddr, SocketAddrV6};
use std::ops::{Deref, DerefMut};
use std::pin::Pin;
use std::sync::Arc;

use futures::future::TryFutureExt;
use rand::Rng;
use tokio::sync::mpsc::error::SendError;
use tokio::sync::OnceCell;
use tokio::time::Instant;

use crate::reply::Reply;
use crate::request::{Request, RequestBody};
use crate::ClientOptions;

#[cfg(unix)]
use std::os::unix::ffi::OsStrExt;
#[cfg(unix)]
use std::os::unix::fs::PermissionsExt;
#[cfg(unix)]
use tokio::net::UnixDatagram;

#[cfg(unix)]
#[derive(Debug)]
struct UnixDatagramClient(UnixDatagram);

#[cfg(unix)]
impl AsRef<UnixDatagram> for UnixDatagramClient {
    fn as_ref(&self) -> &UnixDatagram {
        &self.0
    }
}

#[cfg(unix)]
impl AsMut<UnixDatagram> for UnixDatagramClient {
    fn as_mut(&mut self) -> &mut UnixDatagram {
        &mut self.0
    }
}

#[cfg(unix)]
impl Deref for UnixDatagramClient {
    type Target = UnixDatagram;
    fn deref(&self) -> &UnixDatagram {
        &self.0
    }
}

#[cfg(unix)]
impl DerefMut for UnixDatagramClient {
    fn deref_mut(&mut self) -> &mut UnixDatagram {
        &mut self.0
    }
}

#[cfg(unix)]
impl Drop for UnixDatagramClient {
    fn drop(&mut self) {
        if let Ok(addr) = self.0.local_addr() {
            if let Some(path) = addr.as_pathname() {
                let _ = self.0.shutdown(std::net::Shutdown::Both);
                let _ = std::fs::remove_file(path);
            }
        }
    }
}

#[cfg(unix)]
impl UnixDatagramClient {
    async fn new() -> std::io::Result<UnixDatagramClient> {
        let id: [u8; 16] = rand::random();
        let mut path = b"/var/run/chrony/client-000102030405060708090a0b0c0d0e0f.sock".clone();
        hex::encode_to_slice(id, &mut path[23..55]).unwrap();
        let path_str = OsStr::from_bytes(&path);
        let sock = UnixDatagram::bind(path_str)?;
        let client = UnixDatagramClient(sock);
        std::fs::set_permissions(path_str, Permissions::from_mode(0o777))?;
        client.connect("/var/run/chrony/chronyd.sock")?;
        Ok(client)
    }
}

#[derive(Debug, Hash)]
enum ServerAddr {
    Udp(SocketAddrV6),
    #[cfg(unix)]
    Unix,
}

type ReplySender = tokio::sync::oneshot::Sender<std::io::Result<Reply>>;
type ReplyReceiver = tokio::sync::oneshot::Receiver<std::io::Result<Reply>>;
#[derive(Debug)]
struct RequestMeta {
    body: RequestBody,
    reply_sender: ReplySender,
    server: ServerAddr,
}

type RequestSender = tokio::sync::mpsc::UnboundedSender<RequestMeta>;
type RequestReceiver = tokio::sync::mpsc::UnboundedReceiver<RequestMeta>;

#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
enum ServerKey {
    Udp(SocketAddr),
    #[cfg(unix)]
    Unix,
}

#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
struct InflightKey {
    server_key: ServerKey,
    sequence: u32,
}

#[derive(Debug)]
struct InflightValue {
    request: Vec<u8>,
    attempt: u16,
    reply_sender: ReplySender,
    server: ServerAddr,
}

/// Asynchronously sends requests and receives replies
#[derive(Debug)]
pub struct Client {
    task_handle: tokio::task::JoinHandle<()>,
    sender: RequestSender,
}

/// A future which can be `await`ed to obtain the reply to a sent query
#[derive(Debug)]
pub struct ReplyFuture(ReplyReceiver);

impl Future for ReplyFuture {
    type Output = std::io::Result<Reply>;

    fn poll(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        let receiver = &mut self.get_mut().0;
        let mut result = receiver.unwrap_or_else(|e| {
            Err(std::io::Error::new(
                std::io::ErrorKind::ConnectionAborted,
                e,
            ))
        });
        Pin::new(&mut result).poll(cx)
    }
}

impl Client {
    /// Spawns a task to handle sending of requests and receiving of replies, and returns a `Client`
    /// that communicates with this task.
    pub fn spawn(handle: &tokio::runtime::Handle, options: crate::net::ClientOptions) -> Client {
        let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
        let task_handle = handle.spawn(client_task(options, receiver));
        Client {
            task_handle,
            sender,
        }
    }

    /// Sends a request to the given server, and returns a future which can be `await`ed to obtain the reply.
    ///
    /// Note that this is *not* an async function; it is a synchronous, non-blocking function which returns a future. Even
    /// if you don't immediately await the returned future, the request will still immediately be dispatched
    /// to the task which was spawned to service this client, which will immediately form it into a packet
    /// and send it.
    ///
    /// The type of the `server` parameter is more restrictive than that of this method's blocking counterpart
    /// [crate::net::blocking_query] because converting some implementations of `ToSocketAddrs` into a `SocketAddr` can
    /// involve blocking on a DNS lookup, which here would be unacceptable. If you're communicating with localhost,
    /// can call this method with `LOCAL_SERVER_ADDR.into()`. If you're communicating with a remote server,
    /// you'll need to handle the DNS lookup yourself.
    pub fn query(&self, request: RequestBody, server: SocketAddr) -> ReplyFuture {
        let mapped_server = match server {
            SocketAddr::V4(v4) => SocketAddrV6::new(v4.ip().to_ipv6_mapped(), v4.port(), 0, 0),
            SocketAddr::V6(v6) => v6,
        };

        let (sender, receiver) = tokio::sync::oneshot::channel();
        if let Err(SendError(request_meta)) = self.sender.send(RequestMeta {
            body: request,
            reply_sender: sender,
            server: ServerAddr::Udp(mapped_server),
        }) {
            request_meta
                .reply_sender
                .send(Err(std::io::Error::new(
                    std::io::ErrorKind::ConnectionAborted,
                    "Client task unexpectedly shut down",
                )))
                .expect("Send failed but the receiver is still in scope?!")
        }

        ReplyFuture(receiver)
    }

    /// Sends a request to the local Chrony server via its UNIX domain socket, and returns a future which can be `await`ed to obtain the reply.
    #[cfg(unix)]
    pub fn query_uds(&self, request: RequestBody) -> ReplyFuture {
        let (sender, receiver) = tokio::sync::oneshot::channel();

        if let Err(SendError(request_meta)) = self.sender.send(RequestMeta {
            body: request,
            reply_sender: sender,
            server: ServerAddr::Unix,
        }) {
            request_meta
                .reply_sender
                .send(Err(std::io::Error::new(
                    std::io::ErrorKind::ConnectionAborted,
                    "Client task unexpectedly shut down",
                )))
                .expect("Send failed but the receiver is still in scope?!")
        }

        ReplyFuture(receiver)
    }
}

#[derive(Debug)]
struct ReplyMeta<'a> {
    reply: &'a [u8],
    server_key: ServerKey,
}

#[derive(Debug)]
enum SelectResult<'a> {
    Request(RequestMeta),
    Reply(ReplyMeta<'a>),
    Timeout,
    Error(std::io::Error),
    Shutdown,
}

async fn client_task(options: ClientOptions, mut receiver: RequestReceiver) {
    let mut deadlines: VecDeque<(Instant, InflightKey)> = std::collections::VecDeque::new();
    let mut inflight: HashMap<InflightKey, InflightValue> = std::collections::HashMap::new();

    let udp_init = || tokio::net::UdpSocket::bind((Ipv6Addr::UNSPECIFIED, 0));
    let udp_cell: OnceCell<tokio::net::UdpSocket> = OnceCell::new();
    let mut udp_buf = [0u8; 1500];

    #[cfg(unix)]
    let uds_init = || UnixDatagramClient::new();
    #[cfg(unix)]
    let uds_cell: OnceCell<UnixDatagramClient> = OnceCell::new();
    #[cfg(unix)]
    let mut uds_buf = [0u8; 1500];

    let (mut sequence, key0, key1): (u32, u64, u64) = {
        let mut rng = rand::thread_rng();
        (rng.gen(), rng.gen(), rng.gen())
    };

    loop {
        let now = tokio::time::Instant::now();

        // Find expired deadlines and retransmit those messages, or give up
        // if the attempt limit has been reached
        while let Some((deadline, _)) = deadlines.front() {
            if *deadline > now {
                break;
            }
            //`deadline_key` and `inflight_key` have the same value, but
            // we keep both copies in order to minimize further cloning.
            let (_, deadline_key) = deadlines.pop_front().unwrap();
            if let Some((inflight_key, mut inflight_val)) = inflight.remove_entry(&deadline_key) {
                inflight_val.attempt += 1;
                if inflight_val.attempt > options.n_tries {
                    let _ = inflight_val.reply_sender.send(Err(std::io::Error::new(
                        std::io::ErrorKind::TimedOut,
                        "request timed out and max retries reached",
                    )));
                } else {
                    crate::request::increment_attempt(inflight_val.request.as_mut());
                    let send_result = match inflight_val.server {
                        // These are retries, so we can safely unwrap() since we must have
                        // gotten a socket the first time around.
                        ServerAddr::Udp(addr) => {
                            udp_cell
                                .get()
                                .unwrap()
                                .send_to(inflight_val.request.as_ref(), addr)
                                .await
                        }
                        #[cfg(unix)]
                        ServerAddr::Unix => {
                            uds_cell
                                .get()
                                .unwrap()
                                .send(inflight_val.request.as_ref())
                                .await
                        }
                    };
                    match send_result {
                        Ok(_) => {
                            inflight.insert(inflight_key, inflight_val);
                            let new_deadline = now + options.timeout;
                            deadlines.push_back((new_deadline, deadline_key));
                        }
                        Err(e) => {
                            let _ = inflight_val.reply_sender.send(Err(e));
                        }
                    }
                }
            }
        }

        // Cull unexpired deadlines that have already been met. On every iteration of
        // the main loop, pop off the front of the queue until we reach something
        // still in-flight. This takes amortized O(1), and it ensures that the next
        // coming deadline is "real" and prevents unnecessary wakeups.
        while let Some((_, inflight_key)) = deadlines.front() {
            if inflight.contains_key(inflight_key) {
                break;
            } else {
                deadlines.pop_front();
            }
        }

        // In rare cases, like if the network comes back online in the middle of a
        // big burst of queries, there might be a lot of met deadlines still remaining
        // in the queue. Finding these takes O(N), so we don't want to do it every time.
        // Instead, we do a complete cull only when the length of `deadlines` reaches
        // twice the capacity of `inflight`. This keeps us at amortized O(1) and also
        // keeps the deadline queue from asymptotically outgrowing the inflight table.
        if deadlines.len() >= 2 * inflight.capacity() {
            deadlines.retain(|(_, inflight_key)| inflight.contains_key(inflight_key))
        }

        let timeout = async {
            match deadlines.front() {
                Some((deadline, _)) => tokio::time::sleep_until(*deadline).await,
                None => futures::future::pending().await,
            }
        };

        let udp_recv = async {
            match udp_cell.get() {
                Some(udp) => {
                    let (size, peer) = udp.recv_from(&mut udp_buf).await?;
                    std::io::Result::Ok(ReplyMeta {
                        reply: &udp_buf[0..size],
                        server_key: ServerKey::Udp(peer),
                    })
                }
                _ => futures::future::pending().await,
            }
        };

        #[cfg(unix)]
        let uds_recv = async {
            match uds_cell.get() {
                Some(uds) => {
                    let size = uds.recv(&mut uds_buf).await?;
                    std::io::Result::Ok(ReplyMeta {
                        reply: &uds_buf[0..size],
                        server_key: ServerKey::Unix,
                    })
                }
                _ => futures::future::pending().await,
            }
        };
        #[cfg(not(unix))]
        let uds_recv = futures::future::pending();

        let select_result = tokio::select! {
            result = udp_recv => match result {
                Ok(reply_meta) => {
                    SelectResult::Reply(reply_meta)
                },
                Err(e) => SelectResult::Error(e),
            },
            result = uds_recv => match result {
                Ok(reply_meta) => {
                    SelectResult::Reply(reply_meta)
                },
                Err(e) => SelectResult::Error(e),
            },
            result = receiver.recv() => {
                match result {
                    Some(request) => SelectResult::Request(request),
                    None => SelectResult::Shutdown,
                }
            },
            _ = timeout => SelectResult::Timeout
        };

        match select_result {
            SelectResult::Request(request_meta) => {
                // Sequnce numbers should be unpredictable in order to make
                // off-path blind spoofing harder. A single global sequence
                // number accumulator with a random initial state is mixed
                // with a server-specific key derived using SipHash. This
                // way we don't have to keep state for each individual server,
                // sequence numbers sent to one server can't be used to guess
                // another's, and we avoid the birthday-bound collisions that
                // we'd get from picking a random number for every request.
                let mut hasher = siphasher::sip::SipHasher::new_with_keys(key0, key1);
                request_meta.server.hash(&mut hasher);
                let obfuscated_sequence = sequence.wrapping_add(hasher.finish() as u32);
                sequence = sequence.wrapping_add(1);

                let request = Request {
                    sequence: obfuscated_sequence,
                    attempt: 0,
                    body: request_meta.body,
                };
                let mut send_buf = Vec::with_capacity(request.length());
                request.serialize(&mut send_buf);

                let inflight_key = InflightKey {
                    server_key: match request_meta.server {
                        ServerAddr::Udp(addr) => ServerKey::Udp(addr.into()),
                        #[cfg(unix)]
                        ServerAddr::Unix => ServerKey::Unix,
                    },
                    sequence: obfuscated_sequence,
                };

                let inflight_val = InflightValue {
                    request: send_buf,
                    attempt: 0,
                    reply_sender: request_meta.reply_sender,
                    server: request_meta.server,
                };

                let deadline = now + options.timeout;

                match inflight_val.server {
                    ServerAddr::Udp(addr) => match udp_cell.get_or_try_init(udp_init).await {
                        Ok(udp) => {
                            if let Err(e) = udp.send_to(inflight_val.request.as_ref(), addr).await {
                                let _ = inflight_val.reply_sender.send(Err(e));
                                continue;
                            }
                        }
                        Err(e) => {
                            let _ = inflight_val.reply_sender.send(Err(e));
                            continue;
                        }
                    },
                    #[cfg(unix)]
                    ServerAddr::Unix => match uds_cell.get_or_try_init(uds_init).await {
                        Ok(uds) => {
                            if let Err(e) = uds.send(inflight_val.request.as_ref()).await {
                                let _ = inflight_val.reply_sender.send(Err(e));
                                continue;
                            }
                        }
                        Err(e) => {
                            let _ = inflight_val.reply_sender.send(Err(e));
                            continue;
                        }
                    },
                }

                deadlines.push_back((deadline, inflight_key.clone()));
                inflight.insert(inflight_key.clone(), inflight_val);
            }
            SelectResult::Reply(reply_meta) => {
                let mut reply_buf = reply_meta.reply;
                if let Ok(reply) = Reply::deserialize(&mut reply_buf) {
                    let inflight_key = InflightKey {
                        server_key: reply_meta.server_key,
                        sequence: reply.sequence,
                    };
                    if let Some(inflight_val) = inflight.remove(&inflight_key) {
                        let _ = inflight_val.reply_sender.send(Ok(reply));
                    }
                }
            }
            SelectResult::Timeout => {}
            SelectResult::Error(e) => {
                if e.kind() == std::io::ErrorKind::Interrupted {
                    continue;
                }

                // Any other kind of error should never happen here, but if it
                // does we have to bail.
                let erc = Arc::new(e);
                receiver.close();
                // Drain any requests from the channel and answer them all with the error
                // we got.
                while let Some(request) = receiver.recv().await {
                    let _ = request
                        .reply_sender
                        .send(Err(std::io::Error::new(erc.kind(), erc.clone())));
                }
                // Also answer in-flight requests with an error.
                for v in inflight.into_values() {
                    let _ = v
                        .reply_sender
                        .send(Err(std::io::Error::new(erc.kind(), erc.clone())));
                }
                return;
            }
            SelectResult::Shutdown => {
                for v in inflight.into_values() {
                    let _ = v.reply_sender.send(Err(std::io::Error::new(
                        std::io::ErrorKind::ConnectionAborted,
                        "Client dropped before arrival of reply",
                    )));
                }
                return;
            }
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// Portions derived from Chrony copyright Richard P. Curnow 1997-2003
// and Miroslav Lichvar 2009, 2012-2020
// SPDX-License-Identifier: GPL-2.0-only

//! Data structures occurring both in requests and in replies

use bitflags::bitflags;
use bytes::{Buf, BufMut};
use std::convert::From;
use std::error::Error;
use std::fmt::{Debug, Display};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use chrony_candm_derive::ChronySerialize;

///Error returned when attempting to deserialize a malformed message
#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct DeserializationError(Option<&'static str>);

impl DeserializationError {
    pub(crate) fn new(detail: &'static str) -> DeserializationError {
        DeserializationError(Some(detail))
    }

    #[allow(dead_code)]
    pub(crate) fn generic() -> DeserializationError {
        DeserializationError(None)
    }
}

impl Display for DeserializationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.0 {
            None => write!(f, "Deserialization error"),
            Some(detail) => write!(f, "Deserialization error: {}", detail),
        }
    }
}
impl Error for DeserializationError {}

pub(crate) trait ChronySerialize: Sized {
    fn length() -> usize;
    fn serialize<B: BufMut>(&self, buf: &mut B);
    fn deserialize_unchecked<B: Buf>(buf: &mut B) -> Result<Self, DeserializationError>;

    fn deserialize<B: Buf>(buf: &mut B) -> Result<Self, DeserializationError> {
        if buf.remaining() < Self::length() {
            Err(DeserializationError::new("message too short"))
        } else {
            Self::deserialize_unchecked(buf)
        }
    }
}

impl<T: ChronySerialize + Copy, const N: usize> ChronySerialize for [T; N] {
    fn length() -> usize {
        T::length() * N
    }

    fn serialize<B: BufMut>(&self, buf: &mut B) {
        for i in 0..N {
            self[i].serialize(buf)
        }
    }

    fn deserialize_unchecked<B: Buf>(buf: &mut B) -> Result<Self, DeserializationError> {
        let mut arr = arrayvec::ArrayVec::<T, N>::new_const();
        unsafe {
            for _ in 0..N {
                //Safety: We push at most N elements
                arr.push_unchecked(T::deserialize_unchecked(buf)?);
            }
            //Safety: we've pushed exactly N elements
            Ok(arr.into_inner_unchecked())
        }
    }
}

macro_rules! serialize_primitive {
    ($ty:ty, $get:ident, $put:ident) => {
        impl ChronySerialize for $ty {
            fn length() -> usize {
                std::mem::size_of::<$ty>()
            }

            fn serialize<B: BufMut>(&self, buf: &mut B) {
                buf.$put(*self)
            }

            fn deserialize_unchecked<B: Buf>(buf: &mut B) -> Result<Self, DeserializationError> {
                Ok(buf.$get())
            }
        }
    };
}

serialize_primitive!(u8, get_u8, put_u8);
serialize_primitive!(i8, get_i8, put_i8);
serialize_primitive!(u16, get_u16, put_u16);
serialize_primitive!(i16, get_i16, put_i16);
serialize_primitive!(u32, get_u32, put_u32);
serialize_primitive!(i32, get_i32, put_i32);
serialize_primitive!(u64, get_u64, put_u64);
serialize_primitive!(i64, get_i64, put_i64);

impl ChronySerialize for SystemTime {
    fn length() -> usize {
        12
    }

    fn serialize<B: BufMut>(&self, buf: &mut B) {
        let (secs, nsecs) = match self.duration_since(UNIX_EPOCH) {
            Ok(d) => (d.as_secs() as i64, d.subsec_nanos()),
            Err(e) => {
                let d = e.duration();
                let secs = d.as_secs() as i64;
                let nsecs = d.subsec_nanos();
                if nsecs == 0 {
                    (-secs, 0)
                } else {
                    (-secs - 1, 1_000_000_000 - nsecs)
                }
            }
        };

        let secs_high = (secs >> 32) as i32;
        let secs_low = (secs & 0xffffffff) as u32;

        buf.put_i32(secs_high);
        buf.put_u32(secs_low);
        buf.put_u32(nsecs);
    }

    fn deserialize_unchecked<B: Buf>(buf: &mut B) -> Result<Self, DeserializationError> {
        let secs_high_orig = buf.get_i32();
        let secs_low = buf.get_u32();
        let nsecs = buf.get_u32();

        let secs_high = if secs_high_orig == i32::MAX {
            0
        } else {
            secs_high_orig
        };

        if secs_high >= 0 {
            let d = Duration::from_secs((secs_high as u64) << 32)
                + Duration::from_secs(secs_low as u64)
                + Duration::from_nanos(nsecs as u64);
            Ok(UNIX_EPOCH + d)
        } else {
            let d = Duration::from_secs(((-secs_high) as u64) << 32)
                - Duration::from_secs(secs_low as u64)
                - Duration::from_nanos(nsecs as u64);
            Ok(UNIX_EPOCH - d)
        }
    }
}
pub(crate) trait ChronyMessage: Sized {
    fn body_length(&self) -> usize;
    fn cmd(&self) -> u16;
    fn serialize_body<B: BufMut>(&self, buf: &mut B);
    fn deserialize_body<B: Buf>(cmd: u16, body: &mut B) -> Result<Self, DeserializationError>;
}

/// Floating point number as used in Chrony's wire protocol
///
/// A `ChronyFloat` can be converted infallibly to and from an `f64`,
/// but the conversion into `ChronyFloat` is lossy and sometimes
/// nonsensical, e.g., NaNs are represented as 0. These conversion
/// rules are identical to the ones that Chrony uses internally.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Default)]
pub struct ChronyFloat(u32);

impl Debug for ChronyFloat {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("ChronyFloat")
            .field(&f64::from(*self))
            .finish()
    }
}

impl Display for ChronyFloat {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        Display::fmt(&f64::from(*self), f)
    }
}

impl PartialOrd for ChronyFloat {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        f64::from(*self).partial_cmp(&f64::from(*other))
    }
}

impl Ord for ChronyFloat {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.partial_cmp(other)
            .expect("Conversion from ChronyFloat yielded incomparable f64")
    }
}

const FLOAT_EXP_BITS: u32 = 7;
const FLOAT_EXP_MIN: i32 = -((1 << (FLOAT_EXP_BITS - 1)) as i32);
const FLOAT_EXP_MAX: i32 = -FLOAT_EXP_MIN - 1;
const FLOAT_COEF_BITS: u32 = 32 - FLOAT_EXP_BITS;
const FLOAT_COEF_MIN: i32 = -((1 << (FLOAT_COEF_BITS - 1)) as i32);
const FLOAT_COEF_MAX: i32 = -FLOAT_COEF_MIN - 1;

impl From<ChronyFloat> for f64 {
    fn from(f: ChronyFloat) -> Self {
        let ChronyFloat(x) = f;

        let mut exp = (x >> FLOAT_COEF_BITS) as i32;
        if exp >= 1 << (FLOAT_EXP_BITS - 1) {
            exp -= 1 << FLOAT_EXP_BITS;
        }
        exp -= FLOAT_COEF_BITS as i32;

        let mut coef = (x % (1 << FLOAT_COEF_BITS)) as i32;

        if coef >= 1 << (FLOAT_COEF_BITS - 1) {
            coef -= 1 << FLOAT_COEF_BITS
        }

        (coef as f64) * 2.0f64.powi(exp)
    }
}

impl From<f64> for ChronyFloat {
    fn from(f: f64) -> Self {
        let (neg, x) = if f < 0. {
            (1, -f)
        } else if f >= 0. {
            (0, f)
        } else {
            (0, 0.) //Treat NaN as zero
        };

        let mut exp: i32;
        let mut coef: i32;

        if x < 1e-100 {
            coef = 0;
            exp = 0;
        } else if x > 1e100 {
            coef = FLOAT_COEF_MAX + neg;
            exp = FLOAT_EXP_MAX;
        } else {
            exp = x.log2() as i32 + 1;
            coef = (x * 2.0f64.powi(-exp + FLOAT_COEF_BITS as i32) + 0.5) as i32;
            debug_assert!(coef > 0);

            while coef > FLOAT_COEF_MAX + neg {
                coef >>= 1;
                exp += 1;
            }

            if exp > FLOAT_EXP_MAX {
                exp = FLOAT_EXP_MAX;
                coef = FLOAT_COEF_MAX + neg;
            } else if exp < FLOAT_EXP_MIN {
                if exp + FLOAT_COEF_BITS as i32 >= FLOAT_EXP_MIN {
                    coef >>= FLOAT_EXP_MIN - exp;
                    exp = FLOAT_EXP_MIN;
                } else {
                    exp = 0;
                    coef = 0;
                }
            }
        };

        if neg != 0 {
            coef = ((-coef as u32) << FLOAT_EXP_BITS >> FLOAT_EXP_BITS) as i32;
        }

        ChronyFloat((exp as u32) << FLOAT_COEF_BITS | coef as u32)
    }
}

impl ChronySerialize for ChronyFloat {
    fn length() -> usize {
        4
    }

    fn serialize<B: BufMut>(&self, buf: &mut B) {
        self.0.serialize(buf)
    }

    fn deserialize_unchecked<B: Buf>(buf: &mut B) -> Result<Self, DeserializationError> {
        Ok(ChronyFloat(u32::deserialize_unchecked(buf)?))
    }
}

/// A network address as used in Chrony's wire protocol
#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash)]
pub enum ChronyAddr {
    /// An unknown or unspecified address
    Unspec,
    /// An IPv4 address
    V4(Ipv4Addr),
    /// An IPv6 address
    V6(Ipv6Addr),
    /// A placeholder for an address that has not yet been resolved
    Id(u32),
}

impl Default for ChronyAddr {
    fn default() -> Self {
        Self::Unspec
    }
}

impl From<Ipv4Addr> for ChronyAddr {
    fn from(addr: Ipv4Addr) -> Self {
        Self::V4(addr)
    }
}

impl From<Ipv6Addr> for ChronyAddr {
    fn from(addr: Ipv6Addr) -> Self {
        Self::V6(addr)
    }
}

impl From<IpAddr> for ChronyAddr {
    fn from(addr: IpAddr) -> Self {
        match addr {
            IpAddr::V4(addr4) => Self::V4(addr4),
            IpAddr::V6(addr6) => Self::V6(addr6),
        }
    }
}

impl Display for ChronyAddr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ChronyAddr::Unspec => write!(f, "[UNSPEC]"),
            ChronyAddr::V4(v4) => Display::fmt(v4, f),
            ChronyAddr::V6(v6) => Display::fmt(v6, f),
            ChronyAddr::Id(id) => write!(f, "ID#{}", id),
        }
    }
}

const IPADDR_UNSPEC: u16 = 0;
const IPADDR_INET4: u16 = 1;
const IPADDR_INET6: u16 = 2;
const IPADDR_ID: u16 = 3;

impl ChronySerialize for ChronyAddr {
    fn length() -> usize {
        20
    }

    fn serialize<B: BufMut>(&self, buf: &mut B) {
        match self {
            ChronyAddr::Unspec => {
                buf.put(&[0u8; 16] as &[u8]);
                buf.put_u16(IPADDR_UNSPEC);
                buf.put_u16(0);
            }
            ChronyAddr::V4(addr) => {
                buf.put_u32((*addr).into());
                buf.put(&[0u8; 12] as &[u8]);
                buf.put_u16(IPADDR_INET4);
                buf.put_u16(0);
            }
            ChronyAddr::V6(addr) => {
                buf.put_u128((*addr).into());
                buf.put_u16(IPADDR_INET6);
                buf.put_u16(0);
            }
            ChronyAddr::Id(id) => {
                buf.put_u32(*id);
                buf.put(&[0u8; 12] as &[u8]);
                buf.put_u16(IPADDR_ID);
                buf.put_u16(0);
            }
        }
    }

    fn deserialize_unchecked<B: Buf>(buf: &mut B) -> Result<Self, DeserializationError> {
        let mut addr = buf.copy_to_bytes(16);
        let family = buf.get_u16();
        buf.get_u16();

        match family {
            IPADDR_INET4 => Ok(Self::V4(Ipv4Addr::from(addr.get_u32()))),
            IPADDR_INET6 => Ok(Self::V6(Ipv6Addr::from(addr.get_u128()))),
            IPADDR_ID => Ok(Self::Id(addr.get_u32())),
            _ => Ok(Self::Unspec),
        }
    }
}

bitflags! {
    /// Flags associated with a time source
    #[derive(ChronySerialize)]
    pub struct SourceFlags : u16 {
        const ONLINE = 0x1;
        const AUTOOFFLINE = 0x2;
        const IBURST = 0x4;
        const PREFER = 0x8;
        const NOSELECT = 0x10;
        const TRUST = 0x20;
        const REQUIRE = 0x40;
        const INTERLEAVED = 0x80;
        const BURST = 0x100;
        const NTS = 0x200;
        const COPY = 0x400;
    }
}

pub(crate) const REQUEST_HEADER_LENGTH: usize = 20;
pub(crate) const REPLY_HEADER_LENGTH: usize = 28;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: GPL-2.0-only

//! This crate is a library for communicating with [Chrony](https://chrony.tuxfamily.org)'s
//! control & monitoring interface. It provides programmatic access to information that you'd
//! otherwise have to scrape from the output of `chronyc`.
//!
//! For simple use cases which involve non-privileged queries to the locally-running Chrony
//! daemon, your main entry-point into the library will be the [blocking_query] function with
//! a `server` argument of `&LOCAL_SERVER_ADDR`. For privileged commands, use [blocking_query_uds]
//! instead (this will require permissions to write to the `/var/run/chrony` directory).
//!
//! ```no_run
//! use chrony_candm::request::RequestBody;
//! use chrony_candm::reply::ReplyBody;
//! use chrony_candm::{blocking_query,LOCAL_SERVER_ADDR};
//!
//! let request_body = RequestBody::Tracking;
//! let options = Default::default();
//! let reply = blocking_query(request_body, options, &LOCAL_SERVER_ADDR)?;
//! if let ReplyBody::Tracking(tracking) = reply.body {
//!     println!("The current RMS offset is {} seconds.", tracking.rms_offset);
//! }
//! # Ok::<(), std::io::Error>(())
//! ```
//!
//! Asynchronous applications can use [Client] instead of the standalone
//! `blocking_query` and `blocking_query_uds` functions.

extern crate self as chrony_candm; //We need this so that the fully qualified identifiers generated by macros can resolve

#[cfg(feature = "with_tokio")]
mod async_net;
pub mod common;
mod net;
pub mod reply;
pub mod request;

#[cfg(feature = "with_tokio")]
pub use async_net::*;
pub use net::*;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: GPL-2.0-only

use bytes::BytesMut;
use rand::Rng;

use crate::reply::Reply;
use crate::request::{increment_attempt, Request, RequestBody};

use std::ffi::OsStr;
use std::fs::Permissions;
use std::net::Ipv6Addr;
use std::ops::{Deref, DerefMut};
use std::time::Duration;

#[cfg(unix)]
use std::os::unix::net::UnixDatagram;
#[cfg(unix)]
use std::os::unix::ffi::OsStrExt;
#[cfg(unix)]
use std::os::unix::fs::PermissionsExt;

pub const DEFAULT_PORT: u16 = 323;

/// Options for configuring a Chrony client
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ClientOptions {
    /// How long to wait for a reply before assuming the request was dropped (default: 1s)
    pub timeout: Duration,
    /// Number of times to try a request before giving up (default: 3)
    pub n_tries: u16,
}

impl Default for ClientOptions {
    fn default() -> Self {
        Self {
            n_tries: 3,
            timeout: Duration::from_secs(1),
        }
    }
}

///Common interface to `UdpSocket` and `UnixDatagram` to avoid copypasta between our two `blocking_query` functions
trait DgramSocket {
    fn send(&self, buf: &[u8]) -> std::io::Result<usize>;
    fn recv(&self, buf: &mut [u8]) -> std::io::Result<usize>;
}

impl DgramSocket for std::net::UdpSocket {
    fn send(&self, buf: &[u8]) -> std::io::Result<usize> {
        std::net::UdpSocket::send(self, buf)
    }

    fn recv(&self, buf: &mut [u8]) -> std::io::Result<usize> {
        std::net::UdpSocket::recv(self, buf)
    }
}

#[cfg(unix)]
impl DgramSocket for UnixDatagram {
    fn send(&self, buf: &[u8]) -> std::io::Result<usize> {
        UnixDatagram::send(self, buf)
    }

    fn recv(&self, buf: &mut [u8]) -> std::io::Result<usize> {
        UnixDatagram::recv(self, buf)
    }
}

#[cfg(unix)]
#[derive(Debug)]
struct UnixDatagramClient(UnixDatagram);

#[cfg(unix)]
impl AsRef<UnixDatagram> for UnixDatagramClient {
    fn as_ref(&self) -> &UnixDatagram {
        &self.0
    }
}

#[cfg(unix)]
impl AsMut<UnixDatagram> for UnixDatagramClient {
    fn as_mut(&mut self) -> &mut UnixDatagram {
        &mut self.0
    }
}

#[cfg(unix)]
impl Deref for UnixDatagramClient {
    type Target = UnixDatagram;
    fn deref(&self) -> &UnixDatagram {
        &self.0
    }
}

#[cfg(unix)]
impl DerefMut for UnixDatagramClient {
    fn deref_mut(&mut self) -> &mut UnixDatagram {
        &mut self.0
    }
}

#[cfg(unix)]
impl Drop for UnixDatagramClient {
    fn drop(&mut self) {
        if let Ok(addr) = self.0.local_addr() {
            if let Some(path) = addr.as_pathname() {
                let _ = self.0.shutdown(std::net::Shutdown::Both);
                let _ = std::fs::remove_file(path);
            }
        }
    }
}

#[cfg(unix)]
impl UnixDatagramClient {
    fn new() -> std::io::Result<UnixDatagramClient> {
        let id: [u8; 16] = rand::random();
        let mut path = b"/var/run/chrony/client-000102030405060708090a0b0c0d0e0f.sock".clone();
        hex::encode_to_slice(id, &mut path[23..55]).unwrap();
        let path_str = OsStr::from_bytes(&path);
        let sock = UnixDatagram::bind(path_str)?;
        let client = UnixDatagramClient(sock);
        std::fs::set_permissions(path_str, Permissions::from_mode(0o777))?;
        client.connect("/var/run/chrony/chronyd.sock")?;
        Ok(client)
    }
}

fn blocking_query_loop<Sock: DgramSocket>(
    sock: &Sock,
    request_body: RequestBody,
    options: ClientOptions,
) -> std::io::Result<Reply> {
    let request = Request {
        sequence: rand::thread_rng().gen(),
        attempt: 0,
        body: request_body,
    };

    let mut send_buf = BytesMut::with_capacity(request.length());
    request.serialize(&mut send_buf);
    let mut recv_buf = [0; 1500];
    let mut attempt = 0;

    loop {
        sock.send(send_buf.as_ref())?;
        match sock.recv(&mut recv_buf) {
            Ok(len) => {
                let mut msg = &recv_buf[0..len];
                match Reply::deserialize(&mut msg) {
                    Ok(reply) => {
                        if reply.sequence == request.sequence {
                            return Ok(reply);
                        } else {
                            return Err(std::io::Error::new(
                                std::io::ErrorKind::InvalidData,
                                "Bad sequence number",
                            ));
                        }
                    }
                    Err(e) => return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, e)),
                }
            }
            Err(e) => {
                if e.kind() == std::io::ErrorKind::TimedOut || e.kind() == std::io::ErrorKind::WouldBlock {
                    attempt += 1;
                    if attempt == options.n_tries {
                        return Err(e);
                    }
                    increment_attempt(send_buf.as_mut());
                } else {
                    return Err(e);
                }
            }
        }
    }
}

/// Sends a request to a server via UDP and waits for a reply
pub fn blocking_query<Server: std::net::ToSocketAddrs>(
    request_body: RequestBody,
    options: ClientOptions,
    server: &Server,
) -> std::io::Result<Reply> {
    let sock = std::net::UdpSocket::bind((Ipv6Addr::UNSPECIFIED, 0))?;
    sock.connect(server)?;
    sock.set_read_timeout(Some(options.timeout))?;

    blocking_query_loop(&sock, request_body, options)
}

/// Sends a request to a server via a domain socket and waits for a reply
#[cfg(unix)]
pub fn blocking_query_uds(
    request_body: RequestBody,
    options: ClientOptions
) -> std::io::Result<Reply> {
    let sock = UnixDatagramClient::new()?;
    sock.set_read_timeout(Some(options.timeout))?;
    blocking_query_loop(sock.as_ref(), request_body, options)
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// Portions derived from Chrony copyright Richard P. Curnow 1997-2003
// and Miroslav Lichvar 2009, 2012-2020
// SPDX-License-Identifier: GPL-2.0-only

//! Data structures representing replies

use bitflags::bitflags;
use bytes::{Buf, BufMut};
use chrony_candm_derive::{ChronyMessage, ChronySerialize};
use num_enum::{IntoPrimitive, TryFromPrimitive};
use std::time::{SystemTime, UNIX_EPOCH};

use crate::common::*;

#[repr(u16)]
#[derive(
    Debug,
    Copy,
    Clone,
    PartialEq,
    Eq,
    Ord,
    PartialOrd,
    Hash,
    IntoPrimitive,
    TryFromPrimitive,
    ChronySerialize,
)]
pub enum Status {
    Success,
    Failed,
    Unauth,
    Invalid,
    NoSuchSource,
    InvalidTs,
    NotEnabled,
    BadSubnet,
    AccessAllowed,
    AccessDenied,
    NoHostAccess,
    SourceAlreadyKnown,
    TooManySources,
    NoRtc,
    BadRtcFile,
    Inactive,
    BadSample,
    InvalidAf,
    BadPktVersion,
    BadPktLength,
    InvalidName,
}
#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct NSources {
    pub n_sources: u32,
}

#[repr(u16)]
#[derive(
    Debug,
    Copy,
    Clone,
    PartialEq,
    Eq,
    Ord,
    PartialOrd,
    Hash,
    IntoPrimitive,
    TryFromPrimitive,
    ChronySerialize,
)]
pub enum SourceMode {
    Client,
    Peer,
    Ref,
}

#[repr(u16)]
#[derive(
    Debug,
    Copy,
    Clone,
    PartialEq,
    Eq,
    Ord,
    PartialOrd,
    Hash,
    IntoPrimitive,
    TryFromPrimitive,
    ChronySerialize,
)]
pub enum SourceState {
    Selected,
    NonSelectable,
    Falseticker,
    Jittery,
    Unselected,
    Selectable,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct SourceData {
    pub ip_addr: ChronyAddr,
    pub poll: i16,
    pub stratum: u16,
    pub state: SourceState,
    pub mode: SourceMode,
    pub flags: SourceFlags,
    pub reachability: u16,
    pub since_sample: u32,
    pub orig_latest_meas: ChronyFloat,
    pub latest_meas: ChronyFloat,
    pub latest_meas_err: ChronyFloat,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct Tracking {
    pub ref_id: u32,
    pub ip_addr: ChronyAddr,
    pub stratum: u16,
    pub leap_status: u16,
    pub ref_time: SystemTime,
    pub current_correction: ChronyFloat,
    pub last_offset: ChronyFloat,
    pub rms_offset: ChronyFloat,
    pub freq_ppm: ChronyFloat,
    pub resid_freq_ppm: ChronyFloat,
    pub skew_ppm: ChronyFloat,
    pub root_delay: ChronyFloat,
    pub root_dispersion: ChronyFloat,
    pub last_update_interval: ChronyFloat,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct SourceStats {
    pub ref_id: u32,
    pub ip_addr: ChronyAddr,
    pub n_samples: u32,
    pub n_runs: u32,
    pub span_seconds: u32,
    pub sd: ChronyFloat,
    pub resid_freq_ppm: ChronyFloat,
    pub skew_ppm: ChronyFloat,
    pub est_offset: ChronyFloat,
    pub est_offset_err: ChronyFloat,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct Rtc {
    pub ref_time: SystemTime,
    pub n_samples: u16,
    pub n_runs: u16,
    pub span_seconds: u32,
    pub rtc_seconds_fast: ChronyFloat,
    pub rtc_gain_rate_ppm: ChronyFloat,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct ManualTimestamp {
    pub offset: ChronyFloat,
    pub dfreq_ppm: ChronyFloat,
    pub new_afreq_ppm: ChronyFloat,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize, Default)]
pub struct ClientAccessesClient {
    pub ip: ChronyAddr,
    pub ntp_hits: u32,
    pub nke_hits: u32,
    pub cmd_hits: u32,
    pub ntp_drops: u32,
    pub nke_drops: u32,
    pub cmd_drops: u32,
    pub ntp_interval: u8,
    pub nke_interval: u8,
    pub cmd_interval: u8,
    pub ntp_timeout_interval: u8,
    pub last_ntp_hit_ago: u32,
    pub last_nke_hit_ago: u32,
    pub last_cmd_hit_ago: u32,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct ClientAccessesByIndex {
    pub n_indices: u32,
    pub next_index: u32,
    pub n_clients: u32,
    pub clients: [ClientAccessesClient; 8],
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct ServerStats {
    pub ntp_hits: u32,
    pub nke_hits: u32,
    pub cmd_hits: u32,
    pub ntp_drops: u32,
    pub nke_drops: u32,
    pub cmd_drops: u32,
    pub log_drops: u32,
    pub ntp_auth_hits: u32,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct ManualListSample {
    pub when: SystemTime,
    pub slewed_offset: ChronyFloat,
    pub orig_offset: ChronyFloat,
    pub residual: ChronyFloat,
}

impl Default for ManualListSample {
    fn default() -> Self {
        Self {
            when: UNIX_EPOCH,
            slewed_offset: Default::default(),
            orig_offset: Default::default(),
            residual: Default::default(),
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct ManualList {
    pub n_samples: u32,
    pub samples: [ManualListSample; 16],
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct Activity {
    pub online: u32,
    pub offline: u32,
    pub burst_online: u32,
    pub burst_offline: u32,
    pub unresolved: u32,
}

bitflags! {
    #[derive(ChronySerialize)]
    pub struct SmoothingFlags : u32 {
        const ACTIVE = 0x1;
        const LEAPONLY = 0x2;
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct Smoothing {
    pub flags: SmoothingFlags,
    pub offset: ChronyFloat,
    pub freq_ppm: ChronyFloat,
    pub wander_ppm: ChronyFloat,
    pub last_update_ago: ChronyFloat,
    pub remaining_time: ChronyFloat,
}

bitflags! {
    #[derive(ChronySerialize)]
    pub struct NtpFlags : u16 {
        const TESTS = 0x3ff;
        const INTERLEAVED = 0x4000;
        const AUTHENTICATED = 0x8000;
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct NtpData {
    pub remote_addr: ChronyAddr,
    pub local_addr: ChronyAddr,
    pub remote_port: u16,
    pub leap: u8,
    pub version: u8,
    pub mode: u8,
    pub stratum: u8,
    pub poll: i8,
    pub precision: i8,
    pub root_delay: ChronyFloat,
    pub root_dispersion: ChronyFloat,
    pub ref_id: u32,
    pub ref_time: SystemTime,
    pub offset: ChronyFloat,
    pub peer_delay: ChronyFloat,
    pub peer_dispersion: ChronyFloat,
    pub response_time: ChronyFloat,
    pub jitter_asymmetry: ChronyFloat,
    pub flags: NtpFlags,
    pub tx_tss_char: u8,
    pub rx_tss_char: u8,
    pub total_tx_count: u32,
    pub total_rx_count: u32,
    #[pad = 16]
    pub total_valid_count: u32,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct NtpSourceName {
    pub name: [u8; 256],
}

#[repr(u16)]
#[derive(
    Debug,
    Copy,
    Clone,
    PartialEq,
    Eq,
    Ord,
    PartialOrd,
    Hash,
    IntoPrimitive,
    TryFromPrimitive,
    ChronySerialize,
)]
pub enum AuthDataMode {
    None,
    Symmetric,
    Nts,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct AuthData {
    pub mode: AuthDataMode,
    pub key_type: u16,
    pub key_id: u32,
    pub key_length: u16,
    pub ke_attempts: u16,
    pub last_ke_ago: u32,
    pub cookies: u16,
    pub cookies_length: u16,
    #[pad = 2]
    pub nak: u16,
}

bitflags! {
    #[derive(ChronySerialize)]
    pub struct SelectDataOptions : u16 {
        const NOSELECT = 0x1;
        const PREFER = 0x2;
        const TRUST = 0x4;
        const REQUIRE = 0x8;
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronySerialize)]
pub struct SelectData {
    pub ref_id: u32,
    pub ip_addr: ChronyAddr,
    pub state_char: u8,
    pub authentication: u8,
    pub leap: u8,
    pub pad: u8,
    pub conf_options: SelectDataOptions,
    pub eff_options: SelectDataOptions,
    pub last_sample_ago: u32,
    pub score: ChronyFloat,
    pub lo_limit: ChronyFloat,
    pub hi_limit: ChronyFloat,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, ChronyMessage)]
pub enum ReplyBody {
    #[cmd = 1]
    Null,
    NSources(NSources),
    SourceData(SourceData),
    ManualTimestamp,
    Tracking(Tracking),
    SourceStats(SourceStats),
    Rtc(Rtc),
    #[cmd = 12]
    Activity(Activity),
    Smoothing(Smoothing),
    #[cmd = 16]
    NtpData(NtpData),
    ManualTimestamp2(ManualTimestamp),
    ManualList2(ManualList),
    NtpSourceName(NtpSourceName),
    AuthData(AuthData),
    ClientAccessesByIndex3(ClientAccessesByIndex),
    ServerStats2(ServerStats),
    SelectData(SelectData),
}

/// A reply from Chrony
#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash)]
pub struct Reply {
    /// Status code associated with the reply
    pub status: Status,
    /// Command number associated with the request this is replying to
    pub cmd: u16,
    /// Sequence number associated with the request this is replying to
    pub sequence: u32,
    /// Body of the reply
    pub body: ReplyBody,
}

impl Reply {
    /// Returns the length in bytes of this reply when serialized.
    pub fn length(&self) -> usize {
        REPLY_HEADER_LENGTH + self.body.body_length()
    }

    /// Serializes this reply into `buf`. May panic if `buf` does not
    /// have a capacity of at least `self.length()` and its implementation
    /// does not support automatic resizing.
    pub fn serialize<B: BufMut>(&self, buf: &mut B) {
        let reply = self.body.cmd();

        buf.put_u8(6); //Version
        buf.put_u8(2); //Packet type = reply
        buf.put_u8(0); //Reserved
        buf.put_u8(0); //Reserved
        buf.put_u16(self.cmd); //Requested command, set to 0 since it's just ignored
        buf.put_u16(reply);
        self.status.serialize(buf);
        buf.put_u16(0); //Padding
        buf.put_u16(0); //Padding
        buf.put_u16(0); //Padding
        buf.put_u32(self.sequence);
        buf.put_u32(0); //Padding
        buf.put_u32(0); //Padding
        self.body.serialize_body(buf);
    }

    /// Deserializes a reply from `buf`.
    pub fn deserialize<B: Buf>(buf: &mut B) -> Result<Self, DeserializationError> {
        if buf.remaining() < REPLY_HEADER_LENGTH {
            return Err(DeserializationError::new("message too short"));
        }
        if buf.get_u8() != 6 {
            //Version
            return Err(DeserializationError::new("unsupported version"));
        }

        if buf.get_u8() != 2 {
            //Packet type
            return Err(DeserializationError::new("wrong packet type"));
        }

        buf.get_u8(); //Reserved
        buf.get_u8(); //Reserved
        let cmd = buf.get_u16();
        let reply = buf.get_u16();
        let status = Status::deserialize_unchecked(buf)?;
        buf.get_u16(); //Padding
        buf.get_u16(); //Padding
        buf.get_u16(); //Padding
        let sequence = buf.get_u32();
        buf.get_u32(); //Padding
        buf.get_u32(); //Padding
        let body = ReplyBody::deserialize_body(reply, buf)?;
        Ok(Self {
            status,
            cmd,
            sequence,
            body,
        })
    }
}