eculocate/ecu-esp32/src/wifi.rs

//! This gets an ip address via DHCP and advertises it with mdns

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

use hashbrown::HashMap;

use ed25519_dalek::{VerifyingKey, Verifier, Signature};

use embassy_executor::Spawner;
use embassy_net::{
    Runner,
    StackResources
};
use embassy_futures::join::join3;
use embassy_sync::mutex::Mutex;
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;

use embassy_time::{Duration, Timer, Instant};
use esp_alloc as _;
use esp_backtrace as _;
use esp_hal::{
    rng::Rng,
};
use esp_println::println;
use esp_radio::wifi::{
    ScanConfig,
    WifiController,
    WifiDevice,
    WifiEvent,
    WifiStaState,
    WifiError,
    AccessPointInfo
};

use embassy_net::Stack;
use embedded_io_async::Read;

use log::{warn, info, error};

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
    }};
}

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 struct SessionStore {
    map : Mutex<CriticalSectionRawMutex, HashMap<[u8; 32], Instant>>,
    public_key: VerifyingKey
}

impl SessionStore {
    pub fn new() -> Self {
	let m = HashMap::new();
	let verifying_key_bytes = include_bytes!("../../pubkey.bin");
	let public_key: VerifyingKey = VerifyingKey::from_bytes(verifying_key_bytes).expect("verify");

	Self { map : Mutex::new(m), public_key }
    }

    pub async fn register(&self, key: &[u8], signature: &[u8]) -> bool {
	info!("pretending to register {:?}", key);

	if self.public_key.verify(key, &Signature::try_from(signature).expect("64 bytes")).is_ok() {
	    let mut m = self.map.lock().await;
	    m.insert(key.try_into().unwrap(), Instant::now() + Duration::from_secs(86400));
	    true
	} else {
	    warn!("wrong key");
	    false
	}
    }

    pub async fn is_registered(&self, key: &[u8]) -> bool {
	let m = self.map.lock().await;
	match m.get(key) {
	    Some(expiry) => *expiry > Instant::now(),
	    None => false
	}
    }
}

pub async fn scan_networks(controller: &mut WifiController<'_>) -> Vec<AccessPointInfo> {
    println!("Scan");
    let scan_config = ScanConfig::default().with_max(10);
    let result = controller
        .scan_with_config_async(scan_config)
        .await
        .expect("scan failed");
    result
}

pub async fn start_wifi(spawner: &Spawner,
                        mut controller: WifiController<'static>,
                        wifi_interface: WifiDevice<'static>,
			mut flasher: crate::ota::Flasher<'static>,
                        ecu: &mut crate::ecu::Ecu<'_>)  {

    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,
    );

    if !connect(&mut controller).await.is_ok() {
	flasher.set_var("ssid", None);
	flasher.set_var("secret", None);
	info!("[gatt] deep breath then reboot");
	Timer::after(Duration::from_millis(2000)).await;
	esp_hal::system::software_reset();
    }

    spawner.spawn(maintain_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;

    let session_store = SessionStore::new();

    join3(
        crate::streamer::loop_udp_thing(stack, ecu, &session_store),
        run_tcp_handler(stack, &session_store, flasher),
        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;
}

async fn connect(mut controller: &mut WifiController<'static>) -> Result<(), WifiError> {
    let mut attempts = 0;
    let mut retval = Ok(());
    println!("start connecting");
    while attempts < 5 {
        println!("About to connect...");

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

#[embassy_executor::task]
async fn maintain_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
            }
            _ => {}
        }
	let _ = connect(&mut controller);
    }
}

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

async fn run_tcp_handler(stack : Stack<'_>,
			 sessions : &SessionStore,
			 mut flasher : crate::ota::Flasher<'_>) {
    let mut rx_buffer = [0; 4096];
    let mut tx_buffer = [0; 4096];
    let mut buf = [0u8; 96];

    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());

        if let Ok(_) = socket.read_exact(&mut buf[0..4]).await {
            info!("got a thing");
            match &buf[0..4] {
                b"ROM0" => {
                    info!("OTA request");
                    if let Err(e) = flasher.write(&mut socket).await {
			error!("ota flash failed: {:?}", e);
		    }
                },
                b"KEY0" => {
                    info!("sessionkey");
                    if let Ok(_) = socket.read_exact(&mut buf).await {
                        let (key, sig) = buf.split_at(32);
                        if sessions.register(key, sig).await {
			    info!("registered? {}", sessions.is_registered(key).await);
			    info!("write {:?}",  socket.write(b"BYE\n").await);
			    let _ = socket.flush().await;
			} else {
			    warn!("couldn't register key")
			}
			socket.close();
                    }
                },
                _ => {
                    warn!("unrecognised command");
                }

            }
        }
    }
}