eculocate/ecu-esp32/src/wifi.rs

//! Set SSID and PASSWORD env variable before running this example.
//!
//! This gets an ip address via DHCP and advertises it with mdns

// use alloc::vec;
use alloc::vec::Vec;
use alloc::string::String;
use alloc::format;

use embassy_executor::Spawner;
use embassy_net::{
    Runner,
    StackResources
};
use embassy_futures::join::join;

use embassy_time::{Duration, Timer};
use esp_alloc as _;
use esp_backtrace as _;
use esp_hal::{
    rng::Rng,
};
use esp_println::println;
use esp_radio::{
    Controller,
    wifi::{
        ClientConfig,
        ModeConfig,
        ScanConfig,
        WifiController,
        WifiDevice,
        WifiEvent,
        WifiStaState,
    },
};
use esp_hal_ota::Ota;
use esp_storage::FlashStorage;

use embassy_net::Stack;

use log::{warn,info};

use crate::mdns;

// When you are okay with using a nightly compiler it's better to use https://docs.rs/static_cell/2.1.0/static_cell/macro.make_static.html
macro_rules! mk_static {
    ($t:ty,$val:expr) => {{
        static STATIC_CELL: static_cell::StaticCell<$t> = static_cell::StaticCell::new();
        #[deny(unused_attributes)]
        let x = STATIC_CELL.uninit().write(($val));
        x
    }};
}

const SSID: &str = env!("SSID");
const PASSWORD: &str = env!("PASSWORD");


fn dns_records(my_ip_address: [u8;4], serial: Option<usize>) -> Vec<mdns::Resource> {
    use crate::mdns::*;
    use crate::mdns::RData as R;

    let (hostname, ptr_name) =
	match serial {
	    None => (
		String::from("eculocate.local"),
		String::from("Bike._keihin._udp.local")
	    ),
	    Some(number) => (
		format!("eculocate{number}.local"),
		format!("Bike ({number})._keihin._udp.local")
	    )
	};


    let records = [
	Resource::new(&hostname, R::A(my_ip_address)),
	Resource::new(
	    &ptr_name,
	    R::Srv {
		target: String::from(hostname), port: 5000,
		weight: 0, priority: 0
	    }
	),
	Resource::new(&ptr_name, R::Txt(["txtvers=1".into()].into())),
	Resource::new("_keihin._udp.local", R::Ptr(ptr_name)),
	Resource::new(DISCOVERY_NAME, R::Ptr("_keihin._udp.local".into()))
    ];
    Vec::from(records)
}

pub async fn start_wifi(spawner: &Spawner,
		        wifi_peripheral:  esp_hal::peripherals::WIFI<'static>,
			flash_peripheral: esp_hal::peripherals::FLASH<'static>,
                        ecu: &mut crate::ecu::Ecu<'_>)  {


    let esp_radio_ctrl = &*mk_static!(Controller<'static>, esp_radio::init().unwrap());

    let (controller, interfaces) =
        esp_radio::wifi::new(esp_radio_ctrl, wifi_peripheral, Default::default()).unwrap();

    let wifi_interface = interfaces.sta;

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

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

    // Init network stack
    let (stack, runner) = embassy_net::new(
        wifi_interface,
        config,
        mk_static!(StackResources<5>, StackResources::<5>::new()),
        seed,
    );

    spawner.spawn(connection(controller)).ok();
    spawner.spawn(net_task(runner)).ok();

    let my_address;

    loop {
        if stack.is_link_up() {
            break;
        }
        Timer::after(Duration::from_millis(500)).await;
    }

    println!("Waiting to get IP address...");
    loop {
        if let Some(config) = stack.config_v4() {
            println!("Got IP address: {}", config.address);
	    my_address = config.address.address().octets();
            break;
        }
        Timer::after(Duration::from_millis(500)).await;
    }

    let mut serial = None;

    join(
	join(
            crate::streamer::loop_udp_thing(stack, ecu),
	    run_tcp_handler(stack, flash_peripheral),
	),
        async {
            loop {
	        let mdns_thread = mdns::responder(
	            stack,
	            dns_records(my_address, serial)
	        );
	        mdns_thread.await;
	        serial = match serial { None => Some(2), Some(n) => Some(n+1) };
	        info!("probe failed, restarting with incremented serial {:?}", serial);
                Timer::after(Duration::from_millis(5000)).await;
            };
        }
    ).await;
}


#[embassy_executor::task]
async fn connection(mut controller: WifiController<'static>) {
    println!("start connection task");
    println!("Device capabilities: {:?}", controller.capabilities());
    loop {
        match esp_radio::wifi::sta_state() {
            WifiStaState::Connected => {
                // wait until we're no longer connected
                controller
                    .wait_for_event(WifiEvent::StaDisconnected)
                    .await;
                Timer::after(Duration::from_millis(5000)).await
            }
            _ => {}
        }
        if !matches!(controller.is_started(), Ok(true)) {
            let station_config = ModeConfig::Client(
                ClientConfig::default()
                    .with_ssid(SSID.into())
                    .with_password(PASSWORD.into()),
            );
            controller.set_config(&station_config).unwrap();
            println!("Starting wifi");
            controller.start_async().await.unwrap();
            println!("Wifi started!");

            println!("Scan");
            let scan_config = ScanConfig::default().with_max(10);
            let result = controller
                .scan_with_config_async(scan_config)
                .await
                .unwrap();
            for ap in result {
                println!("{:?}", ap);
            }
        }
        println!("About to connect...");

        match controller.connect_async().await {
            Ok(_) => println!("Wifi connected!"),
            Err(e) => {
                println!("Failed to connect to wifi: {e:?}");
                Timer::after(Duration::from_millis(5000)).await
            }
        }
    }
}

#[embassy_executor::task]
async fn net_task(mut runner: Runner<'static, WifiDevice<'static>>) {
    runner.run().await
}

async fn read_all(socket: &mut embassy_net::tcp::TcpSocket<'_>,
		  buf: &mut [u8],
		  len: usize) -> usize {
    let mut n = 0;
    while n < len {
	match socket.read(&mut buf[n..]).await {
	    Ok(0) => {
		warn!("read EOF");
		break;
	    },
	    Ok(count) => {
		n = n + count;
	    }
	    Err(e) => {
		warn!("read err {:?}", e);
		return 0;	// ew
	    }
	}
    }
    n
}

async fn write_ota(ota : &mut Ota<FlashStorage<'_>>,
		   sock: &mut embassy_net::tcp::TcpSocket<'_>,
		   len: u32) {
    let mut buf = [0u8; 4096];

    ota.ota_begin(len, 0);

    let mut len = len;
    while len > 0 {
	match sock.read(&mut buf).await {
	    Ok(chunksize) => {
		info!("writing {} bytes, {} remaining", chunksize, len);
		let res = ota.ota_write_chunk(&buf[0..chunksize]);
		if res == Ok(true) {
		    info!("eof on write, remaining bytes = {} ", len);

		    if ota.ota_flush(false, false).is_ok() { // ignore crc, no rollback
			Timer::after_millis(1000).await;
			esp_hal::system::software_reset();
		    }
		}
		// let progress = (ota.get_ota_progress() * 100.0) as u8;
		// log::info!("progress: {}%", progress);
		len = len - (chunksize as u32);
	    },
	    Err(x) => {
		warn!("error reading ota image: {:?}", x);
		break;
	    }
	}
    }
    if len == 0 {
	info!("wrote all bytes");
    } else {
	warn!("stream ended early, was expecting another {} bytes", len);
    }
}

fn register_session_key(key : &[u8]) {
    info!("got session key{:?}", key);
}


async fn run_tcp_handler(stack : Stack<'_>, flash_peripheral: esp_hal::peripherals::FLASH<'static>) {
    let mut rx_buffer = [0; 4096];
    let mut tx_buffer = [0; 4096];
    let mut buf = [0u8; 68];
    let flash = FlashStorage::new(flash_peripheral);
    let mut ota = Ota::new(flash).unwrap();

    loop {
	let mut socket = embassy_net::tcp::TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
	socket.set_timeout(Some(Duration::from_secs(10)));

	if let Err(e) = socket.accept(5000).await {
	    warn!("error accepting connection: {:?}", e);
	    continue;
	}

        info!("Received connection from {:?}", socket.remote_endpoint());

	let n = read_all(&mut socket, &mut buf, 68).await;

	if n == 68 {
	    info!("got a thing");
	    match &buf[0..4] {
		b"ROM0" => {
		    info!("OTA request");
		    let mut len_bytes = [0u8; 4];
		    let _n = read_all(&mut socket, &mut len_bytes[..], 4).await;
		    let len = u32::from_be_bytes(len_bytes);
		    info!("reading {} bytes", len);
		    write_ota(&mut ota, &mut socket, len).await;
		    info!("write {:?}",  socket.write(b"BYE\n").await);
		    socket.flush().await;
		    socket.close();
		},
		b"KEY0" => {
		    info!("sessionkey");
		    let _sig = &buf[4..68];
		    let mut key = [0u8; 32]; // 256 bits
		    if read_all(&mut socket, &mut key, 32).await == 32 {
			register_session_key(&key);
		    }
		},
		_ => {
		    warn!("unrecognised command");
		}

	    }
	}
    }
}