k-frame/crates/client-esp32/src/provisioning/portal.rs

use std::net::UdpSocket;
use std::thread;
use esp_idf_hal::io::Write;
use esp_idf_svc::http::server::{Configuration as HttpConfig, EspHttpServer};
use esp_idf_svc::nvs::{EspNvsPartition, NvsDefault};
use client_domain::{BoundingBox, Color, DisplayPort, FontSize};
use log::{info, error};

use super::{DeviceConfig, save_config};
use super::html::SETUP_HTML;

const AP_IP: [u8; 4] = [192, 168, 4, 1];

pub fn run_captive_portal<D: DisplayPort>(
    nvs: EspNvsPartition<NvsDefault>,
    display: &mut D,
) {
    draw_setup_screen(display);
    spawn_dns_responder();

    let nvs_clone = nvs.clone();
    let mut server = EspHttpServer::new(&HttpConfig {
        http_port: 80,
        ..Default::default()
    }).expect("HTTP server start failed");

    server
        .fn_handler::<esp_idf_svc::io::EspIOError, _>("/", esp_idf_svc::http::Method::Get, |req| {
            req.into_ok_response()?
                .write_all(SETUP_HTML.as_bytes())?;
            Ok(())
        })
        .expect("GET / handler failed");

    server
        .fn_handler::<esp_idf_svc::io::EspIOError, _>("/save", esp_idf_svc::http::Method::Post, move |mut req| {
            let mut body = vec![0u8; 512];
            let len = req.read(&mut body).unwrap_or(0);
            let body_str = std::str::from_utf8(&body[..len]).unwrap_or("");

            let config = parse_form(body_str);
            info!("Portal received config: ssid={}", config.wifi_ssid);

            if let Err(e) = save_config(nvs_clone.clone(), &config) {
                error!("Save failed: {e}");
                req.into_ok_response()?.write_all(b"Error saving config")?;
                return Ok(());
            }

            req.into_ok_response()?.write_all(b"OK")?;

            thread::spawn(|| {
                thread::sleep(std::time::Duration::from_secs(1));
                unsafe { esp_idf_svc::sys::esp_restart(); }
            });
            Ok(())
        })
        .expect("POST /save handler failed");

    server
        .fn_handler::<esp_idf_svc::io::EspIOError, _>("/generate_204", esp_idf_svc::http::Method::Get, |req| {
            redirect(req)
        })
        .expect("generate_204 handler failed");

    server
        .fn_handler::<esp_idf_svc::io::EspIOError, _>("/hotspot-detect.html", esp_idf_svc::http::Method::Get, |req| {
            redirect(req)
        })
        .expect("hotspot-detect handler failed");

    server
        .fn_handler::<esp_idf_svc::io::EspIOError, _>("/canonical.html", esp_idf_svc::http::Method::Get, |req| {
            redirect(req)
        })
        .expect("canonical handler failed");

    info!("Captive portal running on http://192.168.4.1");

    loop {
        thread::sleep(std::time::Duration::from_secs(60));
    }
}

fn redirect(req: esp_idf_svc::http::server::Request<&mut esp_idf_svc::http::server::EspHttpConnection>) -> Result<(), esp_idf_svc::io::EspIOError> {
    let mut resp = req.into_response(302, None, &[("Location", "http://192.168.4.1/")])?;
    resp.write_all(b"")?;
    Ok(())
}

fn spawn_dns_responder() {
    thread::Builder::new()
        .stack_size(4096)
        .name("dns".into())
        .spawn(dns_responder)
        .expect("DNS thread spawn failed");
}

fn dns_responder() {
    let socket = match UdpSocket::bind("0.0.0.0:53") {
        Ok(s) => s,
        Err(e) => {
            error!("DNS bind failed: {e}");
            return;
        }
    };

    info!("DNS responder started on :53");
    let mut buf = [0u8; 512];

    loop {
        let (len, src) = match socket.recv_from(&mut buf) {
            Ok(r) => r,
            Err(_) => continue,
        };
        if len < 12 {
            continue;
        }

        let mut resp = Vec::with_capacity(len + 16);

        // Header: copy ID, set response flags
        resp.extend_from_slice(&buf[0..2]);
        resp.extend_from_slice(&[0x81, 0x80]);
        resp.extend_from_slice(&buf[4..6]);
        resp.extend_from_slice(&buf[4..6]);
        resp.extend_from_slice(&[0, 0, 0, 0]);

        // Copy question section
        resp.extend_from_slice(&buf[12..len]);

        // Answer: A record pointing to AP IP
        resp.extend_from_slice(&[0xC0, 0x0C]); // Name pointer
        resp.extend_from_slice(&[0, 1]); // Type A
        resp.extend_from_slice(&[0, 1]); // Class IN
        resp.extend_from_slice(&[0, 0, 0, 60]); // TTL 60s
        resp.extend_from_slice(&[0, 4]); // Data length
        resp.extend_from_slice(&AP_IP);

        let _ = socket.send_to(&resp, src);
    }
}

fn parse_form(body: &str) -> DeviceConfig {
    let mut ssid = String::new();
    let mut pass = String::new();
    let mut server = String::new();

    for pair in body.split('&') {
        let mut parts = pair.splitn(2, '=');
        let key = parts.next().unwrap_or("");
        let val = url_decode(parts.next().unwrap_or(""));
        match key {
            "ssid" => ssid = val,
            "pass" => pass = val,
            "server" => server = val,
            _ => {}
        }
    }

    DeviceConfig { wifi_ssid: ssid, wifi_pass: pass, server_addr: server }
}

fn url_decode(s: &str) -> String {
    let mut result = String::with_capacity(s.len());
    let mut chars = s.bytes();
    while let Some(b) = chars.next() {
        match b {
            b'+' => result.push(' '),
            b'%' => {
                let hi = chars.next().and_then(hex_val);
                let lo = chars.next().and_then(hex_val);
                if let (Some(h), Some(l)) = (hi, lo) {
                    result.push((h << 4 | l) as char);
                }
            }
            _ => result.push(b as char),
        }
    }
    result
}

fn hex_val(b: u8) -> Option<u8> {
    match b {
        b'0'..=b'9' => Some(b - b'0'),
        b'a'..=b'f' => Some(b - b'a' + 10),
        b'A'..=b'F' => Some(b - b'A' + 10),
        _ => None,
    }
}

const FULL_SCREEN: BoundingBox = BoundingBox { x: 0, y: 0, width: 320, height: 240 };
const WHITE: Color = Color(0xFF, 0xFF, 0xFF);
const DARK_BG: Color = Color(0x08, 0x10, 0x18);

fn draw_setup_screen<D: DisplayPort>(display: &mut D) {
    let _ = display.fill_rect(FULL_SCREEN, DARK_BG);
    let _ = display.draw_text_span("K-Frame Setup", 100, 50, WHITE, FontSize::Small);
    let _ = display.draw_text_span("Connect to WiFi:", 60, 90, WHITE, FontSize::Small);
    let _ = display.draw_text_span("KFrame-Setup", 90, 110, WHITE, FontSize::Small);
    let _ = display.draw_text_span("Then open browser", 60, 150, WHITE, FontSize::Small);
    let _ = display.draw_text_span("to configure", 80, 170, WHITE, FontSize::Small);
}