Fix root dir

Author: fangpenlin

src/api/processor.rs

@@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize};
 use std::fs::{read_dir, FileType};
 use std::io::{Read, Seek};
 use std::mem::MaybeUninit;
+use std::path::Path;
 use std::time::SystemTime;
 use time::serde::timestamp::milliseconds;
 use time::OffsetDateTime;
@@ -79,6 +80,7 @@ pub type DeviceInfoProducer = Box<dyn Fn() -> anyhow::Result<DeviceInfo>>;
 
 pub struct Processor {
     pub device_info_producer: DeviceInfoProducer,
+    pub root_dir: String,
 }
 
 impl Processor {
@@ -91,10 +93,11 @@ impl Processor {
     }
 
     fn list_files(&self, path: &str) -> anyhow::Result<Response> {
+        let dir_path = Path::new(&self.root_dir).join(path);
         // Ideally we should find a way to learn the size of all files, but we need to
         // iterate over all files anyway... so.. maybe not? :/
         let mut files: Vec<File> = vec![];
-        for entry in read_dir(path)? {
+        for entry in read_dir(dir_path)? {
             let entry = entry?;
             let path = entry.path().into_os_string().into_string().map_err(|e| {
                 anyhow!(
@@ -183,9 +186,11 @@ impl Processor {
 pub async fn process_events(
     mut client: WebSocketSession<'_>,
     device_info_producer: DeviceInfoProducer,
+    root_dir: &str,
 ) {
     let mut processor: Option<Box<Processor>> = Some(Box::new(Processor {
         device_info_producer,
+        root_dir: root_dir.to_string(),
     }));
     client.connect();
 
@@ -198,9 +203,8 @@ pub async fn process_events(
                 new_state: ConnectionState::Connected,
                 ..
             } => {
-                processor = Some(Box::new(Processor {
-                    device_info_producer: processor.map(|p| p.device_info_producer).unwrap(),
-                }));
+                let last_processor = processor.unwrap();
+                processor = Some(Box::new(Processor { ..*last_processor }));
             }
             SessionEvent::ReceiveText { text } => {
                 let request: serde_json::Result<CommandRequest> = serde_json::from_str(&text);

src/main.rs

@@ -24,10 +24,15 @@ const VERSION: &str = env!("CARGO_PKG_VERSION");
 const SSID: &str = env!("WIFI_SSID");
 const PASSWORD: &str = env!("WIFI_PASS");
 const API_ENDPOINT: &str = env!("API_ENDPOINT");
+const PARTITION_LABEL: Option<&str> = option_env!("PARTITION_LABEL");
+const MOUNT_PATH: Option<&str> = option_env!("MOUNT_PATH");
 
 async fn run_async(spawner: LocalSpawner) -> Result<(), anyhow::Error> {
     log::info!("Start {} - {}", PKG_NAME, VERSION,);
 
+    let partition_label = PARTITION_LABEL.unwrap_or("storage");
+    let mount_path = MOUNT_PATH.unwrap_or("/disk");
+
     let peripherals = Peripherals::take()?;
     let mut wifi = WifiSession::new(SSID, PASSWORD, AuthMethod::WPA2Personal, peripherals.modem)?;
     wifi.connect().await?;
@@ -37,7 +42,7 @@ async fn run_async(spawner: LocalSpawner) -> Result<(), anyhow::Error> {
     let _sntp = EspSntp::new_default()?;
     log::info!("SNTP initialized");
 
-    let mut msc_device = MSCDevice::new("storage", "/disk");
+    let mut msc_device = MSCDevice::new(partition_label, mount_path);
     msc_device.install()?;
 
     let mut button = PinDriver::input(peripherals.pins.gpio14)?;
@@ -60,7 +65,7 @@ async fn run_async(spawner: LocalSpawner) -> Result<(), anyhow::Error> {
     button.wait_for_low().await?;
     log::info!("Button pressed!");
 
-    spawner.spawn_local(process_events(client, device_info_producer))?;
+    spawner.spawn_local(process_events(client, device_info_producer, mount_path))?;
 
     loop {
         // Asynchronously wait for GPIO events, allowing other tasks

src/usb/msc_device.rs

@@ -12,8 +12,8 @@ use std::ffi::CString;
 #[derive(Default)]
 pub struct MSCDevice {
     pub partition_label: String,
-    pub base_path: String,
-    base_path_c_str: CString,
+    pub mount_path: String,
+    mount_path_c_str: CString,
     wl_partition: Option<EspWlPartition<EspPartition>>,
 }
 
@@ -36,7 +36,7 @@ unsafe extern "C" fn storage_mount_changed_cb(event: *mut tinyusb_msc_event_t) {
 impl MSCDevice {
     pub fn new(partition_label: &str, base_path: &str) -> Self {
         Self {
-            base_path: base_path.to_string(),
+            mount_path: base_path.to_string(),
             partition_label: partition_label.to_string(),
             ..Default::default()
         }
@@ -58,7 +58,7 @@ impl MSCDevice {
             Some(EspWlPartition::new(partition.unwrap()).with_context(|| {
                 format!(
                     "Failed to mount partition {} at {}",
-                    self.partition_label, self.base_path
+                    self.partition_label, self.mount_path
                 )
             })?);
 
@@ -78,17 +78,17 @@ impl MSCDevice {
         esp!(unsafe { tinyusb_msc_storage_init_spiflash(&config_spi) })
             .with_context(|| "Failed to initialize spiflash")?;
 
-        let base_path_c_str = CString::new(self.base_path.as_bytes()).unwrap();
+        let base_path_c_str = CString::new(self.mount_path.as_bytes()).unwrap();
         esp!(unsafe { tinyusb_msc_storage_mount(base_path_c_str.as_ptr()) })
-            .with_context(|| format!("Failed to mount storage at {}", self.base_path))?;
+            .with_context(|| format!("Failed to mount storage at {}", self.mount_path))?;
 
         let tusb_cfg = tinyusb_config_t::default();
         esp!(unsafe { tinyusb_driver_install(&tusb_cfg) })
             .with_context(|| "Failed to install TinyUSB driver")?;
 
         log::info!("TinyUSB driver installed.");
 
-        self.base_path_c_str = base_path_c_str;
+        self.mount_path_c_str = base_path_c_str;
         self.wl_partition = wl_partition;
         Ok(())
     }