itpdp/esp32/src/net.rs

use core::net::SocketAddrV4;
use core::str::FromStr;

use alloc::format;
use embassy_futures::join::join;
use embassy_net::tcp::TcpSocket;
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
use embassy_sync::signal::Signal;
use embassy_time::{Duration, Timer};
use esp_hal::peripherals::WIFI;
use esp_hal::rng::Rng;
use esp_println::println;
use esp_radio::wifi::sta::StationConfig;
use esp_radio::wifi::{Config, ControllerConfig, scan::ScanConfig};

use crate::screen::overwrite_lcd;
use crate::{WIFI_NETWORK, WIFI_PASSWORD, main_loop};
use crate::{buffer::wait_for_config, tcp_read_loop, tcp_write_loop};
use crate::{reconnecting_state, reset_state};

pub struct NetworkConfig<'a> {
    pub wifi: WIFI<'a>,
    pub wifi_network: &'a str,
    pub wifi_password: &'a str,
    pub target_ip: &'a str,
    pub target_port: u16,
}

#[embassy_executor::task]
async fn net_task(
    mut runner: embassy_net::Runner<'static, esp_radio::wifi::Interface<'static>>,
) -> ! {
    runner.run().await
}

#[embassy_executor::task]
pub async fn network_setup_task(
    spawner: embassy_executor::Spawner,
    config: NetworkConfig<'static>,
) {
    let wifi_config = Config::Station(
        StationConfig::default()
            .with_ssid(config.wifi_network)
            .with_password(config.wifi_password.into()),
    );

    println!("Starting wifi");
    let (mut controller, interfaces) = esp_radio::wifi::new(
        config.wifi,
        ControllerConfig::default().with_initial_config(wifi_config),
    )
    .unwrap();
    println!("Wifi configured and started!");

    controller
        .set_power_saving(esp_radio::wifi::PowerSaveMode::None)
        .unwrap();

    let wifi_interface = interfaces.station;
    let net_config = embassy_net::Config::dhcpv4(Default::default());
    let rng = Rng::new();
    let seed = (rng.random() as u64) << 32 | rng.random() as u64;

    let (stack, runner) = embassy_net::new(
        wifi_interface,
        net_config,
        crate::screen::stack_resources(),
        seed,
    );
    spawner.spawn(net_task(runner).expect("spawn net task"));
    spawner.spawn(main_loop().expect("spawn main loop"));

    loop {
        loop {
            overwrite_lcd("Connecting", "").await;
            println!("Scan");
            let scan_config = ScanConfig::default().with_max(10);
            let result = controller.scan_async(&scan_config).await.unwrap();
            for ap in result {
                println!(
                    "- {} (CH {} Strength {})",
                    ap.ssid.as_str(),
                    ap.channel,
                    ap.signal_strength
                );
            }

            println!("Scan ended");

            if let Err(e) = controller.set_config(&Config::Station(
                StationConfig::default()
                    // .with_auth_method(esp_radio::wifi::AuthenticationMethod::Wpa2Wpa3Personal)
                    .with_ssid(WIFI_NETWORK)
                    .with_password(WIFI_PASSWORD.into()),
            )) {
                println!("Connection error: {:?}", e);
                overwrite_lcd("Connection error", &format!("{}", e)).await;
                Timer::after_millis(1000).await;
                continue;
            }

            if let Err(e) = controller.connect_async().await {
                println!("Connect error: {:?}", e);
                overwrite_lcd("Connection error", &format!("{}", e)).await;
                Timer::after_millis(1000).await;
                continue;
            }

            println!("Configured");

            stack.wait_config_up().await;

            break;
        }

        println!("Waiting for DHCP");

        let cfg = wait_for_config(stack).await;
        let local_addr = cfg.address.address();
        println!("IP address: {:?}", local_addr);

        loop {
            let host_addr = embassy_net::Ipv4Address::from_str(config.target_ip).unwrap();
            let mut rx_buffer = [0; 4096];
            let mut tx_buffer = [0; 4096];
            let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
            socket.set_timeout(Some(Duration::from_secs(10)));
            socket.set_keep_alive(Some(Duration::from_secs(1)));
            if let Err(e) = socket
                .connect(SocketAddrV4::new(host_addr, config.target_port))
                .await
            {
                println!("tcp connect error: {:?}", e);
                reconnecting_state().await;
                overwrite_lcd("TCP error", &format!("{}", e)).await;
                Timer::after_millis(1000).await;
                continue;
            }
            reset_state().await;
            overwrite_lcd("Connected", "").await;

            let cancel = Signal::<CriticalSectionRawMutex, ()>::new();
            let (read, write) = socket.split();
            let (read_result, write_result) =
                join(tcp_read_loop(read, &cancel), tcp_write_loop(write, &cancel)).await;
            let disconnect = read_result.err().or(write_result.err());

            if let Some(reason) = disconnect {
                println!("tcp disconnected: {:?}", reason);
            } else {
                println!("tcp loops ended");
            }

            if !stack.is_config_up() {
                println!("wifi down, reconnecting wifi");
                reconnecting_state().await;
                break;
            }

            reconnecting_state().await;
            Timer::after_millis(500).await;
        }
    }
}