flashthing/
aml.rs

1use rusb::{Context, DeviceHandle, Direction, UsbContext};
2use std::{io::Read, sync::Arc, thread::sleep, time::Duration};
3
4use crate::{
5  flash::FlashProgress, partitions::PartitionInfo, Callback, Error, Event, Result, ADDR_BL2, ADDR_TMP,
6  AMLC_AMLS_BLOCK_LENGTH, AMLC_MAX_BLOCK_LENGTH, AMLC_MAX_TRANSFER_LENGTH, BL2_BIN, BOOTLOADER_BIN, FLAG_KEEP_POWER_ON,
7  PART_SECTOR_SIZE, PRODUCT_ID, REQ_BULKCMD, REQ_GET_AMLC, REQ_IDENTIFY_HOST, REQ_READ_MEM, REQ_RUN_IN_ADDR,
8  REQ_WRITE_AMLC, REQ_WRITE_MEM, REQ_WR_LARGE_MEM, TRANSFER_BLOCK_SIZE, TRANSFER_SIZE_THRESHOLD, UNBRICK_BIN_ZIP,
9  VENDOR_ID,
10};
11
12const COMMAND_TIMEOUT: Duration = Duration::from_secs(10);
13
14#[derive(Debug)]
15struct AmlInner {
16  handle: DeviceHandle<Context>,
17  interface_number: u8,
18  endpoint_in: u8,
19  endpoint_out: u8,
20}
21
22/// The main interface for interacting with Amlogic-based hardware
23///
24/// This provides low-level access to the Amlogic SoC on the Superbird device,
25/// allowing for memory operations, partition management, and firmware flashing.
26#[derive(Clone)]
27pub struct AmlogicSoC {
28  inner: Arc<AmlInner>,
29}
30
31impl AmlogicSoC {
32  /// Initialize a connection to an Amlogic SoC device
33  ///
34  /// This will search for a connected device, put it in the correct mode if necessary,
35  /// and establish a connection for flashing operations.
36  ///
37  /// # Parameters
38  /// - `callback`: Optional callback function to receive status updates
39  ///
40  /// # Returns
41  /// - `Result<Self>`: A connected AmlogicSoC instance or an error
42  pub fn init(callback: Option<Callback>) -> Result<Self> {
43    if let Some(callback) = &callback {
44      callback(Event::FindingDevice);
45    };
46
47    let mode = find_device();
48    if let Some(callback) = &callback {
49      callback(Event::DeviceMode(mode));
50    };
51
52    match mode {
53      DeviceMode::Usb => {
54        tracing::info!("device booted in usb mode - moving to usb burn mode");
55        let device = Self::connect(callback.clone())?;
56        if let Some(callback) = &callback {
57          callback(Event::Bl2Boot);
58        };
59
60        device.bl2_boot(None, None)?;
61        drop(device);
62
63        if let Some(callback) = &callback {
64          callback(Event::Resetting);
65        };
66
67        tracing::debug!("device successfully moved to usb burn mode, sleeping then grabbing new handle");
68        sleep(Duration::from_millis(5000));
69      }
70      DeviceMode::UsbBurn => tracing::info!("device found!"),
71      DeviceMode::Normal => {
72        tracing::error!(
73          "device is booted in normal mode. make sure to power on the car thing while holding buttons 1 & 4"
74        );
75        return Err(Error::WrongMode);
76      }
77      DeviceMode::NotFound => {
78        tracing::error!("device not found!! make sure to power on the car thing while holding buttons 1 & 4");
79        return Err(Error::NotFound);
80      }
81    };
82
83    let mut attempts = 0;
84    while attempts < 3 {
85      match Self::connect(callback.clone()) {
86        Ok(dev) => return Ok(dev),
87        Err(e) => {
88          tracing::debug!("failed to connect to device: {}. Attempt {}/3", e, attempts + 1);
89          attempts += 1;
90          sleep(Duration::from_secs(1));
91        }
92      }
93    }
94
95    Self::connect(callback)
96  }
97
98  fn connect(callback: Option<Callback>) -> Result<Self> {
99    tracing::debug!("connecting to Amlogic device");
100    if let Some(callback) = &callback {
101      callback(Event::Connecting);
102    };
103
104    let context = Context::new()?;
105    let handle = {
106      let device = context
107        .devices()?
108        .iter()
109        .find(|device| {
110          if let Ok(desc) = device.device_descriptor() {
111            desc.vendor_id() == VENDOR_ID && desc.product_id() == PRODUCT_ID
112          } else {
113            false
114          }
115        })
116        .ok_or_else(|| Error::InvalidOperation("Device not found".into()))?;
117      device.open()?
118    };
119
120    handle.set_active_configuration(1)?;
121    let interface_number: u8 = 0;
122    handle.claim_interface(interface_number)?;
123
124    let device = handle.device();
125    let config_desc = device.active_config_descriptor()?;
126    let interface = config_desc
127      .interfaces()
128      .find(|i| i.number() == interface_number)
129      .ok_or_else(|| Error::InvalidOperation("Interface not found".into()))?;
130    let descriptor = interface
131      .descriptors()
132      .next()
133      .ok_or_else(|| Error::InvalidOperation("No alt setting".into()))?;
134    let mut endpoint_in = None;
135    let mut endpoint_out = None;
136    for ep in descriptor.endpoint_descriptors() {
137      match ep.direction() {
138        Direction::In => endpoint_in = Some(ep.address()),
139        Direction::Out => endpoint_out = Some(ep.address()),
140      }
141    }
142    let endpoint_in = endpoint_in.ok_or_else(|| Error::InvalidOperation("IN endpoint not found".into()))?;
143    let endpoint_out = endpoint_out.ok_or_else(|| Error::InvalidOperation("OUT endpoint not found".into()))?;
144    tracing::info!("device connected, claiming interface {}", interface_number);
145    if let Some(callback) = &callback {
146      callback(Event::Connected);
147    };
148
149    Ok(Self {
150      inner: Arc::new(AmlInner {
151        handle,
152        interface_number,
153        endpoint_in,
154        endpoint_out,
155      }),
156    })
157  }
158
159  /// Write data to device memory
160  ///
161  /// This writes a small amount of data (up to 64 bytes) to device memory.
162  /// For larger transfers, use `write_large_memory` instead.
163  ///
164  /// # Parameters
165  /// - `address`: The memory address to write to
166  /// - `data`: The data to write, must be <= 64 bytes
167  ///
168  /// # Returns
169  /// - `Result<()>`: Success or an error
170  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
171  pub fn write_simple_memory(&self, address: u32, data: &[u8]) -> Result<()> {
172    tracing::debug!(
173      "writing simple memory at address: {:#X}, length: {}",
174      address,
175      data.len()
176    );
177    if data.len() > 64 {
178      return Err(Error::InvalidOperation("Maximum size of 64 bytes".into()));
179    }
180    let value = (address >> 16) as u16;
181    let index = (address & 0xffff) as u16;
182    self
183      .inner
184      .handle
185      .write_control(0x40, REQ_WRITE_MEM, value, index, data, COMMAND_TIMEOUT)?;
186    tracing::trace!(
187      "write_control completed for write_simple_memory at address: {:#X}",
188      address
189    );
190    Ok(())
191  }
192
193  /// Write arbitrary size data to device memory
194  ///
195  /// This breaks down larger transfers into multiple write_simple_memory operations.
196  ///
197  /// # Parameters
198  /// - `address`: The memory address to write to
199  /// - `data`: The data to write
200  ///
201  /// # Returns
202  /// - `Result<()>`: Success or an error
203  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
204  pub fn write_memory(&self, address: u32, data: &[u8]) -> Result<()> {
205    tracing::debug!(
206      "writing memory starting at address: {:#X} with total length: {}",
207      address,
208      data.len()
209    );
210    let mut offset = 0;
211    let length = data.len();
212    while offset < length {
213      let chunk_size = std::cmp::min(64, length - offset);
214      self.write_simple_memory(address + offset as u32, &data[offset..offset + chunk_size])?;
215      tracing::trace!(
216        "chunk written for write_memory at address: {:#X}, new offset: {}",
217        address,
218        offset + chunk_size
219      );
220      offset += chunk_size;
221    }
222    Ok(())
223  }
224
225  /// Read a small amount of data from device memory
226  ///
227  /// This reads up to 64 bytes from device memory.
228  /// For larger transfers, use `read_memory` instead.
229  ///
230  /// # Parameters
231  /// - `address`: The memory address to read from
232  /// - `length`: The number of bytes to read (must be <= 64)
233  ///
234  /// # Returns
235  /// - `Result<Vec<u8>>`: The read data or an error
236  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
237  pub fn read_simple_memory(&self, address: u32, length: usize) -> Result<Vec<u8>> {
238    tracing::debug!(
239      "reading simple memory at address: {:#X} with length: {}",
240      address,
241      length
242    );
243    if length == 0 {
244      return Ok(vec![]);
245    }
246    if length > 64 {
247      return Err(Error::InvalidOperation("Maximum size of 64 bytes".into()));
248    }
249    let value = (address >> 16) as u16;
250    let index = (address & 0xffff) as u16;
251    let mut buf = vec![0u8; length];
252    let read = self
253      .inner
254      .handle
255      .read_control(0xC0, REQ_READ_MEM, value, index, &mut buf, COMMAND_TIMEOUT)?;
256    tracing::trace!(
257      "read_control completed for read_simple_memory at address: {:#X}, bytes read: {}",
258      address,
259      read
260    );
261    if read != length {
262      return Err(Error::InvalidOperation("Incomplete read".into()));
263    }
264    Ok(buf)
265  }
266
267  /// Read arbitrary size data from device memory
268  ///
269  /// This breaks down larger transfers into multiple read_simple_memory operations.
270  ///
271  /// # Parameters
272  /// - `address`: The memory address to read from
273  /// - `length`: The number of bytes to read
274  ///
275  /// # Returns
276  /// - `Result<Vec<u8>>`: The read data or an error
277  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
278  pub fn read_memory(&self, address: u32, length: usize) -> Result<Vec<u8>> {
279    tracing::debug!("reading memory at address: {:#X} with length: {}", address, length);
280    let mut data = vec![0u8; length];
281    let mut offset = 0;
282    while offset < length {
283      let read_length = std::cmp::min(64, length - offset);
284      let chunk = self.read_simple_memory(address + offset as u32, read_length)?;
285      data[offset..offset + read_length].copy_from_slice(&chunk);
286      tracing::trace!(
287        "chunk read for read_memory at address: {:#X}, offset: {}",
288        address,
289        offset
290      );
291      offset += read_length;
292    }
293    Ok(data)
294  }
295
296  /// Execute code at the specified memory address
297  ///
298  /// # Parameters
299  /// - `address`: The memory address to execute code from
300  /// - `keep_power`: Whether to keep power on after execution
301  ///
302  /// # Returns
303  /// - `Result<()>`: Success or an error
304  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
305  pub fn run(&self, address: u32, keep_power: Option<bool>) -> Result<()> {
306    let keep_power = keep_power.unwrap_or(true);
307    tracing::debug!("running at address: {:#X} with keep_power: {}", address, keep_power);
308    let data = if keep_power {
309      address | FLAG_KEEP_POWER_ON
310    } else {
311      address
312    };
313    let buffer = data.to_le_bytes();
314    let value = (address >> 16) as u16;
315    let index = (address & 0xffff) as u16;
316    self
317      .inner
318      .handle
319      .write_control(0x40, REQ_RUN_IN_ADDR, value, index, &buffer, COMMAND_TIMEOUT)?;
320    tracing::trace!("run command sent at address: {:#X}", address);
321    Ok(())
322  }
323
324  /// Identify the device
325  ///
326  /// # Returns
327  /// - `Result<String>`: The device identification string or an error
328  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
329  pub fn identify(&self) -> Result<String> {
330    tracing::debug!("identifying device");
331    let mut buf = [0u8; 8];
332    let read = self
333      .inner
334      .handle
335      .read_control(0xC0, REQ_IDENTIFY_HOST, 0, 0, &mut buf, COMMAND_TIMEOUT)?;
336    tracing::trace!("identify response received: {:?} ({} bytes)", &buf, read);
337    if read != 8 {
338      return Err(Error::InvalidOperation("Failed to read identify data".into()));
339    }
340    Ok(String::from_utf8(buf.to_vec())?)
341  }
342
343  /// Write large blocks of data to device memory
344  ///
345  /// This is used for writing firmware images and other large data blocks.
346  ///
347  /// # Parameters
348  /// - `memory_address`: The memory address to write to
349  /// - `data`: The data to write
350  /// - `block_length`: The size of each block to transfer
351  /// - `append_zeros`: Whether to pad data with zeros to match block_length
352  ///
353  /// # Returns
354  /// - `Result<()>`: Success or an error
355  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
356  pub fn write_large_memory(
357    &self,
358    memory_address: u32,
359    data: &[u8],
360    block_length: usize,
361    append_zeros: bool,
362  ) -> Result<()> {
363    tracing::debug!(
364      "writing large memory to address: {:#X} with data length: {}",
365      memory_address,
366      data.len()
367    );
368
369    let mut data_vec = data.to_vec();
370    if append_zeros {
371      let remainder = data_vec.len() % block_length;
372      if remainder != 0 {
373        let padding = block_length - remainder;
374        data_vec.extend(vec![0u8; padding]);
375      }
376    } else if data_vec.len() % block_length != 0 {
377      return Err(Error::InvalidOperation(
378        "Large Data must be a multiple of block length".into(),
379      ));
380    }
381
382    let total_bytes = data_vec.len() as u32;
383    let block_count = (data_vec.len() / block_length) as u16;
384    let mut control_data = Vec::with_capacity(16);
385    control_data.extend_from_slice(&memory_address.to_le_bytes());
386    control_data.extend_from_slice(&total_bytes.to_le_bytes());
387    control_data.extend_from_slice(&0u32.to_le_bytes());
388    control_data.extend_from_slice(&0u32.to_le_bytes());
389
390    tracing::trace!("writing control data: {:?}", &control_data);
391    self.inner.handle.write_control(
392      0x40,
393      REQ_WR_LARGE_MEM,
394      block_length as u16,
395      block_count,
396      &control_data,
397      COMMAND_TIMEOUT,
398    )?;
399
400    let mut data_offset = 0;
401    while data_offset < data_vec.len() {
402      let end = data_offset + block_length;
403      let chunk = &data_vec[data_offset..end];
404      tracing::trace!(target: "flashthing::aml::write_large_memory", "writing actual data from offset: {:#X}", &data_offset);
405
406      self
407        .inner
408        .handle
409        .write_bulk(self.inner.endpoint_out, chunk, Duration::from_millis(2000))?;
410
411      tracing::trace!(target: "flashthing::aml::write_large_memory", "wrote actual data from offset: {:#X}", &data_offset);
412
413      data_offset += block_length;
414    }
415
416    Ok(())
417  }
418
419  /// Write large blocks of data directly to a disk address with progress tracking
420  ///
421  /// # Parameters
422  /// - `disk_address`: The disk address to write to
423  /// - `reader`: A reader providing the data to write
424  /// - `data_size`: The total size of data to write
425  /// - `block_length`: The size of each block to transfer
426  /// - `append_zeros`: Whether to pad data with zeros to match block_length
427  /// - `progress_callback`: Function to call with progress updates
428  ///
429  /// # Returns
430  /// - `Result<()>`: Success or an error
431  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
432  pub fn write_large_memory_to_disk<R: std::io::Read, F: Fn(FlashProgress)>(
433    &self,
434    disk_address: u32,
435    reader: &mut R,
436    data_size: usize,
437    block_length: usize,
438    append_zeros: bool,
439    progress_callback: F,
440  ) -> Result<()> {
441    tracing::debug!("streaming {} bytes to disk address: {:#X}", data_size, disk_address);
442
443    let start_time = std::time::Instant::now();
444    let mut total_chunks = 0;
445    let mut avg_chunk_time_secs = 0.0;
446
447    // needed for write operations
448    self.bulkcmd("mmc dev 1")?;
449    self.bulkcmd("amlmmc key")?;
450
451    let total_len = data_size;
452    let max_bytes_per_transfer = TRANSFER_SIZE_THRESHOLD;
453    let mut offset = 0;
454    let mut buffer = vec![0u8; max_bytes_per_transfer];
455
456    while offset < total_len {
457      let chunk_start_time = std::time::Instant::now();
458
459      let remaining = total_len - offset;
460      let write_length = std::cmp::min(remaining, max_bytes_per_transfer);
461
462      let data_slice = &mut buffer[..write_length];
463      reader.read_exact(data_slice)?;
464
465      self.write_large_memory(ADDR_TMP, &buffer[..write_length], block_length, append_zeros)?;
466
467      let start_time_cmd = std::time::Instant::now();
468      let mut retries = 0;
469      let max_retries = 3;
470
471      loop {
472        match self.bulkcmd(&format!(
473          "mmc write {:#X} {:#X} {:#X}",
474          ADDR_TMP,
475          (disk_address as usize + offset) / 512,
476          write_length / 512
477        )) {
478          Ok(_) => {
479            let elapsed = start_time_cmd.elapsed();
480            if elapsed > Duration::from_millis(3000) {
481              tracing::debug!("mmc write command took {}ms, cooling down for 5s", elapsed.as_millis());
482              sleep(Duration::from_secs(5));
483            }
484            break;
485          }
486          Err(e) => {
487            retries += 1;
488            if retries >= max_retries {
489              return Err(e);
490            }
491            sleep(Duration::from_secs(5)); // cooldown after error
492          }
493        }
494      }
495
496      let chunk_time = chunk_start_time.elapsed();
497      let chunk_time_secs = chunk_time.as_secs_f64();
498      total_chunks += 1;
499      if total_chunks == 1 {
500        avg_chunk_time_secs = chunk_time_secs;
501      } else {
502        avg_chunk_time_secs = avg_chunk_time_secs + (chunk_time_secs - avg_chunk_time_secs) / total_chunks as f64;
503      }
504
505      offset += write_length;
506      let progress_percent = offset as f64 / total_len as f64 * 100.0;
507
508      let elapsed = start_time.elapsed();
509      let elapsed_secs = elapsed.as_secs_f64();
510      let bytes_per_sec = if elapsed_secs > 0.0 {
511        offset as f64 / elapsed_secs
512      } else {
513        offset as f64
514      };
515
516      let remaining_bytes = total_len - offset;
517      let eta_secs = if bytes_per_sec > 0.0 {
518        remaining_bytes as f64 / bytes_per_sec
519      } else {
520        0.0
521      };
522
523      tracing::info!(
524        "progress: {:.1}% | elapsed: {:.1}s | eta: {:.1}s | rate: {:.2} KB/s | avg chunk: {:.1}s | avg rate: {:.2} KB/s",
525        progress_percent,
526        elapsed_secs,
527        eta_secs,
528        write_length as f64 / chunk_time_secs / 1024.0,
529        avg_chunk_time_secs,
530        bytes_per_sec / 1024.0
531      );
532
533      progress_callback(FlashProgress {
534        percent: progress_percent,
535        elapsed: elapsed_secs * 1000.0,
536        eta: eta_secs * 1000.0,
537        rate: write_length as f64 / chunk_time_secs / 1024.0,
538        avg_chunk_time: avg_chunk_time_secs * 1000.0,
539        avg_rate: bytes_per_sec / 1024.0,
540      });
541    }
542
543    let total_elapsed = start_time.elapsed();
544    let total_elapsed_secs = total_elapsed.as_secs_f64();
545    let avg_bytes_per_sec = if total_elapsed_secs > 0.0 {
546      total_len as f64 / total_elapsed_secs
547    } else {
548      total_len as f64
549    };
550
551    tracing::info!(
552      "Transfer complete | total time: {:?} | avg rate: {:.2} KB/s",
553      total_elapsed,
554      avg_bytes_per_sec / 1024.0
555    );
556
557    Ok(())
558  }
559
560  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
561  pub fn write_amlc_data(&self, offset: u32, data: &[u8]) -> Result<()> {
562    tracing::debug!("writing amlc data at offset: {:#X} with length: {}", offset, data.len());
563
564    self.inner.handle.write_control(
565      0x40,
566      REQ_WRITE_AMLC,
567      (offset / AMLC_AMLS_BLOCK_LENGTH as u32) as u16,
568      (data.len() - 1) as u16,
569      &[],
570      COMMAND_TIMEOUT,
571    )?;
572    tracing::trace!("amlc header sent for data write at offset: {:#X}", offset);
573
574    let max_chunk_size = AMLC_MAX_BLOCK_LENGTH;
575    let mut data_offset = 0;
576    let write_length = data.len();
577    let mut remaining = write_length;
578
579    let bulk_timeout = Duration::from_millis(1000);
580
581    while remaining > 0 {
582      let block_length = std::cmp::min(remaining, max_chunk_size);
583      let chunk = &data[data_offset..data_offset + block_length];
584
585      let mut retries = 0;
586      let max_retries = 3;
587      let mut success = false;
588
589      while !success && retries < max_retries {
590        match self
591          .inner
592          .handle
593          .write_bulk(self.inner.endpoint_out, chunk, bulk_timeout)
594        {
595          Ok(written) => {
596            if written == block_length {
597              success = true;
598              tracing::trace!(
599                "bulk write in AMLC data, data_offset: {}, chunk: {}",
600                data_offset,
601                block_length
602              );
603            } else {
604              tracing::warn!(
605                "Incomplete bulk write: {} of {} bytes. Retry {}/{}",
606                written,
607                block_length,
608                retries + 1,
609                max_retries
610              );
611              retries += 1;
612              sleep(Duration::from_millis(100));
613            }
614          }
615          Err(e) => {
616            tracing::warn!("Error in bulk write: {}. Retry {}/{}", e, retries + 1, max_retries);
617            retries += 1;
618            sleep(Duration::from_millis(100));
619
620            if retries >= max_retries {
621              return Err(Error::UsbError(e));
622            }
623          }
624        }
625      }
626
627      data_offset += block_length;
628      remaining -= block_length;
629
630      sleep(Duration::from_millis(10));
631    }
632
633    let mut ack_buf = [0u8; 16];
634    let mut retries = 0;
635    let max_retries = 3;
636    let mut read = 0;
637
638    while retries < max_retries {
639      match self
640        .inner
641        .handle
642        .read_bulk(self.inner.endpoint_in, &mut ack_buf, bulk_timeout)
643      {
644        Ok(bytes_read) => {
645          read = bytes_read;
646          if read >= 4 {
647            break;
648          }
649          tracing::warn!("short ack read: {} bytes. retry {}/{}", read, retries + 1, max_retries);
650        }
651        Err(e) => {
652          tracing::warn!("error reading ack: {}. retry {}/{}", e, retries + 1, max_retries);
653        }
654      }
655      retries += 1;
656      sleep(Duration::from_millis(100));
657    }
658
659    tracing::trace!("received amlc ack: {:?} ({} bytes)", &ack_buf[..read], read);
660
661    if read < 4 {
662      return Err(Error::InvalidOperation("no acknowledgment received".into()));
663    }
664
665    let ack = String::from_utf8(ack_buf[0..4].to_vec())?;
666    if ack != "OKAY" {
667      return Err(Error::InvalidOperation(format!("invalid amlc data write ack: {}", ack)));
668    }
669
670    Ok(())
671  }
672
673  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
674  pub fn write_amlc_data_packet(&self, seq: u8, amlc_offset: u32, data: &[u8]) -> Result<()> {
675    tracing::debug!("writing amlc data packet, seq: {}, offset: {:#X}", seq, amlc_offset);
676
677    let data_len = data.len();
678    let max_transfer_length = AMLC_MAX_TRANSFER_LENGTH;
679    let transfer_count = data_len.div_ceil(max_transfer_length);
680
681    if data_len > 0 {
682      let mut offset = 0;
683      for i in 0..transfer_count {
684        let write_length = std::cmp::min(max_transfer_length, data_len - offset);
685        tracing::trace!(
686          "sending amlc data packet chunk {}/{} at offset: {} with length: {}",
687          i + 1,
688          transfer_count,
689          offset,
690          write_length
691        );
692
693        self.write_amlc_data(offset as u32, &data[offset..offset + write_length])?;
694        sleep(Duration::from_millis(50));
695
696        offset += write_length;
697      }
698    }
699
700    let checksum = self.amlc_checksum(data)?;
701
702    let mut amlc_header = [0u8; 16];
703    amlc_header[0..4].copy_from_slice(b"AMLS"); // ! This is AMLS not AMLC for final packet - do not change
704    amlc_header[4] = seq;
705    amlc_header[8..12].copy_from_slice(&checksum.to_le_bytes());
706
707    let mut amlc_data = vec![0u8; AMLC_AMLS_BLOCK_LENGTH];
708    amlc_data[0..16].copy_from_slice(&amlc_header);
709
710    if data.len() > 16 {
711      let copy_len = std::cmp::min(AMLC_AMLS_BLOCK_LENGTH - 16, data.len() - 16);
712      amlc_data[16..16 + copy_len].copy_from_slice(&data[16..16 + copy_len]);
713    }
714
715    tracing::debug!("sending AMLS block with seq {} to offset {:#X}", seq, amlc_offset);
716    self.write_amlc_data(amlc_offset, &amlc_data)?;
717
718    Ok(())
719  }
720
721  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
722  pub fn get_boot_amlc(&self) -> Result<(u32, u32)> {
723    tracing::debug!("getting boot amlc data");
724    self.inner.handle.write_control(
725      0x40,
726      REQ_GET_AMLC,
727      AMLC_AMLS_BLOCK_LENGTH as u16,
728      0,
729      &[],
730      COMMAND_TIMEOUT,
731    )?;
732    tracing::trace!("amlc get request sent");
733    let mut buf = vec![0u8; AMLC_AMLS_BLOCK_LENGTH];
734    let read = self
735      .inner
736      .handle
737      .read_bulk(self.inner.endpoint_in, &mut buf, Duration::from_secs(2))?;
738    tracing::trace!("amlc data received, length: {}", read);
739    if read < AMLC_AMLS_BLOCK_LENGTH {
740      return Err(Error::InvalidOperation("No amlc data received".into()));
741    }
742    let tag = String::from_utf8(buf[0..4].to_vec())?;
743    if tag != "AMLC" {
744      return Err(Error::InvalidOperation(format!("invalid amlc request: {}", tag)));
745    }
746    let length = u32::from_le_bytes(buf[8..12].try_into()?);
747    let offset = u32::from_le_bytes(buf[12..16].try_into()?);
748    let mut ack = [0u8; 16];
749    ack[..4].copy_from_slice(b"OKAY");
750    self
751      .inner
752      .handle
753      .write_bulk(self.inner.endpoint_out, &ack, Duration::from_secs(2))?;
754    tracing::trace!("acknowledgment sent for amlc data");
755    Ok((length, offset))
756  }
757
758  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
759  fn amlc_checksum(&self, data: &[u8]) -> Result<u32> {
760    let mut checksum: u32 = 0;
761    let mut offset = 0;
762    let uint32_max = u32::MAX as u64 + 1;
763    while offset < data.len() {
764      let remaining = data.len() - offset;
765      let val: u32 = if remaining >= 4 {
766        let v = u32::from_le_bytes(data[offset..offset + 4].try_into()?);
767        offset += 4;
768        v
769      } else if remaining >= 3 {
770        let mut temp = [0u8; 4];
771        temp[..remaining].copy_from_slice(&data[offset..]);
772        offset += 3;
773        u32::from_le_bytes(temp) & 0xffffff
774      } else if remaining >= 2 {
775        let v = u16::from_le_bytes(data[offset..offset + 2].try_into()?) as u32;
776        offset += 2;
777        v
778      } else {
779        let v = data[offset] as u32;
780        offset += 1;
781        v
782      };
783      checksum = ((checksum as u64 + (val as i64).unsigned_abs()) % uint32_max) as u32;
784    }
785    Ok(checksum)
786  }
787
788  /// Execute the BL2 boot sequence
789  ///
790  /// This boots the device using the specified BL2 and bootloader binaries.
791  ///
792  /// # Parameters
793  /// - `bl2`: Optional BL2 binary data (uses built-in if None)
794  /// - `bootloader`: Optional bootloader binary data (uses built-in if None)
795  ///
796  /// # Returns
797  /// - `Result<()>`: Success or an error
798  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
799  pub fn bl2_boot(&self, bl2: Option<&[u8]>, bootloader: Option<&[u8]>) -> Result<()> {
800    let bl2 = bl2.unwrap_or(BL2_BIN);
801    let bootloader = bootloader.unwrap_or(BOOTLOADER_BIN);
802
803    tracing::info!("sending bl2 binary to address {:#X}...", ADDR_BL2);
804    self.write_large_memory(ADDR_BL2, bl2, 4096, true)?;
805
806    tracing::info!("booting from bl2...");
807    self.run(ADDR_BL2, Some(true))?;
808
809    tracing::debug!("waiting for bootloader to initialize...");
810    sleep(Duration::from_secs(2));
811
812    let mut prev_length: u32 = 0;
813    let mut prev_offset: u32 = 0;
814    let mut seq: u8 = 0;
815
816    let max_retries = 3;
817    let max_iterations = 50;
818    let mut iterations = 0;
819
820    tracing::info!("starting AMLC data transfer sequence...");
821
822    loop {
823      if iterations >= max_iterations {
824        return Err(Error::InvalidOperation("maximum iterations reached in bl2_boot".into()));
825      }
826      iterations += 1;
827
828      let mut retry_count = 0;
829      let (length, offset) = loop {
830        match self.get_boot_amlc() {
831          Ok(result) => break result,
832          Err(e) => {
833            retry_count += 1;
834            if retry_count >= max_retries {
835              tracing::error!("failed to get boot amlc data after {} attempts: {}", max_retries, e);
836              return Err(e);
837            }
838            tracing::warn!("failed to get boot amlc, retry {}/{}: {}", retry_count, max_retries, e);
839            sleep(Duration::from_millis(500));
840          }
841        }
842      };
843
844      tracing::debug!("amlc request: dataSize={}, offset={}, seq={}", length, offset, seq);
845
846      if length == prev_length && offset == prev_offset {
847        tracing::debug!("amlc transfer complete - received same length/offset twice");
848        break;
849      }
850
851      prev_length = length;
852      prev_offset = offset;
853
854      if offset as usize >= bootloader.len() {
855        tracing::warn!(
856          "amlc requested offset {} exceeds bootloader size {}",
857          offset,
858          bootloader.len()
859        );
860        let empty_slice = &[];
861        self.write_amlc_data_packet(seq, offset, empty_slice)?;
862      } else {
863        let actual_length = std::cmp::min(length as usize, bootloader.len() - offset as usize);
864        let data_slice = &bootloader[offset as usize..offset as usize + actual_length];
865
866        tracing::debug!("sending {} bytes at offset {} with seq {}", actual_length, offset, seq);
867        self.write_amlc_data_packet(seq, offset, data_slice)?;
868      }
869
870      seq = seq.wrapping_add(1);
871      sleep(Duration::from_millis(100));
872    }
873
874    tracing::info!("bl2 boot sequence completed successfully!");
875    Ok(())
876  }
877
878  /// Send a bulk command to the device
879  ///
880  /// # Parameters
881  /// - `command`: The command string to send
882  ///
883  /// # Returns
884  /// - `Result<String>`: The command response or an error
885  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
886  pub fn bulkcmd(&self, command: &str) -> Result<String> {
887    tracing::debug!("sending bulk command: {:?}", command);
888    let mut command = command.as_bytes().to_vec();
889    command.push(0x00);
890    self
891      .inner
892      .handle
893      .write_control(0x40, REQ_BULKCMD, 0, 0, &command, COMMAND_TIMEOUT)?;
894    tracing::trace!("bulk command control write completed");
895
896    let mut buf = vec![0u8; 512];
897    let read = self
898      .inner
899      .handle
900      .read_bulk(self.inner.endpoint_in, &mut buf, COMMAND_TIMEOUT)?;
901    tracing::trace!("bulk command response received, length: {}", read);
902
903    if read == 0 {
904      return Err(Error::InvalidOperation("No response received for bulk command".into()));
905    }
906    let slice = &buf[..read];
907    let start = slice.iter().position(|&b| b != 0).unwrap_or(0);
908    let end = slice.iter().rposition(|&b| b != 0).map(|pos| pos + 1).unwrap_or(0);
909    let trimmed = &slice[start..end];
910    let response = String::from_utf8(trimmed.to_vec())?;
911    if !response.to_lowercase().contains("success") {
912      return Err(Error::InvalidOperation(format!(
913        "Bulk command failed, response did not contain 'success': {}",
914        response
915      )));
916    }
917    Ok(response)
918  }
919
920  /// Validate the size of a partition
921  ///
922  /// # Parameters
923  /// - `part_name`: The name of the partition
924  /// - `part_info`: Partition information
925  ///
926  /// # Returns
927  /// - `Result<usize>`: The validated partition size or an error
928  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
929  pub fn validate_partition_size(&self, part_name: &str, part_info: &PartitionInfo) -> Result<usize> {
930    tracing::debug!("validating partition size for partition: {}", part_name);
931
932    if part_name == "cache" {
933      tracing::warn!("The \"cache\" partition is zero-length on superbird, you cannot read or write to it!");
934      return Err(Error::InvalidOperation("Cache partition is zero-length".into()));
935    }
936
937    if part_name == "reserved" {
938      tracing::warn!("The \"reserved\" partition cannot be read or written!");
939      return Err(Error::InvalidOperation("Reserved partition cannot be accessed".into()));
940    }
941
942    let part_size = part_info.size * PART_SECTOR_SIZE;
943    tracing::info!(
944      "Validating size of partition: {} size: {:#x} {}MB - ...",
945      part_name,
946      part_size,
947      part_size / 1024 / 1024
948    );
949
950    // Try to read the last sector
951    match self.bulkcmd(&format!(
952      "amlmmc read {} {:#x} {:#x} {:#x}",
953      part_name,
954      ADDR_TMP,
955      part_size - PART_SECTOR_SIZE,
956      PART_SECTOR_SIZE
957    )) {
958      Ok(_) => {
959        tracing::info!(
960          "Validating size of partition: {} size: {:#x} {}MB - OK",
961          part_name,
962          part_size,
963          part_size / 1024 / 1024
964        );
965        Ok(part_size)
966      }
967      Err(e) => {
968        tracing::warn!(
969          "Validating size of partition: {} size: {:#x} {}MB - FAIL",
970          part_name,
971          part_size,
972          part_size / 1024 / 1024
973        );
974
975        // Check if it's the data partition which can have an alternate size
976        if part_name == "data" && part_info.size_alt.is_some() {
977          let alt_size = part_info.size_alt.unwrap() * PART_SECTOR_SIZE;
978          tracing::info!(
979            "Failed while fetching last chunk of partition: {}, trying alternate size: {:#x} {}MB",
980            part_name,
981            alt_size,
982            alt_size / 1024 / 1024
983          );
984
985          tracing::info!(
986            "Validating size of partition: {} size: {:#x} {}MB - ...",
987            part_name,
988            alt_size,
989            alt_size / 1024 / 1024
990          );
991
992          match self.bulkcmd(&format!(
993            "amlmmc read {} {:#x} {:#x} {:#x}",
994            part_name,
995            ADDR_TMP,
996            alt_size - PART_SECTOR_SIZE,
997            PART_SECTOR_SIZE
998          )) {
999            Ok(_) => {
1000              tracing::info!(
1001                "Validating size of partition: {} size: {:#x} {}MB - OK",
1002                part_name,
1003                alt_size,
1004                alt_size / 1024 / 1024
1005              );
1006              Ok(alt_size)
1007            }
1008            Err(e2) => {
1009              tracing::error!(
1010                "Validating size of partition: {} size: {:#x} {}MB - FAIL",
1011                part_name,
1012                alt_size,
1013                alt_size / 1024 / 1024
1014              );
1015              tracing::error!(
1016                "Failed while validating size of partition: {}, is partition size {:#x} correct? error: {}",
1017                part_name,
1018                alt_size,
1019                e2
1020              );
1021              Err(e2)
1022            }
1023          }
1024        } else {
1025          tracing::error!(
1026            "Failed while validating size of partition: {}, is partition size {:#x} correct? error: {}",
1027            part_name,
1028            part_size,
1029            e
1030          );
1031          Err(e)
1032        }
1033      }
1034    }
1035  }
1036
1037  /// Restore a partition from a data source
1038  ///
1039  /// # Parameters
1040  /// - `part_name`: The name of the partition to restore
1041  /// - `part_size`: The size of the partition
1042  /// - `reader`: A reader providing the partition data
1043  /// - `file_size`: The size of the data being read
1044  /// - `progress_callback`: Function to call with progress updates
1045  ///
1046  /// # Returns
1047  /// - `Result<()>`: Success or an error
1048  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
1049  pub fn restore_partition<R: Read, F: Fn(FlashProgress)>(
1050    &self,
1051    part_name: &str,
1052    part_size: usize,
1053    mut reader: R,
1054    file_size: usize,
1055    progress_callback: F,
1056  ) -> Result<()> {
1057    tracing::debug!("restoring partition: {} with file size: {}", part_name, file_size);
1058
1059    let adjusted_part_size = if part_name == "bootloader" {
1060      // Bootloader is only 2MB, though dumps may be zero-padded to 4MB
1061      2 * 1024 * 1024
1062    } else {
1063      part_size
1064    };
1065
1066    if file_size > adjusted_part_size && part_name != "bootloader" {
1067      return Err(Error::InvalidOperation(format!(
1068        "file is larger than target partition: {} bytes vs {} bytes",
1069        file_size, adjusted_part_size
1070      )));
1071    }
1072
1073    let start_time = std::time::Instant::now();
1074    let mut total_chunks = 0;
1075    let mut avg_chunk_time_secs = 0.0;
1076
1077    self.bulkcmd("amlmmc key")?;
1078
1079    let total_len = file_size;
1080    let max_bytes_per_transfer = TRANSFER_SIZE_THRESHOLD;
1081    let mut offset = 0;
1082    let mut buffer = vec![0u8; max_bytes_per_transfer];
1083
1084    while offset < total_len {
1085      let chunk_start_time = std::time::Instant::now();
1086
1087      let remaining = total_len - offset;
1088      let write_length = std::cmp::min(remaining, max_bytes_per_transfer);
1089
1090      let data_slice = &mut buffer[..write_length];
1091      reader.read_exact(data_slice)?;
1092
1093      self.write_large_memory(ADDR_TMP, &buffer[..write_length], TRANSFER_BLOCK_SIZE, true)?;
1094
1095      let start_time_cmd = std::time::Instant::now();
1096      let mut retries = 0;
1097      let max_retries = 3;
1098
1099      // Special handling for bootloader partition
1100      if part_name == "bootloader" {
1101        // Bootloader writes always cause timeout - this is expected
1102        match self.bulkcmd(&format!(
1103          "amlmmc write {} {:#x} {:#x} {:#x}",
1104          part_name, ADDR_TMP, offset, write_length
1105        )) {
1106          Ok(_) => tracing::debug!("bootloader write succeeded unexpectedly"),
1107          Err(e) => tracing::debug!("expected timeout for bootloader write: {}", e),
1108        }
1109        sleep(Duration::from_secs(2)); // Allow time for write to complete
1110      } else {
1111        loop {
1112          match self.bulkcmd(&format!(
1113            "amlmmc write {} {:#x} {:#x} {:#x}",
1114            part_name, ADDR_TMP, offset, write_length
1115          )) {
1116            Ok(_) => {
1117              let elapsed = start_time_cmd.elapsed();
1118              if elapsed > Duration::from_millis(3000) {
1119                tracing::debug!("write command took {}ms, cooling down for 5s", elapsed.as_millis());
1120                sleep(Duration::from_secs(5));
1121              }
1122              break;
1123            }
1124            Err(e) => {
1125              retries += 1;
1126              if retries >= max_retries {
1127                return Err(e);
1128              }
1129              tracing::warn!("write command failed, retrying ({}/{}): {}", retries, max_retries, e);
1130              sleep(Duration::from_secs(5)); // cooldown after error
1131            }
1132          }
1133        }
1134      }
1135
1136      let chunk_time = chunk_start_time.elapsed();
1137      let chunk_time_secs = chunk_time.as_secs_f64();
1138      total_chunks += 1;
1139      if total_chunks == 1 {
1140        avg_chunk_time_secs = chunk_time_secs;
1141      } else {
1142        avg_chunk_time_secs = avg_chunk_time_secs + (chunk_time_secs - avg_chunk_time_secs) / total_chunks as f64;
1143      }
1144
1145      offset += write_length;
1146      let progress_percent = offset as f64 / total_len as f64 * 100.0;
1147
1148      let elapsed = start_time.elapsed();
1149      let elapsed_secs = elapsed.as_secs_f64();
1150      let bytes_per_sec = if elapsed_secs > 0.0 {
1151        offset as f64 / elapsed_secs
1152      } else {
1153        offset as f64
1154      };
1155
1156      let remaining_bytes = total_len - offset;
1157      let eta_secs = if bytes_per_sec > 0.0 {
1158        remaining_bytes as f64 / bytes_per_sec
1159      } else {
1160        0.0
1161      };
1162
1163      tracing::info!(
1164        "progress: {:.1}% | elapsed: {:.1}s | eta: {:.1}s | rate: {:.2} KB/s | avg chunk: {:.1}s | avg rate: {:.2} KB/s",
1165        progress_percent,
1166        elapsed_secs,
1167        eta_secs,
1168        write_length as f64 / chunk_time_secs / 1024.0,
1169        avg_chunk_time_secs,
1170        bytes_per_sec / 1024.0
1171      );
1172
1173      progress_callback(FlashProgress {
1174        percent: progress_percent,
1175        elapsed: elapsed_secs * 1000.0,
1176        eta: eta_secs * 1000.0,
1177        rate: write_length as f64 / chunk_time_secs / 1024.0,
1178        avg_chunk_time: avg_chunk_time_secs * 1000.0,
1179        avg_rate: bytes_per_sec / 1024.0,
1180      });
1181    }
1182
1183    let total_elapsed = start_time.elapsed();
1184    let total_elapsed_secs = total_elapsed.as_secs_f64();
1185    let avg_bytes_per_sec = if total_elapsed_secs > 0.0 {
1186      total_len as f64 / total_elapsed_secs
1187    } else {
1188      total_len as f64
1189    };
1190
1191    tracing::info!(
1192      "partition restore complete | total time: {:?} | avg rate: {:.2} KB/s",
1193      total_elapsed,
1194      avg_bytes_per_sec / 1024.0
1195    );
1196
1197    Ok(())
1198  }
1199
1200  /// Execute the unbrick procedure
1201  ///
1202  /// This writes the emergency unbrick image to the device.
1203  ///
1204  /// # Returns
1205  /// - `Result<()>`: Success or an error
1206  #[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
1207  pub fn unbrick(&self) -> Result<()> {
1208    tracing::info!("starting unbrick procedure...");
1209
1210    let cursor = std::io::Cursor::new(UNBRICK_BIN_ZIP);
1211
1212    let mut archive = match zip::ZipArchive::new(cursor) {
1213      Ok(archive) => archive,
1214      Err(e) => {
1215        tracing::error!("failed to open unbrick zip archive: {}", e);
1216        return Err(Error::Zip(e));
1217      }
1218    };
1219
1220    let mut file = match archive.by_name("unbrick.bin") {
1221      Ok(file) => file,
1222      Err(e) => {
1223        tracing::error!("failed to find unbrick.bin in zip archive: {}", e);
1224        return Err(Error::Zip(e));
1225      }
1226    };
1227
1228    let file_size = file.size() as usize;
1229    self.write_large_memory_to_disk(0, &mut file, file_size, TRANSFER_BLOCK_SIZE, true, |progress| {
1230      tracing::info!(
1231        "unbrick progress: {:.1}% | elapsed: {:.1}s | eta: {:.1}s | rate: {:.2} KB/s | avg rate: {:.2} KB/s",
1232        progress.percent,
1233        progress.elapsed,
1234        progress.eta,
1235        progress.rate,
1236        progress.avg_rate
1237      );
1238    })?;
1239
1240    tracing::info!("unbrick procedure completed successfully!");
1241    Ok(())
1242  }
1243
1244  /// Set up the host environment for USB access
1245  ///
1246  /// On Linux, this creates udev rules to allow access to the device.
1247  ///
1248  /// # Returns
1249  /// - `Result<()>`: Success or an error
1250  pub fn host_setup() -> Result<()> {
1251    #[cfg(target_os = "linux")]
1252    crate::setup::setup_host_linux()?;
1253
1254    Ok(())
1255  }
1256}
1257
1258impl Drop for AmlogicSoC {
1259  fn drop(&mut self) {
1260    match self.inner.handle.release_interface(self.inner.interface_number) {
1261      Ok(()) => tracing::trace!("successfully dropped usb interface"),
1262      Err(err) => tracing::warn!("failed to release usb interface: {:?}", err),
1263    }
1264  }
1265}
1266
1267/// The current mode of the Superbird device
1268///
1269/// The device can be in different modes depending on how it was powered on
1270/// and what stage of the boot process it's in.
1271#[derive(Debug, Clone, Copy, PartialEq)]
1272pub enum DeviceMode {
1273  /// Normal operating mode (running regular firmware)
1274  Normal,
1275  /// USB mode (entered by holding buttons 1 & 4 during power-on)
1276  Usb,
1277  /// USB Burn mode (ready for flashing operations)
1278  UsbBurn,
1279  /// Device not detected
1280  NotFound,
1281}
1282
1283#[cfg_attr(feature = "instrument", tracing::instrument(level = "trace", skip_all))]
1284fn find_device() -> DeviceMode {
1285  let context = match Context::new() {
1286    Ok(c) => c,
1287    Err(_) => return DeviceMode::NotFound,
1288  };
1289  let devices = match context.devices() {
1290    Ok(d) => d,
1291    Err(_) => return DeviceMode::NotFound,
1292  };
1293  for device in devices.iter() {
1294    let desc = match device.device_descriptor() {
1295      Ok(d) => d,
1296      Err(_) => continue,
1297    };
1298    // Match normal mode: vendor=0x18d1, product=0x4e40
1299    if desc.vendor_id() == 0x18d1 && desc.product_id() == 0x4e40 {
1300      tracing::debug!("Found device booted normally, with USB Gadget (adb/usbnet) enabled");
1301      return DeviceMode::Normal;
1302    }
1303    // Match USB burn/usb mode: vendor=0x1b8e, product=0xc003
1304    if desc.vendor_id() == 0x1b8e && desc.product_id() == 0xc003 {
1305      // Attempt to open device and read product string
1306      match device.open() {
1307        Ok(handle) => {
1308          // Common language ID
1309          let lang = handle.read_languages(COMMAND_TIMEOUT).unwrap_or_default();
1310          let Some(lang) = lang.first() else {
1311            tracing::debug!("Found device in USB Burn Mode (unable to read product string)");
1312            return DeviceMode::UsbBurn;
1313          };
1314
1315          let prod = handle
1316            .read_product_string(*lang, &desc, Duration::from_millis(100))
1317            .ok();
1318          if prod.as_deref() == Some("GX-CHIP") {
1319            tracing::debug!("Found device booted in USB Mode (buttons 1 & 4 held at boot)");
1320            return DeviceMode::Usb;
1321          } else {
1322            tracing::debug!("Found device booted in USB Burn Mode (ready for commands)");
1323            return DeviceMode::UsbBurn;
1324          }
1325        }
1326        Err(_) => {
1327          tracing::debug!("Found device in USB Burn Mode (unable to read product string)");
1328          return DeviceMode::UsbBurn;
1329        }
1330      }
1331    }
1332  }
1333
1334  tracing::debug!("No device found!");
1335  DeviceMode::NotFound
1336}
1337
1338#[cfg(test)]
1339mod tests {
1340  use super::*;
1341
1342  #[test]
1343  fn test_amlogic_soc_connect() {
1344    let soc = AmlogicSoC::init(None);
1345    // This test will only pass if the device is connected
1346    assert!(soc.is_ok());
1347  }
1348}