PlantCtrl/Software/MainBoard/rust/src/network.rs

use crate::bail;
use crate::config::NetworkConfig;
use crate::fat_error::{ContextExt, FatError, FatResult};
use crate::hal::{PlantHal, HAL};
use crate::mqtt;
use crate::util::mk_static;
use alloc::string::{String, ToString};
use alloc::sync::Arc;
use chrono::{DateTime, Utc};
use core::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
use embassy_executor::Spawner;
use embassy_net::dns::DnsQueryType;
use embassy_net::udp::{PacketMetadata, UdpSocket};
use embassy_net::{DhcpConfig, Runner, Stack, StackResources, StaticConfigV4};
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
use embassy_sync::mutex::{Mutex, MutexGuard};
use embassy_time::{Duration, Timer, WithTimeout};
use option_lock::OptionLock;
use edge_dhcp::{
    io::{self, DEFAULT_SERVER_PORT},
    server::{Server, ServerOptions},
};
use edge_nal::UdpBind;
use edge_nal_embassy::{Udp, UdpBuffers};
use esp_hal::rng::Rng;
use esp_println::println;
use esp_radio::wifi::ap::AccessPointConfig;
use esp_radio::wifi::sta::StationConfig;
use esp_radio::wifi::{AuthenticationMethod, Config, Interface};
use log::{info, warn, error};
use serde::Serialize;
use sntpc::{NtpContext, NtpTimestampGenerator, NtpUdpSocket, get_time};

const NTP_SERVER: &str = "pool.ntp.org";

#[derive(Copy, Clone, Default)]
struct Timestamp {
    stamp: DateTime<Utc>,
}

impl NtpTimestampGenerator for Timestamp {
    fn init(&mut self) {
        self.stamp = DateTime::default();
    }

    fn timestamp_sec(&self) -> u64 {
        self.stamp.timestamp() as u64
    }

    fn timestamp_subsec_micros(&self) -> u32 {
        self.stamp.timestamp_subsec_micros()
    }
}

struct EmbassyNtpSocket<'a, 'b> {
    socket: &'a UdpSocket<'b>,
}

impl<'a, 'b> EmbassyNtpSocket<'a, 'b> {
    fn new(socket: &'a UdpSocket<'b>) -> Self {
        Self { socket }
    }
}

impl NtpUdpSocket for EmbassyNtpSocket<'_, '_> {
    async fn send_to(&self, buf: &[u8], addr: SocketAddr) -> sntpc::Result<usize> {
        self.socket
            .send_to(buf, addr)
            .await
            .map_err(|_| sntpc::Error::Network)?;
        Ok(buf.len())
    }

    async fn recv_from(&self, buf: &mut [u8]) -> sntpc::Result<(usize, SocketAddr)> {
        let (len, metadata) = self
            .socket
            .recv_from(buf)
            .await
            .map_err(|_| sntpc::Error::Network)?;
        let addr = match metadata.endpoint.addr {
            embassy_net::IpAddress::Ipv4(ip) => IpAddr::V4(ip),
            embassy_net::IpAddress::Ipv6(ip) => IpAddr::V6(ip),
        };
        Ok((len, SocketAddr::new(addr, metadata.endpoint.port)))
    }
}

pub async fn sntp(max_wait_ms: u32, stack: Stack<'_>) -> FatResult<DateTime<Utc>> {
    println!("start sntp");
    let mut rx_meta = [PacketMetadata::EMPTY; 16];
    let mut rx_buffer = [0; 4096];
    let mut tx_meta = [PacketMetadata::EMPTY; 16];
    let mut tx_buffer = [0; 4096];

    let mut socket = UdpSocket::new(
        stack,
        &mut rx_meta,
        &mut rx_buffer,
        &mut tx_meta,
        &mut tx_buffer,
    );
    socket.bind(123).context("Could not bind UDP socket")?;

    let context = NtpContext::new(Timestamp::default());
    let ntp_socket = EmbassyNtpSocket::new(&socket);

    let ntp_addrs = stack
        .dns_query(NTP_SERVER, DnsQueryType::A)
        .await
        .context("Failed to resolve DNS")?;

    if ntp_addrs.is_empty() {
        bail!("No IP addresses found for NTP server");
    }
    let ntp = ntp_addrs[0];
    info!("NTP server: {ntp:?}");

    let mut counter = 0;
    loop {
        let addr: IpAddr = ntp.into();
        let timeout = get_time(SocketAddr::from((addr, 123)), &ntp_socket, context)
            .with_timeout(Duration::from_millis((max_wait_ms / 10) as u64))
            .await;

        match timeout {
            Ok(result) => {
                let time = result?;
                info!("Time: {time:?}");
                return DateTime::from_timestamp(time.seconds as i64, 0)
                    .context("Could not convert Sntp result");
            }
            Err(err) => {
                warn!("sntp timeout, retry: {err:?}");
                counter += 1;
                if counter > 10 {
                    bail!("Failed to get time from NTP server");
                }
                Timer::after(Duration::from_millis(100)).await;
            }
        }
    }
}

#[derive(Serialize, Debug, PartialEq)]
pub enum SntpMode {
    OFFLINE,
    SYNC { current: DateTime<Utc> },
}

#[derive(Serialize, Debug, PartialEq)]
pub enum NetworkMode {
    WIFI {
        sntp: SntpMode,
        mqtt: bool,
        ip_address: String,
    },
    OFFLINE,
}

#[embassy_executor::task(pool_size = 2)]
pub(crate) async fn net_task(mut runner: Runner<'static, Interface<'static>>) {
    runner.run().await;
}

#[embassy_executor::task]
pub(crate) async fn run_dhcp(stack: Stack<'static>, ip: Ipv4Addr) {
    let mut buf = [0u8; 1500];

    let mut gw_buf = [Ipv4Addr::UNSPECIFIED];

    let buffers = UdpBuffers::<3, 1024, 1024, 10>::new();
    let unbound_socket = Udp::new(stack, &buffers);
    let mut bound_socket = match unbound_socket
        .bind(SocketAddr::V4(SocketAddrV4::new(
            Ipv4Addr::UNSPECIFIED,
            DEFAULT_SERVER_PORT,
        )))
        .await
    {
        Ok(s) => s,
        Err(e) => {
            error!("dhcp task failed to bind socket: {:?}", e);
            return;
        }
    };

    loop {
        _ = io::server::run(
            &mut Server::<_, 64>::new_with_et(ip),
            &ServerOptions::new(ip, Some(&mut gw_buf)),
            &mut bound_socket,
            &mut buf,
        )
        .await
        .inspect_err(|e| warn!("DHCP server error: {e:?}"));
        Timer::after(Duration::from_millis(500)).await;
    }
}

pub async fn wifi_ap(
    ssid: String,
    interface_ap: Interface<'static>,
    controller: &Arc<Mutex<CriticalSectionRawMutex, esp_radio::wifi::WifiController<'static>>>,
    rng: &mut Rng,
    spawner: Spawner,
) -> FatResult<Stack<'static>> {
    let gw_ip_addr = Ipv4Addr::new(192, 168, 71, 1);

    let config = embassy_net::Config::ipv4_static(StaticConfigV4 {
        address: embassy_net::Ipv4Cidr::new(gw_ip_addr, 24),
        gateway: Some(gw_ip_addr),
        dns_servers: Default::default(),
    });

    let seed = (rng.random() as u64) << 32 | rng.random() as u64;

    println!("init secondary stack");
    let (stack, runner) = embassy_net::new(
        interface_ap,
        config,
        mk_static!(StackResources<4>, StackResources::<4>::new()),
        seed,
    );
    let stack = mk_static!(Stack, stack);

    let client_config =
        Config::AccessPoint(AccessPointConfig::default().with_ssid(ssid.clone()));
    controller.lock().await.set_config(&client_config)?;

    println!("start net task");
    spawner.spawn(net_task(runner)?);
    println!("run dhcp");
    spawner.spawn(run_dhcp(*stack, gw_ip_addr)?);

    loop {
        if stack.is_link_up() {
            break;
        }
        Timer::after(Duration::from_millis(500)).await;
    }
    while !stack.is_config_up() {
        Timer::after(Duration::from_millis(100)).await
    }
    println!("Connect to the AP `${ssid}` and point your browser to http://{gw_ip_addr}/");
    stack
        .config_v4()
        .inspect(|c| println!("ipv4 config: {c:?}"));

    Ok(*stack)
}

pub async fn wifi(
    network_config: &NetworkConfig,
    interface_sta: Interface<'static>,
    controller: &Arc<Mutex<CriticalSectionRawMutex, esp_radio::wifi::WifiController<'static>>>,
    rng: &mut Rng,
    spawner: Spawner,
) -> FatResult<Stack<'static>> {
    esp_radio::wifi_set_log_verbose();
    let ssid = match &network_config.ssid {
        Some(ssid) => {
            if ssid.is_empty() {
                bail!("Wifi ssid was empty")
            }
            ssid.as_str().to_string()
        }
        None => {
            bail!("Wifi ssid was empty")
        }
    };
    info!("attempting to connect wifi {ssid}");
    let password = match network_config.password {
        Some(ref password) => password.as_str().to_string(),
        None => "".to_string(),
    };
    let max_wait = network_config.max_wait;

    let config = embassy_net::Config::dhcpv4(DhcpConfig::default());

    let seed = (rng.random() as u64) << 32 | rng.random() as u64;

    let (stack, runner) = embassy_net::new(
        interface_sta,
        config,
        mk_static!(StackResources<8>, StackResources::<8>::new()),
        seed,
    );
    let stack = mk_static!(Stack, stack);

    let auth_method = if password.is_empty() {
        AuthenticationMethod::None
    } else {
        AuthenticationMethod::Wpa2Personal
    };
    let client_config = StationConfig::default()
        .with_ssid(ssid)
        .with_auth_method(auth_method)
        .with_scan_method(esp_radio::wifi::sta::ScanMethod::AllChannels)
        .with_listen_interval(10)
        .with_beacon_timeout(10)
        .with_failure_retry_cnt(3)
        .with_password(password);

    controller
        .lock()
        .await
        .set_config(&Config::Station(client_config))?;
    spawner.spawn(net_task(runner)?);
    controller
        .lock()
        .await
        .connect_async()
        .with_timeout(Duration::from_millis(max_wait as u64 * 1000))
        .await
        .context("Timeout waiting for wifi sta connected")??;

    let res = async {
        while !stack.is_link_up() {
            Timer::after(Duration::from_millis(500)).await;
        }
        Ok::<(), FatError>(())
    }
    .with_timeout(Duration::from_millis(max_wait as u64 * 1000))
    .await;

    if res.is_err() {
        bail!("Timeout waiting for wifi link up")
    }

    let res = async {
        while !stack.is_config_up() {
            Timer::after(Duration::from_millis(100)).await
        }
        Ok::<(), FatError>(())
    }
    .with_timeout(Duration::from_millis(max_wait as u64 * 1000))
    .await;

    if res.is_err() {
        bail!("Timeout waiting for wifi config up")
    }

    info!("Connected WIFI, dhcp: {:?}", stack.config_v4());
    Ok(*stack)
}

pub async fn try_connect_wifi_sntp_mqtt(
    board: &mut MutexGuard<'static, CriticalSectionRawMutex, HAL<'static>>,
    stack_store: &mut OptionLock<Stack<'static>>,
    spawner: Spawner,
) -> NetworkMode {
    let nw_conf = &board.board_hal.get_config().network.clone();
    let esp = board.board_hal.get_esp();
    let device = match esp.interface_sta.take() {
        Some(d) => d,
        None => {
            info!("Offline mode due to STA interface already taken");
            board.board_hal.general_fault(true).await;
            return NetworkMode::OFFLINE;
        }
    };
    match wifi(nw_conf, device, &esp.controller, &mut esp.rng, spawner).await {
        Ok(stack) => {
            stack_store.replace(stack);

            let sntp_mode: SntpMode = match sntp(1000 * 10, stack).await {
                Ok(new_time) => {
                    info!("Using time from sntp {}", new_time.to_rfc3339());
                    let _ = board
                        .board_hal
                        .get_rtc_module()
                        .set_rtc_time(&new_time)
                        .await;
                    SntpMode::SYNC { current: new_time }
                }
                Err(err) => {
                    warn!("sntp error: {err}");
                    board.board_hal.general_fault(true).await;
                    SntpMode::OFFLINE
                }
            };

            let mqtt_connected = if board.board_hal.get_config().network.mqtt_url.is_some() {
                let nw_config = board.board_hal.get_config().network.clone();
                let nw_config = mk_static!(NetworkConfig, nw_config);
                match mqtt::mqtt_init(nw_config, stack, spawner).await {
                    Ok(_) => {
                        info!("Mqtt connection ready");
                        true
                    }
                    Err(err) => {
                        warn!("Could not connect mqtt due to {err}");
                        false
                    }
                }
            } else {
                false
            };

            let ip = match stack.config_v4() {
                Some(config) => config.address.address().to_string(),
                None => match stack.config_v6() {
                    Some(config) => config.address.address().to_string(),
                    None => String::from("No IP"),
                },
            };
            NetworkMode::WIFI {
                sntp: sntp_mode,
                mqtt: mqtt_connected,
                ip_address: ip,
            }
        }
        Err(err) => {
            info!("Offline mode due to {err}");
            board.board_hal.general_fault(true).await;
            NetworkMode::OFFLINE
        }
    }
}