xref: /aosp_15_r20/bootable/libbootloader/gbl/libgbl/src/fastboot/mod.rs (revision 5225e6b173e52d2efc6bcf950c27374fd72adabc)
1 // Copyright 2024, The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 //! Fastboot backend for libgbl.
16 
17 use crate::{
18     fuchsia_boot::GblAbrOps,
19     gbl_print, gbl_println,
20     partition::{check_part_unique, GblDisk, PartitionIo},
21     GblOps,
22 };
23 pub use abr::{mark_slot_active, set_one_shot_bootloader, set_one_shot_recovery, SlotIndex};
24 use core::{
25     array::from_fn, cmp::min, ffi::CStr, fmt::Write, future::Future, marker::PhantomData,
26     mem::take, ops::DerefMut, pin::Pin, str::from_utf8,
27 };
28 use fastboot::{
29     next_arg, next_arg_u64, process_next_command, run_tcp_session, snprintf, CommandError,
30     CommandResult, FastbootImplementation, FormattedBytes, InfoSender, OkaySender, RebootMode,
31     UploadBuilder, Uploader, VarInfoSender,
32 };
33 use gbl_async::{join, yield_now};
34 use gbl_storage::{BlockIo, Disk, Gpt};
35 use liberror::Error;
36 use safemath::SafeNum;
37 use zbi::{ZbiContainer, ZbiType};
38 
39 mod vars;
40 
41 pub(crate) mod sparse;
42 use sparse::is_sparse_image;
43 
44 mod shared;
45 pub use shared::Shared;
46 
47 mod buffer_pool;
48 pub use buffer_pool::BufferPool;
49 use buffer_pool::ScopedBuffer;
50 
51 mod pin_fut_container;
52 pub use pin_fut_container::PinFutContainer;
53 use pin_fut_container::{PinFutContainerTyped, PinFutSlice};
54 
55 // Re-exports dependency types
56 pub use fastboot::{TcpStream, Transport};
57 
58 /// Reserved name for indicating flashing GPT.
59 const FLASH_GPT_PART: &str = "gpt";
60 
61 /// Represents a GBL Fastboot async task.
62 enum Task<'a, 'b, B: BlockIo, P: BufferPool> {
63     /// Image flashing task. (partition io, downloaded data, data size)
64     Flash(PartitionIo<'a, B>, ScopedBuffer<'b, P>, usize),
65     // Image erase task.
66     Erase(PartitionIo<'a, B>, ScopedBuffer<'b, P>),
67     None,
68 }
69 
70 impl<'a, 'b, B: BlockIo, P: BufferPool> Task<'a, 'b, B, P> {
71     // Runs the task.
run(self)72     async fn run(self) {
73         let _ = async {
74             match self {
75                 Self::Flash(mut part_io, mut download, data_size) => {
76                     match is_sparse_image(&download) {
77                         Ok(_) => part_io.write_sparse(0, &mut download).await,
78                         _ => part_io.write(0, &mut download[..data_size]).await,
79                     }
80                 }
81                 Self::Erase(mut part_io, mut buffer) => part_io.zeroize(&mut buffer).await,
82                 _ => Ok(()),
83             }
84         }
85         .await;
86     }
87 }
88 
89 /// `GblFastboot` implements fastboot commands in the GBL context.
90 ///
91 /// # Lifetimes
92 ///
93 /// * `'a`: [GblOps] and disks lifetime.
94 /// * `'b`: Lifetime for the buffer allocated by `P`.
95 /// * `'c`: Lifetime of the pinned [Future]s in task container `task`.
96 /// * `'d`: Lifetime of the `tasks` and `gbl_ops` objects borrowed.
97 /// * `'e`: Lifetime of the ImageBuffers returned by `get_image_buffer()`.
98 ///
99 /// # Generics
100 ///
101 /// * `G`: Type of `Self::gbl_ops` which implements [GblOps].
102 /// * `B`: Type that implements [BlockIo] in the [Disk] parameter of [GblDisk] for `Self::disks`.
103 /// * `S`: Type of scratch buffer in the [Disk] parameter of [GblDisk] for `Self::disks`.
104 /// * `T`: Type of gpt buffer in the [Gpt] parameter of [GblDisk] for `Self::disks`.
105 /// * `P`: Type of `Self::buffer_pool` which implements [BufferPool].
106 /// * `C`: Type of `Self::tasks` which implements [PinFutContainerTyped].
107 /// * `F`: Type of [Future] stored by `Self::Tasks`.
108 struct GblFastboot<'a, 'b, 'c, 'd, 'e, G, B, S, T, P, C, F>
109 where
110     G: GblOps<'a, 'e>,
111     B: BlockIo,
112     S: DerefMut<Target = [u8]>,
113     T: DerefMut<Target = [u8]>,
114     P: BufferPool,
115 {
116     pub(crate) gbl_ops: &'d mut G,
117     // We store the partition devices returned by `gbl_ops.disks()` directly instead of getting it
118     // from `gbl_ops` later because we need to establish to the compiler that the hidden type of
119     // [BlockIo] in `GblDisk<Disk<impl BlockIO...>...>` returned by `gbl_ops.disks()` will be the
120     // same as the [BlockIo] type (denoted as B) in the function pointer
121     // `task_mapper`: fn(Task<'a, 'b, B, P>) -> F`. Otherwise, compiler won't allow `fn flash()`
122     // to call `task_mapper` with a `Task` constructed from `GblDisk<Disk<impl BlockIO...>...>`.
123     disks: &'a [GblDisk<Disk<B, S>, Gpt<T>>],
124     buffer_pool: &'b Shared<P>,
125     task_mapper: fn(Task<'a, 'b, B, P>) -> F,
126     tasks: &'d Shared<C>,
127     current_download_buffer: Option<ScopedBuffer<'b, P>>,
128     current_download_size: usize,
129     enable_async_task: bool,
130     default_block: Option<usize>,
131     bootimg_buf: &'b mut [u8],
132     // Introduces marker type so that we can enforce constraint 'd <= min('b, 'c).
133     // The constraint is expressed in the implementation block for the `FastbootImplementation`
134     // trait.
135     _tasks_context_lifetime: PhantomData<&'c P>,
136     _get_image_buffer_lifetime: PhantomData<&'e ()>,
137 }
138 
139 // See definition of [GblFastboot] for docs on lifetimes and generics parameters.
140 impl<'a: 'c, 'b: 'c, 'c, 'd, 'e, G, B, S, T, P, C, F>
141     GblFastboot<'a, 'b, 'c, 'd, 'e, G, B, S, T, P, C, F>
142 where
143     G: GblOps<'a, 'e>,
144     B: BlockIo,
145     S: DerefMut<Target = [u8]>,
146     T: DerefMut<Target = [u8]>,
147     P: BufferPool,
148     C: PinFutContainerTyped<'c, F>,
149     F: Future<Output = ()> + 'c,
150 {
151     /// Creates a new [GblFastboot].
152     ///
153     /// # Args
154     ///
155     /// * `gbl_ops`: An implementation of `GblOps`.
156     /// * `disks`: The disk devices returned by `gbl_ops.disks()`. This is needed for expressing the
157     ///   property that the hidden [BlockIo] type is the same as that in `task_mapper`.
158     /// * `task_mapper`: A function pointer that maps `Task<'a, 'b, G, B>` to the target [Future]
159     ///   type `F` for input to `PinFutContainerTyped<F>::add_with()`.
160     /// * `tasks`: A shared instance of `PinFutContainerTyped<F>`.
161     /// * `buffer_pool`: A shared instance of `BufferPool`.
162     ///
163     /// The combination of `task_mapper` and `tasks` allows type `F`, which will be running the
164     /// async function `Task::run()`, to be defined at the callsite. This is necessary for the
165     /// usage of preallocated pinned futures (by `run_gbl_fastboot_stack()`) because the returned
166     /// type of a `async fn` is compiler-generated and can't be named. The only way to create a
167     /// preallocated slice of anonymous future is to keep the type generic and pass in the
168     /// anonymous future instance at the initialization callsite (aka defining use) and let compiler
169     /// infer and propagate it.
new( gbl_ops: &'d mut G, disks: &'a [GblDisk<Disk<B, S>, Gpt<T>>], task_mapper: fn(Task<'a, 'b, B, P>) -> F, tasks: &'d Shared<C>, buffer_pool: &'b Shared<P>, bootimg_buf: &'b mut [u8], ) -> Self170     fn new(
171         gbl_ops: &'d mut G,
172         disks: &'a [GblDisk<Disk<B, S>, Gpt<T>>],
173         task_mapper: fn(Task<'a, 'b, B, P>) -> F,
174         tasks: &'d Shared<C>,
175         buffer_pool: &'b Shared<P>,
176         bootimg_buf: &'b mut [u8],
177     ) -> Self {
178         Self {
179             gbl_ops,
180             disks,
181             task_mapper,
182             tasks,
183             buffer_pool,
184             current_download_buffer: None,
185             current_download_size: 0,
186             enable_async_task: false,
187             default_block: None,
188             bootimg_buf,
189             _tasks_context_lifetime: PhantomData,
190             _get_image_buffer_lifetime: PhantomData,
191         }
192     }
193 
194     /// Returns the shared task container.
tasks(&self) -> &'d Shared<impl PinFutContainerTyped<'c, F>>195     fn tasks(&self) -> &'d Shared<impl PinFutContainerTyped<'c, F>> {
196         self.tasks
197     }
198 
199     /// Listens on the given USB/TCP channels and runs fastboot.
run( &mut self, mut usb: Option<impl GblUsbTransport>, mut tcp: Option<impl GblTcpStream>, )200     async fn run(
201         &mut self,
202         mut usb: Option<impl GblUsbTransport>,
203         mut tcp: Option<impl GblTcpStream>,
204     ) {
205         if usb.is_none() && tcp.is_none() {
206             gbl_println!(self.gbl_ops, "No USB or TCP found for GBL Fastboot");
207             return;
208         }
209         let tasks = self.tasks();
210         // The fastboot command loop task for interacting with the remote host.
211         let cmd_loop_end = Shared::from(false);
212         let cmd_loop_task = async {
213             loop {
214                 if let Some(v) = usb.as_mut() {
215                     if v.has_packet() {
216                         let res = match process_next_command(v, self).await {
217                             Ok(true) => break,
218                             v => v,
219                         };
220                         if res.is_err() {
221                             gbl_println!(self.gbl_ops, "GBL Fastboot USB session error: {:?}", res);
222                         }
223                     }
224                 }
225 
226                 if let Some(v) = tcp.as_mut() {
227                     if v.accept_new() {
228                         let res = match run_tcp_session(v, self).await {
229                             Ok(()) => break,
230                             v => v,
231                         };
232                         if res.is_err_and(|e| e != Error::Disconnected) {
233                             gbl_println!(self.gbl_ops, "GBL Fastboot TCP session error: {:?}", res);
234                         }
235                     }
236                 }
237 
238                 yield_now().await;
239             }
240             *cmd_loop_end.borrow_mut() = true;
241         };
242 
243         // Schedules [Task] spawned by GBL fastboot.
244         let gbl_fb_tasks = async {
245             while tasks.borrow_mut().poll_all() > 0 || !*cmd_loop_end.borrow_mut() {
246                 yield_now().await;
247             }
248         };
249 
250         let _ = join(cmd_loop_task, gbl_fb_tasks).await;
251     }
252 
253     /// Extracts the next argument and verifies that it is a valid block device ID if present.
254     ///
255     /// # Returns
256     ///
257     /// * Returns `Ok(Some(blk_id))` if next argument is present and is a valid block device ID.
258     /// * Returns `None` if next argument is not available and there are more than one block
259     ///   devices.
260     /// * Returns `Err(())` if next argument is present but is an invalid block device ID.
check_next_arg_blk_id<'s>( &self, args: &mut impl Iterator<Item = &'s str>, ) -> CommandResult<Option<usize>>261     fn check_next_arg_blk_id<'s>(
262         &self,
263         args: &mut impl Iterator<Item = &'s str>,
264     ) -> CommandResult<Option<usize>> {
265         let devs = self.disks;
266         let blk_id = match next_arg_u64(args)? {
267             Some(v) => {
268                 let v = usize::try_from(v)?;
269                 // Checks out of range.
270                 devs.get(v).ok_or("Invalid block ID")?;
271                 Some(v)
272             }
273             _ => None,
274         };
275         let blk_id = blk_id.or(self.default_block);
276         let blk_id = blk_id.or((devs.len() == 1).then_some(0));
277         Ok(blk_id)
278     }
279 
280     /// Parses and checks the argument for "fastboot flash gpt/<blk_idx>/"resize".
281     ///
282     /// # Returns
283     ///
284     /// * Returns `Ok(Some((blk_idx, resize)))` if command is a GPT flashing command.
285     /// * Returns `Ok(None)` if command is not a GPT flashing command.
286     /// * Returns `Err()` otherwise.
parse_flash_gpt_args(&self, part: &str) -> CommandResult<Option<(usize, bool)>>287     pub(crate) fn parse_flash_gpt_args(&self, part: &str) -> CommandResult<Option<(usize, bool)>> {
288         // Syntax: flash gpt/<blk_idx>/"resize"
289         let mut args = part.split('/');
290         if next_arg(&mut args).filter(|v| *v == FLASH_GPT_PART).is_none() {
291             return Ok(None);
292         }
293         // Parses block device ID.
294         let blk_id = self
295             .check_next_arg_blk_id(&mut args)?
296             .ok_or("Block ID is required for flashing GPT")?;
297         // Parses resize option.
298         let resize = match next_arg(&mut args) {
299             Some("resize") => true,
300             Some(_) => return Err("Unknown argument".into()),
301             _ => false,
302         };
303         Ok(Some((blk_id, resize)))
304     }
305 
306     /// Parses and checks the partition argument and returns the partition name, block device
307     /// index, start offset and size.
parse_partition<'s>( &self, part: &'s str, ) -> CommandResult<(Option<&'s str>, usize, u64, u64)>308     pub(crate) fn parse_partition<'s>(
309         &self,
310         part: &'s str,
311     ) -> CommandResult<(Option<&'s str>, usize, u64, u64)> {
312         let devs = self.disks;
313         let mut args = part.split('/');
314         // Parses partition name.
315         let part = next_arg(&mut args);
316         // Parses block device ID.
317         let blk_id = self.check_next_arg_blk_id(&mut args)?;
318         // Parses sub window offset.
319         let window_offset = next_arg_u64(&mut args)?.unwrap_or(0);
320         // Parses sub window size.
321         let window_size = next_arg_u64(&mut args)?;
322         // Checks uniqueness of the partition and resolves its block device ID.
323         let find = |p: Option<&'s str>| match blk_id {
324             None => Ok((check_part_unique(devs, p.ok_or("Must provide a partition")?)?, p)),
325             Some(v) => Ok(((v, devs[v].find_partition(p)?), p)),
326         };
327         let ((blk_id, partition), actual) = match find(part) {
328             // Some legacy Fuchsia devices in the field uses name "fuchsia-fvm" for the standard
329             // "fvm" partition. However all of our infra uses the standard name "fvm" when flashing.
330             // Here we do a one off mapping if the device falls into this case. Once we have a
331             // solution for migrating those devices off the legacy name, we can remove this.
332             //
333             // If we run into more of such legacy aliases that we can't migrate, consider adding
334             // interfaces in GblOps for this.
335             Err(Error::NotFound) if part == Some("fvm") => find(Some("fuchsia-fvm"))?,
336             v => v?,
337         };
338         let part_sz = SafeNum::from(partition.size()?);
339         let window_size = window_size.unwrap_or((part_sz - window_offset).try_into()?);
340         u64::try_from(part_sz - window_size - window_offset)?;
341         Ok((actual, blk_id, window_offset, window_size))
342     }
343 
344     /// Takes the download data and resets download size.
take_download(&mut self) -> Option<(ScopedBuffer<'b, P>, usize)>345     fn take_download(&mut self) -> Option<(ScopedBuffer<'b, P>, usize)> {
346         Some((self.current_download_buffer.take()?, take(&mut self.current_download_size)))
347     }
348 
349     /// Waits until a Disk device is ready and get the [PartitionIo] for `part`.
wait_partition_io( &self, blk: usize, part: Option<&str>, ) -> CommandResult<PartitionIo<'a, B>>350     pub async fn wait_partition_io(
351         &self,
352         blk: usize,
353         part: Option<&str>,
354     ) -> CommandResult<PartitionIo<'a, B>> {
355         loop {
356             match self.disks[blk].partition_io(part) {
357                 Err(Error::NotReady) => yield_now().await,
358                 Ok(v) => {
359                     v.last_err()?;
360                     return Ok(v);
361                 }
362                 Err(e) => return Err(e.into()),
363             }
364         }
365     }
366 
367     /// An internal helper for parsing a partition and getting the partition IO
parse_and_get_partition_io( &self, part: &str, ) -> CommandResult<(usize, PartitionIo<'a, B>)>368     async fn parse_and_get_partition_io(
369         &self,
370         part: &str,
371     ) -> CommandResult<(usize, PartitionIo<'a, B>)> {
372         let (part, blk_idx, start, sz) = self.parse_partition(part)?;
373         Ok((blk_idx, self.wait_partition_io(blk_idx, part).await?.sub(start, sz)?))
374     }
375 
376     /// Helper for scheduiling an async task.
377     ///
378     /// * If `Self::enable_async_task` is true, the method will add the task to the background task
379     ///   list. Otherwise it simply runs the task.
schedule_task( &mut self, task: Task<'a, 'b, B, P>, responder: &mut impl InfoSender, ) -> CommandResult<()>380     async fn schedule_task(
381         &mut self,
382         task: Task<'a, 'b, B, P>,
383         responder: &mut impl InfoSender,
384     ) -> CommandResult<()> {
385         match self.enable_async_task {
386             true => {
387                 let mut t = Some((self.task_mapper)(task));
388                 self.tasks.borrow_mut().add_with(|| t.take().unwrap());
389                 while t.is_some() {
390                     yield_now().await;
391                     self.tasks.borrow_mut().add_with(|| t.take().unwrap());
392                 }
393                 self.tasks.borrow_mut().poll_all();
394                 let info = "An IO task is launched. To sync manually, run \"oem gbl-sync-tasks\".";
395                 responder.send_info(info).await?
396             }
397             _ => task.run().await,
398         };
399         Ok(())
400     }
401 
402     /// Waits for all block devices to be ready.
sync_all_blocks(&self) -> CommandResult<()>403     async fn sync_all_blocks(&self) -> CommandResult<()> {
404         for (idx, _) in self.disks.iter().enumerate() {
405             let _ = self.wait_partition_io(idx, None).await;
406         }
407         Ok(())
408     }
409 
410     /// Implementation for "fastboot oem gbl-sync-tasks".
oem_sync_blocks<'s>( &self, mut responder: impl InfoSender, res: &'s mut [u8], ) -> CommandResult<&'s [u8]>411     async fn oem_sync_blocks<'s>(
412         &self,
413         mut responder: impl InfoSender,
414         res: &'s mut [u8],
415     ) -> CommandResult<&'s [u8]> {
416         self.sync_all_blocks().await?;
417         // Checks error.
418         let mut has_error = false;
419         for (i, ele) in self.disks.iter().enumerate() {
420             match ele.partition_io(None)?.last_err() {
421                 Ok(_) => {}
422                 Err(e) => {
423                     has_error = true;
424                     responder.send_info(snprintf!(res, "Block #{} error: {:?}.", i, e)).await?;
425                 }
426             }
427         }
428         match has_error {
429             true => Err("Errors during async block IO. Please reset device.".into()),
430             _ => Ok(b""),
431         }
432     }
433 
434     /// Syncs all storage devices and reboots.
sync_block_and_reboot( &mut self, mode: RebootMode, mut resp: impl InfoSender + OkaySender, ) -> CommandResult<()>435     async fn sync_block_and_reboot(
436         &mut self,
437         mode: RebootMode,
438         mut resp: impl InfoSender + OkaySender,
439     ) -> CommandResult<()> {
440         resp.send_info("Syncing storage...").await?;
441         self.sync_all_blocks().await?;
442         match mode {
443             RebootMode::Normal => {
444                 resp.send_info("Rebooting...").await?;
445                 resp.send_okay("").await?;
446                 self.gbl_ops.reboot();
447             }
448             RebootMode::Bootloader => {
449                 let f = self.gbl_ops.reboot_bootloader()?;
450                 resp.send_info("Rebooting to bootloader...").await?;
451                 resp.send_okay("").await?;
452                 f()
453             }
454             RebootMode::Recovery => {
455                 let f = self.gbl_ops.reboot_recovery()?;
456                 resp.send_info("Rebooting to recovery...").await?;
457                 resp.send_okay("").await?;
458                 f()
459             }
460             _ => return Err("Unsupported".into()),
461         }
462         Ok(())
463     }
464 
465     /// Appends a staged payload as bootloader file.
add_staged_bootloader_file(&mut self, file_name: &str) -> CommandResult<()>466     async fn add_staged_bootloader_file(&mut self, file_name: &str) -> CommandResult<()> {
467         let buffer = self
468             .gbl_ops
469             .get_zbi_bootloader_files_buffer_aligned()
470             .ok_or("No ZBI bootloader file buffer is provided")?;
471         let data = self.current_download_buffer.as_mut().ok_or("No file staged")?;
472         let data = &mut data[..self.current_download_size];
473         let mut zbi = match ZbiContainer::parse(&mut buffer[..]) {
474             Ok(v) => v,
475             _ => ZbiContainer::new(&mut buffer[..])?,
476         };
477         let next_payload = zbi.get_next_payload()?;
478         // Format: name length (1 byte) | name | file content.
479         let (name_len, rest) = next_payload.split_at_mut_checked(1).ok_or("Buffer too small")?;
480         let (name, rest) = rest.split_at_mut_checked(file_name.len()).ok_or("Buffer too small")?;
481         let file_content = rest.get_mut(..data.len()).ok_or("Buffer too small")?;
482         name_len[0] = file_name.len().try_into().map_err(|_| "File name length overflows 256")?;
483         name.clone_from_slice(file_name.as_bytes());
484         file_content.clone_from_slice(data);
485         // Creates the entry;
486         zbi.create_entry(
487             ZbiType::BootloaderFile,
488             0,
489             Default::default(),
490             1 + file_name.len() + data.len(),
491         )?;
492         Ok(())
493     }
494 }
495 
496 // See definition of [GblFastboot] for docs on lifetimes and generics parameters.
497 impl<'a: 'c, 'b: 'c, 'c, 'e, G, B, S, T, P, C, F> FastbootImplementation
498     for GblFastboot<'a, 'b, 'c, '_, 'e, G, B, S, T, P, C, F>
499 where
500     G: GblOps<'a, 'e>,
501     B: BlockIo,
502     S: DerefMut<Target = [u8]>,
503     T: DerefMut<Target = [u8]>,
504     P: BufferPool,
505     C: PinFutContainerTyped<'c, F>,
506     F: Future<Output = ()> + 'c,
507 {
get_var( &mut self, var: &CStr, args: impl Iterator<Item = &'_ CStr> + Clone, out: &mut [u8], _: impl InfoSender, ) -> CommandResult<usize>508     async fn get_var(
509         &mut self,
510         var: &CStr,
511         args: impl Iterator<Item = &'_ CStr> + Clone,
512         out: &mut [u8],
513         _: impl InfoSender,
514     ) -> CommandResult<usize> {
515         Ok(self.get_var_internal(var, args, out)?.len())
516     }
517 
get_var_all(&mut self, mut resp: impl VarInfoSender) -> CommandResult<()>518     async fn get_var_all(&mut self, mut resp: impl VarInfoSender) -> CommandResult<()> {
519         self.get_var_all_internal(&mut resp).await
520     }
521 
get_download_buffer(&mut self) -> &mut [u8]522     async fn get_download_buffer(&mut self) -> &mut [u8] {
523         if self.current_download_buffer.is_none() {
524             self.current_download_buffer = Some(self.buffer_pool.allocate_async().await);
525         }
526         self.current_download_buffer.as_mut().unwrap()
527     }
528 
download_complete( &mut self, download_size: usize, _: impl InfoSender, ) -> CommandResult<()>529     async fn download_complete(
530         &mut self,
531         download_size: usize,
532         _: impl InfoSender,
533     ) -> CommandResult<()> {
534         self.current_download_size = download_size;
535         Ok(())
536     }
537 
flash(&mut self, part: &str, mut responder: impl InfoSender) -> CommandResult<()>538     async fn flash(&mut self, part: &str, mut responder: impl InfoSender) -> CommandResult<()> {
539         let disks = self.disks;
540 
541         // Checks if we are flashing new GPT partition table
542         if let Some((blk_idx, resize)) = self.parse_flash_gpt_args(part)? {
543             self.wait_partition_io(blk_idx, None).await?;
544             let (mut gpt, size) = self.take_download().ok_or("No GPT downloaded")?;
545             responder.send_info("Updating GPT...").await?;
546             return match disks[blk_idx].update_gpt(&mut gpt[..size], resize).await {
547                 Err(Error::NotReady) => panic!("Should not be busy"),
548                 Err(Error::Unsupported) => Err("Block device is not for GPT".into()),
549                 v => Ok(v?),
550             };
551         }
552 
553         let (blk_idx, part_io) = self.parse_and_get_partition_io(part).await?;
554         let (download_buffer, data_size) = self.take_download().ok_or("No download")?;
555         let write_task = Task::Flash(part_io, download_buffer, data_size);
556         self.schedule_task(write_task, &mut responder).await?;
557         // Checks if block is ready already and returns errors. This can be the case when the
558         // operation is synchronous or runs into early errors.
559         Ok(disks[blk_idx].status().result()?)
560     }
561 
erase(&mut self, part: &str, mut responder: impl InfoSender) -> CommandResult<()>562     async fn erase(&mut self, part: &str, mut responder: impl InfoSender) -> CommandResult<()> {
563         let disks = self.disks;
564 
565         // Checks if we are erasing GPT partition table.
566         if let Some((blk_idx, _)) = self.parse_flash_gpt_args(part)? {
567             self.wait_partition_io(blk_idx, None).await?;
568             return match disks[blk_idx].erase_gpt().await {
569                 Err(Error::NotReady) => panic!("Should not be busy"),
570                 Err(Error::Unsupported) => Err("Block device is not for GPT".into()),
571                 v => Ok(v?),
572             };
573         }
574 
575         let (blk_idx, part_io) = self.parse_and_get_partition_io(part).await?;
576         self.get_download_buffer().await;
577         let erase_task = Task::Erase(part_io, self.take_download().unwrap().0);
578         self.schedule_task(erase_task, &mut responder).await?;
579         // Checks if block is ready already and returns errors. This can be the case when the
580         // operation is synchronous or runs into early errors.
581         Ok(disks[blk_idx].status().result()?)
582     }
583 
upload(&mut self, _: impl UploadBuilder) -> CommandResult<()>584     async fn upload(&mut self, _: impl UploadBuilder) -> CommandResult<()> {
585         Err("Unimplemented".into())
586     }
587 
fetch( &mut self, part: &str, offset: u64, size: u64, mut responder: impl UploadBuilder + InfoSender, ) -> CommandResult<()>588     async fn fetch(
589         &mut self,
590         part: &str,
591         offset: u64,
592         size: u64,
593         mut responder: impl UploadBuilder + InfoSender,
594     ) -> CommandResult<()> {
595         let (_, mut part_io) = self.parse_and_get_partition_io(part).await?;
596         let buffer = self.get_download_buffer().await;
597         let end = u64::try_from(SafeNum::from(offset) + size)?;
598         let mut curr = offset;
599         responder
600             .send_formatted_info(|v| write!(v, "Uploading {} bytes...", size).unwrap())
601             .await?;
602         let mut uploader = responder.initiate_upload(size).await?;
603         while curr < end {
604             let to_send = min(usize::try_from(end - curr)?, buffer.len());
605             part_io.read(curr, &mut buffer[..to_send]).await?;
606             uploader.upload(&mut buffer[..to_send]).await?;
607             curr += u64::try_from(to_send)?;
608         }
609         Ok(())
610     }
611 
reboot( &mut self, mode: RebootMode, resp: impl InfoSender + OkaySender, ) -> CommandError612     async fn reboot(
613         &mut self,
614         mode: RebootMode,
615         resp: impl InfoSender + OkaySender,
616     ) -> CommandError {
617         match self.sync_block_and_reboot(mode, resp).await {
618             Err(e) => e,
619             _ => "Unknown".into(),
620         }
621     }
622 
623     async fn r#continue(&mut self, mut resp: impl InfoSender) -> CommandResult<()> {
624         resp.send_info("Syncing storage...").await?;
625         Ok(self.sync_all_blocks().await?)
626     }
627 
set_active(&mut self, slot: &str, _: impl InfoSender) -> CommandResult<()>628     async fn set_active(&mut self, slot: &str, _: impl InfoSender) -> CommandResult<()> {
629         self.sync_all_blocks().await?;
630         match self.gbl_ops.expected_os_is_fuchsia()? {
631             // TODO(b/374776896): Prioritizes platform specific `set_active_slot`  if available.
632             true => Ok(mark_slot_active(
633                 &mut GblAbrOps(self.gbl_ops),
634                 match slot {
635                     "a" => SlotIndex::A,
636                     "b" => SlotIndex::B,
637                     _ => return Err("Invalid slot index for Fuchsia A/B/R".into()),
638                 },
639             )?),
640             _ => Err("Not supported".into()),
641         }
642     }
643 
oem<'s>( &mut self, cmd: &str, mut responder: impl InfoSender, res: &'s mut [u8], ) -> CommandResult<&'s [u8]>644     async fn oem<'s>(
645         &mut self,
646         cmd: &str,
647         mut responder: impl InfoSender,
648         res: &'s mut [u8],
649     ) -> CommandResult<&'s [u8]> {
650         let mut args = cmd.split(' ');
651         let cmd = args.next().ok_or("Missing command")?;
652         match cmd {
653             "gbl-sync-tasks" => self.oem_sync_blocks(responder, res).await,
654             "gbl-enable-async-task" => {
655                 self.enable_async_task = true;
656                 Ok(b"")
657             }
658             "gbl-disable-async-task" => {
659                 self.enable_async_task = false;
660                 Ok(b"")
661             }
662             "gbl-unset-default-block" => {
663                 self.default_block = None;
664                 Ok(b"")
665             }
666             "gbl-set-default-block" => {
667                 let id = next_arg_u64(&mut args)?.ok_or("Missing block device ID")?;
668                 let id = usize::try_from(id)?;
669                 self.disks.get(id).ok_or("Out of range")?;
670                 self.default_block = Some(id.try_into()?);
671                 responder
672                     .send_formatted_info(|f| write!(f, "Default block device: {id:#x}").unwrap())
673                     .await?;
674                 Ok(b"")
675             }
676             "add-staged-bootloader-file" => {
677                 let file_name = next_arg(&mut args).ok_or("Missing file name")?;
678                 self.add_staged_bootloader_file(file_name).await?;
679                 Ok(b"")
680             }
681             _ => Err("Unknown oem command".into()),
682         }
683     }
684 
boot(&mut self, mut resp: impl InfoSender + OkaySender) -> CommandResult<()>685     async fn boot(&mut self, mut resp: impl InfoSender + OkaySender) -> CommandResult<()> {
686         let len = core::cmp::min(self.bootimg_buf.len(), self.current_download_size);
687         let data = self.current_download_buffer.as_mut().ok_or("No file staged")?;
688         let data = &mut data[..self.current_download_size];
689 
690         self.bootimg_buf[..len].copy_from_slice(&data[..len]);
691         resp.send_info("Boot into boot.img").await?;
692         Ok(())
693     }
694 }
695 
696 /// `GblUsbTransport` defines transport interfaces for running GBL fastboot over USB.
697 pub trait GblUsbTransport: Transport {
698     /// Checks whether there is a new USB packet.
has_packet(&mut self) -> bool699     fn has_packet(&mut self) -> bool;
700 }
701 
702 /// `GblTcpStream` defines transport interfaces for running GBL fastboot over TCP.
703 pub trait GblTcpStream: TcpStream {
704     /// Accepts a new TCP connection.
705     ///
706     /// If a connection is in progress, it should be aborted first.
707     ///
708     /// Returns true if a new connection is established, false otherwise.
accept_new(&mut self) -> bool709     fn accept_new(&mut self) -> bool;
710 }
711 
712 /// Runs GBL fastboot on the given USB/TCP channels.
713 ///
714 /// # Args:
715 ///
716 /// * `gbl_ops`: An instance of [GblOps].
717 /// * `buffer_pool`: An implementation of [BufferPool].
718 /// * `tasks`: An implementation of [PinFutContainer]
719 /// * `usb`: An optional implementation of [GblUsbTransport].
720 /// * `tcp`: An optional implementation of [GblTcpStream].
721 ///
722 /// # Lifetimes
723 /// * `'a`: Lifetime of [GblOps].
724 /// * `'b`: Lifetime of `download_buffers`.
725 /// * `'c`: Lifetime of `tasks`.
run_gbl_fastboot<'a: 'c, 'b: 'c, 'c, 'd>( gbl_ops: &mut impl GblOps<'a, 'd>, buffer_pool: &'b Shared<impl BufferPool>, tasks: impl PinFutContainer<'c> + 'c, usb: Option<impl GblUsbTransport>, tcp: Option<impl GblTcpStream>, bootimg_buf: &'b mut [u8], )726 pub async fn run_gbl_fastboot<'a: 'c, 'b: 'c, 'c, 'd>(
727     gbl_ops: &mut impl GblOps<'a, 'd>,
728     buffer_pool: &'b Shared<impl BufferPool>,
729     tasks: impl PinFutContainer<'c> + 'c,
730     usb: Option<impl GblUsbTransport>,
731     tcp: Option<impl GblTcpStream>,
732     bootimg_buf: &'b mut [u8],
733 ) {
734     let tasks = tasks.into();
735     let disks = gbl_ops.disks();
736     GblFastboot::new(gbl_ops, disks, Task::run, &tasks, buffer_pool, bootimg_buf)
737         .run(usb, tcp)
738         .await;
739 }
740 
741 /// Runs GBL fastboot on the given USB/TCP channels with N stack allocated worker tasks.
742 ///
743 /// The choice of N depends on the level of parallelism the platform can support. For platform with
744 /// `n` storage devices that can independently perform non-blocking IO, it will required `N = n`
745 /// and a `buffer_pool` that can allocate at least n+1 buffers at the same time in order to achieve
746 /// parallel flashing to all storages plus a parallel downloading. However, it is common for
747 /// disks that need to be flashed to be on the same block deviece so flashing of them becomes
748 /// sequential, in which case N can be smaller. Caller should take into consideration usage pattern
749 /// for determining N.
750 ///
751 /// # Args:
752 ///
753 /// * `gbl_ops`: An instance of [GblOps].
754 /// * `buffer_pool`: An implementation of [BufferPool].
755 /// * `usb`: An optional implementation of [GblUsbTransport].
756 /// * `tcp`: An optional implementation of [GblTcpStream].
run_gbl_fastboot_stack<'a, 'b, const N: usize>( gbl_ops: &mut impl GblOps<'a, 'b>, buffer_pool: impl BufferPool, usb: Option<impl GblUsbTransport>, tcp: Option<impl GblTcpStream>, bootimg_buf: &mut [u8], )757 pub async fn run_gbl_fastboot_stack<'a, 'b, const N: usize>(
758     gbl_ops: &mut impl GblOps<'a, 'b>,
759     buffer_pool: impl BufferPool,
760     usb: Option<impl GblUsbTransport>,
761     tcp: Option<impl GblTcpStream>,
762     bootimg_buf: &mut [u8],
763 ) {
764     let buffer_pool = buffer_pool.into();
765     // Creates N worker tasks.
766     let mut tasks: [_; N] = from_fn(|_| Task::None.run());
767     // It is possible to avoid the use of the unsafe `Pin::new_unchecked` by delaring the array and
768     // manually pinning each element i.e.
769     //
770     // ```
771     // let mut tasks = [
772     //     core::pin::pin!(Task::None.run()),
773     //     core::pin::pin!(Task::None.run()),
774     //     core::pin::pin!(Task::None.run()),
775     // ];
776     // ```
777     //
778     // Parameterization of `N` will be an issue, but might be solvable with procedural macro.
779     // SAFETY: `tasks` is immediately shadowed and thus guaranteed not moved for the rest of its
780     // lifetime.
781     let mut tasks: [_; N] = tasks.each_mut().map(|v| unsafe { Pin::new_unchecked(v) });
782     let tasks = PinFutSlice::new(&mut tasks[..]).into();
783     let disks = gbl_ops.disks();
784     GblFastboot::new(gbl_ops, disks, Task::run, &tasks, &buffer_pool, bootimg_buf)
785         .run(usb, tcp)
786         .await;
787 }
788 
789 /// Pre-generates a Fuchsia Fastboot MDNS service broadcast packet.
790 ///
791 /// Fuchsia ffx development flow can detect fastboot devices that broadcast a "_fastboot·_tcp·local"
792 /// MDNS service. This API generates the broadcast MDNS packet for Ipv6. Caller is reponsible for
793 /// sending this packet via UDP at the following address and port (defined by MDNS):
794 ///
795 /// * ipv6: ff02::fb
796 /// * port: 5353
797 ///
798 /// # Args
799 ///
800 /// * `node_name`: The Fuchsia node name for the service. Must be a 22 character ASCII string in the
801 ///   format "fuchsia-xxxx-xxxx-xxxx".
802 /// * `ipv6_addr`: The Ipv6 address bytes.
803 ///
804 /// The packet generated by the API contains the given IPv6 address and a fuchsia node name derived
805 /// from the given ethernet mac address `eth_mac`.
fuchsia_fastboot_mdns_packet(node_name: &str, ipv6_addr: &[u8]) -> Result<[u8; 140], Error>806 pub fn fuchsia_fastboot_mdns_packet(node_name: &str, ipv6_addr: &[u8]) -> Result<[u8; 140], Error> {
807     // Pre-generated Fuchsia fastboot MDNS service packet template.
808     // It contains the node name and ipv6 address. We simply replace with the device's node name and
809     // ipv6 address.
810     let mut packet: [u8; 140] = [
811         0x00, 0x00, 0x84, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x09, 0x5f, 0x66,
812         0x61, 0x73, 0x74, 0x62, 0x6f, 0x6f, 0x74, 0x04, 0x5f, 0x74, 0x63, 0x70, 0x05, 0x6c, 0x6f,
813         0x63, 0x61, 0x6c, 0x00, 0x00, 0x0c, 0x80, 0x01, 0x00, 0x00, 0x00, 0x78, 0x00, 0x19, 0x16,
814         0x66, 0x75, 0x63, 0x68, 0x73, 0x69, 0x61, 0x2d, 0x34, 0x38, 0x32, 0x31, 0x2d, 0x30, 0x62,
815         0x33, 0x31, 0x2d, 0x65, 0x61, 0x66, 0x38, 0xc0, 0x0c, 0xc0, 0x2c, 0x00, 0x21, 0x80, 0x01,
816         0x00, 0x00, 0x00, 0x78, 0x00, 0x1f, 0x00, 0x00, 0x00, 0x00, 0x15, 0xb2, 0x16, 0x66, 0x75,
817         0x63, 0x68, 0x73, 0x69, 0x61, 0x2d, 0x34, 0x38, 0x32, 0x31, 0x2d, 0x30, 0x62, 0x33, 0x31,
818         0x2d, 0x65, 0x61, 0x66, 0x38, 0xc0, 0x1b, 0xc0, 0x57, 0x00, 0x1c, 0x80, 0x01, 0x00, 0x00,
819         0x00, 0x78, 0x00, 0x10, 0xfe, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0x21, 0x0b,
820         0xff, 0xfe, 0x31, 0xea, 0xf8,
821     ];
822     // Offsets to the fuchsia node name field.
823     const NODE_NAME_OFFSETS: &[usize; 2] = &[45, 88];
824     // Offset to the IPv6 address field.
825     const IP6_ADDR_OFFSET: usize = 124;
826 
827     if node_name.as_bytes().len() != 22 {
828         return Err(Error::InvalidInput);
829     }
830 
831     for off in NODE_NAME_OFFSETS {
832         packet[*off..][..node_name.len()].clone_from_slice(node_name.as_bytes());
833     }
834     packet[IP6_ADDR_OFFSET..][..ipv6_addr.len()].clone_from_slice(ipv6_addr);
835     Ok(packet)
836 }
837 
838 #[cfg(test)]
839 mod test {
840     use super::*;
841     use crate::{
842         ops::test::{FakeGblOps, FakeGblOpsStorage},
843         Os,
844     };
845     use abr::{
846         get_and_clear_one_shot_bootloader, get_boot_slot, mark_slot_unbootable, ABR_DATA_SIZE,
847     };
848     use core::{
849         mem::size_of,
850         pin::{pin, Pin},
851         str::from_utf8,
852     };
853     use fastboot::{test_utils::TestUploadBuilder, MAX_RESPONSE_SIZE};
854     use gbl_async::{block_on, poll, poll_n_times};
855     use gbl_storage::GPT_GUID_LEN;
856     use liberror::Error;
857     use spin::{Mutex, MutexGuard};
858     use std::ffi::CString;
859     use std::{collections::VecDeque, io::Read};
860     use zerocopy::AsBytes;
861 
862     /// A test implementation of [InfoSender] and [OkaySender].
863     #[derive(Default)]
864     struct TestResponder {
865         okay_sent: Mutex<bool>,
866         info_messages: Mutex<Vec<String>>,
867     }
868 
869     impl InfoSender for &TestResponder {
send_formatted_info<F: FnOnce(&mut dyn Write)>( &mut self, cb: F, ) -> Result<(), Error>870         async fn send_formatted_info<F: FnOnce(&mut dyn Write)>(
871             &mut self,
872             cb: F,
873         ) -> Result<(), Error> {
874             let mut msg: String = "".into();
875             cb(&mut msg);
876             self.info_messages.try_lock().unwrap().push(msg);
877             Ok(())
878         }
879     }
880 
881     impl OkaySender for &TestResponder {
882         /// Sends a Fastboot "INFO<`msg`>" packet.
send_formatted_okay<F: FnOnce(&mut dyn Write)>(self, _: F) -> Result<(), Error>883         async fn send_formatted_okay<F: FnOnce(&mut dyn Write)>(self, _: F) -> Result<(), Error> {
884             *self.okay_sent.try_lock().unwrap() = true;
885             Ok(())
886         }
887     }
888 
889     /// Helper to test fastboot variable value.
check_var(gbl_fb: &mut impl FastbootImplementation, var: &str, args: &str, expected: &str)890     fn check_var(gbl_fb: &mut impl FastbootImplementation, var: &str, args: &str, expected: &str) {
891         let resp: TestResponder = Default::default();
892         let args_c = args.split(':').map(|v| CString::new(v).unwrap()).collect::<Vec<_>>();
893         let args_c = args_c.iter().map(|v| v.as_c_str());
894         let var_c = CString::new(var).unwrap();
895         let mut out = vec![0u8; MAX_RESPONSE_SIZE];
896         let val =
897             block_on(gbl_fb.get_var_as_str(var_c.as_c_str(), args_c, &resp, &mut out[..])).unwrap();
898         assert_eq!(val, expected, "var {}:{} = {} != {}", var, args, val, expected,);
899     }
900 
901     /// A helper to set the download content.
set_download(gbl_fb: &mut impl FastbootImplementation, data: &[u8])902     fn set_download(gbl_fb: &mut impl FastbootImplementation, data: &[u8]) {
903         block_on(gbl_fb.get_download_buffer())[..data.len()].clone_from_slice(data);
904         block_on(gbl_fb.download_complete(data.len(), &TestResponder::default())).unwrap();
905     }
906 
907     impl<'a> PinFutContainer<'a> for Vec<Pin<Box<dyn Future<Output = ()> + 'a>>> {
add_with<F: Future<Output = ()> + 'a>(&mut self, f: impl FnOnce() -> F)908         fn add_with<F: Future<Output = ()> + 'a>(&mut self, f: impl FnOnce() -> F) {
909             self.push(Box::pin(f()));
910         }
911 
for_each_remove_if( &mut self, mut cb: impl FnMut(&mut Pin<&mut (dyn Future<Output = ()> + 'a)>) -> bool, )912         fn for_each_remove_if(
913             &mut self,
914             mut cb: impl FnMut(&mut Pin<&mut (dyn Future<Output = ()> + 'a)>) -> bool,
915         ) {
916             for idx in (0..self.len()).rev() {
917                 cb(&mut self[idx].as_mut()).then(|| self.swap_remove(idx));
918             }
919         }
920     }
921 
922     #[test]
test_get_var_gbl()923     fn test_get_var_gbl() {
924         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 1]);
925         let storage = FakeGblOpsStorage::default();
926         let mut gbl_ops = FakeGblOps::new(&storage);
927         let tasks = vec![].into();
928         let parts = gbl_ops.disks();
929         let mut gbl_fb =
930             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
931         check_var(
932             &mut gbl_fb,
933             FakeGblOps::GBL_TEST_VAR,
934             "arg",
935             format!("{}:Some(\"arg\")", FakeGblOps::GBL_TEST_VAR_VAL).as_str(),
936         );
937     }
938 
939     #[test]
test_get_var_partition_info()940     fn test_get_var_partition_info() {
941         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 1]);
942         let mut storage = FakeGblOpsStorage::default();
943         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_1.bin"));
944         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_2.bin"));
945         storage.add_raw_device(c"raw_0", [0xaau8; 4 * 1024]);
946         storage.add_raw_device(c"raw_1", [0x55u8; 8 * 1024]);
947         let mut gbl_ops = FakeGblOps::new(&storage);
948         let tasks = vec![].into();
949         let parts = gbl_ops.disks();
950         let mut gbl_fb =
951             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
952 
953         // Check different semantics
954         check_var(&mut gbl_fb, "partition-size", "boot_a", "0x2000");
955         check_var(&mut gbl_fb, "partition-size", "boot_a/", "0x2000");
956         check_var(&mut gbl_fb, "partition-size", "boot_a//", "0x2000");
957         check_var(&mut gbl_fb, "partition-size", "boot_a///", "0x2000");
958         check_var(&mut gbl_fb, "partition-size", "boot_a/0", "0x2000");
959         check_var(&mut gbl_fb, "partition-size", "boot_a/0/", "0x2000");
960         check_var(&mut gbl_fb, "partition-size", "boot_a//0", "0x2000");
961         check_var(&mut gbl_fb, "partition-size", "boot_a/0/0", "0x2000");
962         check_var(&mut gbl_fb, "partition-size", "boot_a//0x1000", "0x1000");
963 
964         check_var(&mut gbl_fb, "partition-size", "boot_b/0", "0x3000");
965         check_var(&mut gbl_fb, "partition-size", "vendor_boot_a/1", "0x1000");
966         check_var(&mut gbl_fb, "partition-size", "vendor_boot_b/1", "0x1800");
967         check_var(&mut gbl_fb, "partition-size", "boot_a//0x1000", "0x1000");
968         check_var(&mut gbl_fb, "partition-size", "raw_0", "0x1000");
969         check_var(&mut gbl_fb, "partition-size", "raw_1", "0x2000");
970 
971         let resp: TestResponder = Default::default();
972         let mut out = vec![0u8; MAX_RESPONSE_SIZE];
973         assert!(block_on(gbl_fb.get_var_as_str(
974             c"partition",
975             [c"non-existent"].into_iter(),
976             &resp,
977             &mut out[..],
978         ))
979         .is_err());
980     }
981 
982     /// `TestVarSender` implements `TestVarSender`. It stores outputs in a vector of string.
983     struct TestVarSender(Vec<String>);
984 
985     impl VarInfoSender for &mut TestVarSender {
send_var_info( &mut self, name: &str, args: impl IntoIterator<Item = &'_ str>, val: &str, ) -> Result<(), Error>986         async fn send_var_info(
987             &mut self,
988             name: &str,
989             args: impl IntoIterator<Item = &'_ str>,
990             val: &str,
991         ) -> Result<(), Error> {
992             let args = args.into_iter().collect::<Vec<_>>();
993             self.0.push(format!("{}:{}: {}", name, args.join(":"), val));
994             Ok(())
995         }
996     }
997 
998     #[test]
test_get_var_all()999     fn test_get_var_all() {
1000         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 1]);
1001         let mut storage = FakeGblOpsStorage::default();
1002         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_1.bin"));
1003         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_2.bin"));
1004         storage.add_raw_device(c"raw_0", [0xaau8; 4 * 1024]);
1005         storage.add_raw_device(c"raw_1", [0x55u8; 8 * 1024]);
1006         let mut gbl_ops = FakeGblOps::new(&storage);
1007         let tasks = vec![].into();
1008         let parts = gbl_ops.disks();
1009         let mut gbl_fb =
1010             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1011 
1012         let mut logger = TestVarSender(vec![]);
1013         block_on(gbl_fb.get_var_all(&mut logger)).unwrap();
1014         assert_eq!(
1015             logger.0,
1016             [
1017                 "version-bootloader:: 1.0",
1018                 "max-fetch-size:: 0xffffffffffffffff",
1019                 "block-device:0:total-blocks: 0x80",
1020                 "block-device:0:block-size: 0x200",
1021                 "block-device:0:status: idle",
1022                 "block-device:1:total-blocks: 0x100",
1023                 "block-device:1:block-size: 0x200",
1024                 "block-device:1:status: idle",
1025                 "block-device:2:total-blocks: 0x1000",
1026                 "block-device:2:block-size: 0x1",
1027                 "block-device:2:status: idle",
1028                 "block-device:3:total-blocks: 0x2000",
1029                 "block-device:3:block-size: 0x1",
1030                 "block-device:3:status: idle",
1031                 "gbl-default-block:: None",
1032                 "partition-size:boot_a/0: 0x2000",
1033                 "partition-type:boot_a/0: raw",
1034                 "partition-size:boot_b/0: 0x3000",
1035                 "partition-type:boot_b/0: raw",
1036                 "partition-size:vendor_boot_a/1: 0x1000",
1037                 "partition-type:vendor_boot_a/1: raw",
1038                 "partition-size:vendor_boot_b/1: 0x1800",
1039                 "partition-type:vendor_boot_b/1: raw",
1040                 "partition-size:raw_0/2: 0x1000",
1041                 "partition-type:raw_0/2: raw",
1042                 "partition-size:raw_1/3: 0x2000",
1043                 "partition-type:raw_1/3: raw",
1044                 format!("{}:1: {}:1", FakeGblOps::GBL_TEST_VAR, FakeGblOps::GBL_TEST_VAR_VAL)
1045                     .as_str(),
1046                 format!("{}:2: {}:2", FakeGblOps::GBL_TEST_VAR, FakeGblOps::GBL_TEST_VAR_VAL)
1047                     .as_str(),
1048             ]
1049         );
1050     }
1051 
1052     /// A helper for fetching partition from a `GblFastboot`
fetch<EOff: core::fmt::Debug, ESz: core::fmt::Debug>( fb: &mut impl FastbootImplementation, part: String, off: impl TryInto<u64, Error = EOff>, size: impl TryInto<u64, Error = ESz>, ) -> CommandResult<Vec<u8>>1053     fn fetch<EOff: core::fmt::Debug, ESz: core::fmt::Debug>(
1054         fb: &mut impl FastbootImplementation,
1055         part: String,
1056         off: impl TryInto<u64, Error = EOff>,
1057         size: impl TryInto<u64, Error = ESz>,
1058     ) -> CommandResult<Vec<u8>> {
1059         let off = off.try_into().unwrap();
1060         let size = size.try_into().unwrap();
1061         let mut upload_out = vec![0u8; usize::try_from(size).unwrap()];
1062         let test_uploader = TestUploadBuilder(&mut upload_out[..]);
1063         block_on(fb.fetch(part.as_str(), off, size, test_uploader))?;
1064         Ok(upload_out)
1065     }
1066 
1067     #[test]
test_fetch_invalid_partition_arg()1068     fn test_fetch_invalid_partition_arg() {
1069         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 1]);
1070         let mut storage = FakeGblOpsStorage::default();
1071         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_1.bin"));
1072         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_2.bin"));
1073         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_2.bin"));
1074         let mut gbl_ops = FakeGblOps::new(&storage);
1075         let tasks = vec![].into();
1076         let parts = gbl_ops.disks();
1077         let mut gbl_fb =
1078             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1079 
1080         // Missing mandatory block device ID for raw block partition.
1081         assert!(fetch(&mut gbl_fb, "//0/0".into(), 0, 0).is_err());
1082 
1083         // GPT partition does not exist.
1084         assert!(fetch(&mut gbl_fb, "non///".into(), 0, 0).is_err());
1085 
1086         // GPT Partition is not unique.
1087         assert!(fetch(&mut gbl_fb, "vendor_boot_a///".into(), 0, 0).is_err());
1088 
1089         // Offset overflows.
1090         assert!(fetch(&mut gbl_fb, "boot_a//0x2001/".into(), 0, 1).is_err());
1091         assert!(fetch(&mut gbl_fb, "boot_a".into(), 0x2000, 1).is_err());
1092 
1093         // Size overflows.
1094         assert!(fetch(&mut gbl_fb, "boot_a///0x2001".into(), 0, 0).is_err());
1095         assert!(fetch(&mut gbl_fb, "boot_a".into(), 0, 0x2001).is_err());
1096     }
1097 
1098     /// A helper for testing raw block upload. It verifies that data read from block device
1099     /// `blk_id` in range [`off`, `off`+`size`) is the same as `disk[off..][..size]`
check_blk_upload( fb: &mut impl FastbootImplementation, blk_id: u64, off: u64, size: u64, disk: &[u8], )1100     fn check_blk_upload(
1101         fb: &mut impl FastbootImplementation,
1102         blk_id: u64,
1103         off: u64,
1104         size: u64,
1105         disk: &[u8],
1106     ) {
1107         let expected = disk[off.try_into().unwrap()..][..size.try_into().unwrap()].to_vec();
1108         // offset/size as part of the partition string.
1109         let part = format!("/{:#x}/{:#x}/{:#x}", blk_id, off, size);
1110         assert_eq!(fetch(fb, part, 0, size).unwrap(), expected);
1111         // offset/size as separate fetch arguments.
1112         let part = format!("/{:#x}", blk_id);
1113         assert_eq!(fetch(fb, part, off, size).unwrap(), expected);
1114     }
1115 
1116     #[test]
test_fetch_raw_block()1117     fn test_fetch_raw_block() {
1118         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 1]);
1119         let mut storage = FakeGblOpsStorage::default();
1120         let disk_0 = include_bytes!("../../../libstorage/test/gpt_test_1.bin");
1121         let disk_1 = include_bytes!("../../../libstorage/test/gpt_test_2.bin");
1122         storage.add_gpt_device(disk_0);
1123         storage.add_gpt_device(disk_1);
1124         let mut gbl_ops = FakeGblOps::new(&storage);
1125         let tasks = vec![].into();
1126         let parts = gbl_ops.disks();
1127         let mut gbl_fb =
1128             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1129 
1130         let off = 512;
1131         let size = 512;
1132         check_blk_upload(&mut gbl_fb, 0, off, size, disk_0);
1133         check_blk_upload(&mut gbl_fb, 1, off, size, disk_1);
1134     }
1135 
1136     /// A helper for testing uploading GPT partition. It verifies that data read from GPT partition
1137     /// `part` at disk `blk_id` in range [`off`, `off`+`size`) is the same as
1138     /// `partition_data[off..][..size]`.
check_part_upload( fb: &mut impl FastbootImplementation, part: &str, off: u64, size: u64, blk_id: Option<u64>, partition_data: &[u8], )1139     fn check_part_upload(
1140         fb: &mut impl FastbootImplementation,
1141         part: &str,
1142         off: u64,
1143         size: u64,
1144         blk_id: Option<u64>,
1145         partition_data: &[u8],
1146     ) {
1147         let expected =
1148             partition_data[off.try_into().unwrap()..][..size.try_into().unwrap()].to_vec();
1149         let blk_id = blk_id.map_or("".to_string(), |v| format!("{:#x}", v));
1150         // offset/size as part of the partition string.
1151         let gpt_part = format!("{}/{}/{:#x}/{:#x}", part, blk_id, off, size);
1152         assert_eq!(fetch(fb, gpt_part, 0, size).unwrap(), expected);
1153         // offset/size as separate fetch arguments.
1154         let gpt_part = format!("{}/{}", part, blk_id);
1155         assert_eq!(fetch(fb, gpt_part, off, size).unwrap(), expected);
1156     }
1157 
1158     #[test]
test_fetch_partition()1159     fn test_fetch_partition() {
1160         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 1]);
1161         let mut storage = FakeGblOpsStorage::default();
1162         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_1.bin"));
1163         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_2.bin"));
1164         storage.add_raw_device(c"raw_0", [0xaau8; 4 * 1024]);
1165         storage.add_raw_device(c"raw_1", [0x55u8; 8 * 1024]);
1166         let mut gbl_ops = FakeGblOps::new(&storage);
1167         let tasks = vec![].into();
1168         let parts = gbl_ops.disks();
1169         let mut gbl_fb =
1170             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1171 
1172         let expect_boot_a = include_bytes!("../../../libstorage/test/boot_a.bin");
1173         let expect_boot_b = include_bytes!("../../../libstorage/test/boot_b.bin");
1174         let expect_vendor_boot_a = include_bytes!("../../../libstorage/test/vendor_boot_a.bin");
1175         let expect_vendor_boot_b = include_bytes!("../../../libstorage/test/vendor_boot_b.bin");
1176 
1177         let size = 512;
1178         let off = 512;
1179 
1180         check_part_upload(&mut gbl_fb, "boot_a", off, size, Some(0), expect_boot_a);
1181         check_part_upload(&mut gbl_fb, "boot_b", off, size, Some(0), expect_boot_b);
1182         check_part_upload(&mut gbl_fb, "vendor_boot_a", off, size, Some(1), expect_vendor_boot_a);
1183         check_part_upload(&mut gbl_fb, "vendor_boot_b", off, size, Some(1), expect_vendor_boot_b);
1184         check_part_upload(&mut gbl_fb, "raw_0", off, size, Some(2), &[0xaau8; 4 * 1024]);
1185         check_part_upload(&mut gbl_fb, "raw_1", off, size, Some(3), &[0x55u8; 8 * 1024]);
1186 
1187         // No block device id
1188         check_part_upload(&mut gbl_fb, "boot_a", off, size, None, expect_boot_a);
1189         check_part_upload(&mut gbl_fb, "boot_b", off, size, None, expect_boot_b);
1190         check_part_upload(&mut gbl_fb, "vendor_boot_a", off, size, None, expect_vendor_boot_a);
1191         check_part_upload(&mut gbl_fb, "vendor_boot_b", off, size, None, expect_vendor_boot_b);
1192         check_part_upload(&mut gbl_fb, "raw_0", off, size, None, &[0xaau8; 4 * 1024]);
1193         check_part_upload(&mut gbl_fb, "raw_1", off, size, None, &[0x55u8; 8 * 1024]);
1194     }
1195 
1196     /// A helper function to get a bit-flipped copy of the input data.
flipped_bits(data: &[u8]) -> Vec<u8>1197     fn flipped_bits(data: &[u8]) -> Vec<u8> {
1198         data.iter().map(|v| !(*v)).collect::<Vec<_>>()
1199     }
1200 
1201     /// A helper function to flash data to a partition
flash_part(fb: &mut impl FastbootImplementation, part: &str, data: &[u8])1202     fn flash_part(fb: &mut impl FastbootImplementation, part: &str, data: &[u8]) {
1203         // Prepare a download buffer.
1204         let dl_size = data.len();
1205         let download = data.to_vec();
1206         let resp: TestResponder = Default::default();
1207         set_download(fb, &download[..]);
1208         block_on(fb.flash(part, &resp)).unwrap();
1209         assert_eq!(fetch(fb, part.into(), 0, dl_size).unwrap(), download);
1210     }
1211 
1212     /// A helper for testing partition flashing.
check_flash_part(fb: &mut impl FastbootImplementation, part: &str, expected: &[u8])1213     fn check_flash_part(fb: &mut impl FastbootImplementation, part: &str, expected: &[u8]) {
1214         flash_part(fb, part, expected);
1215         // Also flashes bit-wise reversed version in case the initial content is the same.
1216         flash_part(fb, part, &flipped_bits(expected));
1217     }
1218 
1219     #[test]
test_flash_partition()1220     fn test_flash_partition() {
1221         let disk_0 = include_bytes!("../../../libstorage/test/gpt_test_1.bin");
1222         let disk_1 = include_bytes!("../../../libstorage/test/gpt_test_2.bin");
1223         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 1]);
1224         let mut storage = FakeGblOpsStorage::default();
1225         storage.add_gpt_device(disk_0);
1226         storage.add_gpt_device(disk_1);
1227         storage.add_raw_device(c"raw_0", [0xaau8; 4 * 1024]);
1228         storage.add_raw_device(c"raw_1", [0x55u8; 8 * 1024]);
1229         let mut gbl_ops = FakeGblOps::new(&storage);
1230         let tasks = vec![].into();
1231         let parts = gbl_ops.disks();
1232         let mut gbl_fb =
1233             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1234 
1235         let expect_boot_a = include_bytes!("../../../libstorage/test/boot_a.bin");
1236         let expect_boot_b = include_bytes!("../../../libstorage/test/boot_b.bin");
1237         check_flash_part(&mut gbl_fb, "boot_a", expect_boot_a);
1238         check_flash_part(&mut gbl_fb, "boot_b", expect_boot_b);
1239         check_flash_part(&mut gbl_fb, "raw_0", &[0xaau8; 4 * 1024]);
1240         check_flash_part(&mut gbl_fb, "raw_1", &[0x55u8; 8 * 1024]);
1241         check_flash_part(&mut gbl_fb, "/0", disk_0);
1242         check_flash_part(&mut gbl_fb, "/1", disk_1);
1243 
1244         // Partital flash
1245         let off = 0x200;
1246         let size = 1024;
1247         check_flash_part(&mut gbl_fb, "boot_a//200", &expect_boot_a[off..size]);
1248         check_flash_part(&mut gbl_fb, "boot_b//200", &expect_boot_b[off..size]);
1249         check_flash_part(&mut gbl_fb, "/0/200", &disk_0[off..size]);
1250         check_flash_part(&mut gbl_fb, "/1/200", &disk_1[off..size]);
1251     }
1252 
1253     #[test]
test_flash_partition_sparse()1254     fn test_flash_partition_sparse() {
1255         let raw = include_bytes!("../../testdata/sparse_test_raw.bin");
1256         let sparse = include_bytes!("../../testdata/sparse_test.bin");
1257         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 1]);
1258         let mut storage = FakeGblOpsStorage::default();
1259         storage.add_raw_device(c"raw", vec![0u8; raw.len()]);
1260         let mut gbl_ops = FakeGblOps::new(&storage);
1261         let tasks = vec![].into();
1262         let parts = gbl_ops.disks();
1263         let mut gbl_fb =
1264             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1265 
1266         let download = sparse.to_vec();
1267         let resp: TestResponder = Default::default();
1268         set_download(&mut gbl_fb, &download[..]);
1269         block_on(gbl_fb.flash("/0", &resp)).unwrap();
1270         assert_eq!(fetch(&mut gbl_fb, "/0".into(), 0, raw.len()).unwrap(), raw);
1271     }
1272 
1273     /// A helper to invoke OEM commands.
1274     ///
1275     /// Returns the result and INFO strings.
oem( fb: &mut impl FastbootImplementation, oem_cmd: &str, resp: impl InfoSender, ) -> CommandResult<String>1276     async fn oem(
1277         fb: &mut impl FastbootImplementation,
1278         oem_cmd: &str,
1279         resp: impl InfoSender,
1280     ) -> CommandResult<String> {
1281         let mut res = [0u8; MAX_RESPONSE_SIZE];
1282         fb.oem(oem_cmd, resp, &mut res[..]).await?;
1283         Ok(from_utf8(&mut res[..]).unwrap().into())
1284     }
1285 
1286     #[test]
test_async_flash()1287     fn test_async_flash() {
1288         // Creates two block devices for writing raw and sparse image.
1289         let sparse_raw = include_bytes!("../../testdata/sparse_test_raw.bin");
1290         let sparse = include_bytes!("../../testdata/sparse_test.bin");
1291         let mut storage = FakeGblOpsStorage::default();
1292         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_1.bin"));
1293         storage.add_gpt_device(vec![0u8; sparse_raw.len() + 67 * 512]);
1294         let mut gpt_builder = storage[1].gpt_builder().unwrap();
1295         gpt_builder.add("sparse", [1u8; GPT_GUID_LEN], [1u8; GPT_GUID_LEN], 0, None).unwrap();
1296         block_on(gpt_builder.persist()).unwrap();
1297         let mut gbl_ops = FakeGblOps::new(&storage);
1298         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 2]);
1299         let tasks = vec![].into();
1300         let parts = gbl_ops.disks();
1301         let mut gbl_fb =
1302             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1303         let tasks = gbl_fb.tasks();
1304         let resp: TestResponder = Default::default();
1305 
1306         // "oem gbl-sync-tasks" should return immediately when there is no pending IOs.
1307         assert!(poll(&mut pin!(oem(&mut gbl_fb, "gbl-sync-tasks", &resp))).unwrap().is_ok());
1308         // Enable async IO.
1309         assert!(poll(&mut pin!(oem(&mut gbl_fb, "gbl-enable-async-task", &resp))).unwrap().is_ok());
1310 
1311         // Flashes "boot_a".
1312         let expect_boot_a = flipped_bits(include_bytes!("../../../libstorage/test/boot_a.bin"));
1313         set_download(&mut gbl_fb, expect_boot_a.as_slice());
1314         block_on(gbl_fb.flash("boot_a", &resp)).unwrap();
1315         check_var(&mut gbl_fb, "block-device", "0:status", "IO pending");
1316 
1317         // Flashes the "sparse" partition on the different block device.
1318         set_download(&mut gbl_fb, sparse);
1319         block_on(gbl_fb.flash("sparse", &resp)).unwrap();
1320         check_var(&mut gbl_fb, "block-device", "1:status", "IO pending");
1321 
1322         {
1323             // "oem gbl-sync-tasks" should block.
1324             let oem_sync_blk_fut = &mut pin!(oem(&mut gbl_fb, "gbl-sync-tasks", &resp));
1325             assert!(poll(oem_sync_blk_fut).is_none());
1326             // Schedules the disk IO tasks to completion.
1327             tasks.borrow_mut().run();
1328             // "oem gbl-sync-tasks" should now be able to finish.
1329             assert!(poll(oem_sync_blk_fut).unwrap().is_ok());
1330         }
1331 
1332         // The two blocks should be in the idle state.
1333         check_var(&mut gbl_fb, "block-device", "0:status", "idle");
1334         check_var(&mut gbl_fb, "block-device", "1:status", "idle");
1335 
1336         // Verifies flashed image.
1337         assert_eq!(
1338             fetch(&mut gbl_fb, "boot_a".into(), 0, expect_boot_a.len()).unwrap(),
1339             expect_boot_a
1340         );
1341         assert_eq!(fetch(&mut gbl_fb, "sparse".into(), 0, sparse_raw.len()).unwrap(), sparse_raw);
1342     }
1343 
1344     #[test]
test_async_flash_block_on_busy_blk()1345     fn test_async_flash_block_on_busy_blk() {
1346         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 2]);
1347         let mut storage = FakeGblOpsStorage::default();
1348         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_1.bin"));
1349         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_2.bin"));
1350         let mut gbl_ops = FakeGblOps::new(&storage);
1351         let tasks = vec![].into();
1352         let parts = gbl_ops.disks();
1353         let mut gbl_fb =
1354             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1355         let tasks = gbl_fb.tasks();
1356         let resp: TestResponder = Default::default();
1357 
1358         // Enable async IO.
1359         assert!(poll(&mut pin!(oem(&mut gbl_fb, "gbl-enable-async-task", &resp))).unwrap().is_ok());
1360 
1361         // Flashes boot_a partition.
1362         let expect_boot_a = flipped_bits(include_bytes!("../../../libstorage/test/boot_a.bin"));
1363         set_download(&mut gbl_fb, expect_boot_a.as_slice());
1364         block_on(gbl_fb.flash("boot_a", &resp)).unwrap();
1365 
1366         // Flashes boot_b partition.
1367         let expect_boot_b = flipped_bits(include_bytes!("../../../libstorage/test/boot_b.bin"));
1368         set_download(&mut gbl_fb, expect_boot_b.as_slice());
1369         {
1370             let flash_boot_b_fut = &mut pin!(gbl_fb.flash("boot_b", &resp));
1371             // Previous IO has not completed. Block is busy.
1372             assert!(poll(flash_boot_b_fut).is_none());
1373             // There should only be the previous disk IO task for "boot_a".
1374             assert_eq!(tasks.borrow_mut().size(), 1);
1375             // Schedule the disk IO task for "flash boot_a" to completion.
1376             tasks.borrow_mut().run();
1377             // The blocked "flash boot_b" should now be able to finish.
1378             assert!(poll(flash_boot_b_fut).is_some());
1379             // There should be a disk IO task spawned for "flash boot_b".
1380             assert_eq!(tasks.borrow_mut().size(), 1);
1381             // Schedule the disk IO tasks for "flash boot_b" to completion.
1382             tasks.borrow_mut().run();
1383         }
1384 
1385         // Verifies flashed image.
1386         assert_eq!(
1387             fetch(&mut gbl_fb, "boot_a".into(), 0, expect_boot_a.len()).unwrap(),
1388             expect_boot_a
1389         );
1390         assert_eq!(
1391             fetch(&mut gbl_fb, "boot_b".into(), 0, expect_boot_b.len()).unwrap(),
1392             expect_boot_b
1393         );
1394     }
1395 
1396     #[test]
test_async_flash_error()1397     fn test_async_flash_error() {
1398         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 2]);
1399         let mut storage = FakeGblOpsStorage::default();
1400         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_1.bin"));
1401         let mut gbl_ops = FakeGblOps::new(&storage);
1402         // Injects an error.
1403         storage[0].partition_io(None).unwrap().dev().io().error =
1404             liberror::Error::Other(None).into();
1405         let tasks = vec![].into();
1406         let parts = gbl_ops.disks();
1407         let mut gbl_fb =
1408             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1409         let tasks = gbl_fb.tasks();
1410         let resp: TestResponder = Default::default();
1411 
1412         // Enable async IO.
1413         assert!(poll(&mut pin!(oem(&mut gbl_fb, "gbl-enable-async-task", &resp))).unwrap().is_ok());
1414         // Flashes boot_a partition.
1415         let expect_boot_a = flipped_bits(include_bytes!("../../../libstorage/test/boot_a.bin"));
1416         set_download(&mut gbl_fb, expect_boot_a.as_slice());
1417         block_on(gbl_fb.flash("boot_a", &resp)).unwrap();
1418         // Schedules the disk IO tasks to completion.
1419         tasks.borrow_mut().run();
1420         // New flash to "boot_a" should fail due to previous error
1421         set_download(&mut gbl_fb, expect_boot_a.as_slice());
1422         assert!(block_on(gbl_fb.flash("boot_a", &resp)).is_err());
1423         // "oem gbl-sync-tasks" should fail.
1424         assert!(block_on(oem(&mut gbl_fb, "gbl-sync-tasks", &resp)).is_err());
1425     }
1426 
1427     #[test]
test_async_erase()1428     fn test_async_erase() {
1429         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 2]);
1430         let mut storage = FakeGblOpsStorage::default();
1431         storage.add_raw_device(c"raw_0", [0xaau8; 4096]);
1432         storage.add_raw_device(c"raw_1", [0x55u8; 4096]);
1433         let mut gbl_ops = FakeGblOps::new(&storage);
1434         let tasks = vec![].into();
1435         let parts = gbl_ops.disks();
1436         let mut gbl_fb =
1437             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1438         let tasks = gbl_fb.tasks();
1439         let resp: TestResponder = Default::default();
1440 
1441         // Enable async IO.
1442         assert!(poll(&mut pin!(oem(&mut gbl_fb, "gbl-enable-async-task", &resp))).unwrap().is_ok());
1443 
1444         // Erases "raw_0".
1445         block_on(gbl_fb.erase("raw_0", &resp)).unwrap();
1446         check_var(&mut gbl_fb, "block-device", "0:status", "IO pending");
1447 
1448         // Erases second half of "raw_1"
1449         block_on(gbl_fb.erase("raw_1//800", &resp)).unwrap();
1450         check_var(&mut gbl_fb, "block-device", "1:status", "IO pending");
1451 
1452         {
1453             // "oem gbl-sync-tasks" should block.
1454             let oem_sync_blk_fut = &mut pin!(oem(&mut gbl_fb, "gbl-sync-tasks", &resp));
1455             assert!(poll(oem_sync_blk_fut).is_none());
1456             // Schedules the disk IO tasks to completion.
1457             tasks.borrow_mut().run();
1458             // "oem gbl-sync-tasks" should now be able to finish.
1459             assert!(poll(oem_sync_blk_fut).unwrap().is_ok());
1460         }
1461 
1462         // The two blocks should be in the idle state.
1463         check_var(&mut gbl_fb, "block-device", "0:status", "idle");
1464         check_var(&mut gbl_fb, "block-device", "1:status", "idle");
1465 
1466         assert_eq!(storage[0].partition_io(None).unwrap().dev().io().storage, [0u8; 4096]);
1467         assert_eq!(
1468             storage[1].partition_io(None).unwrap().dev().io().storage,
1469             [[0x55u8; 2048], [0u8; 2048]].concat()
1470         );
1471     }
1472 
1473     #[test]
test_default_block()1474     fn test_default_block() {
1475         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 1]);
1476         let mut storage = FakeGblOpsStorage::default();
1477         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_1.bin"));
1478         let disk_dup = include_bytes!("../../../libstorage/test/gpt_test_2.bin");
1479         storage.add_gpt_device(disk_dup);
1480         storage.add_gpt_device(disk_dup);
1481         let raw_a = [0xaau8; 4 * 1024];
1482         let raw_b = [0x55u8; 8 * 1024];
1483         storage.add_raw_device(c"raw", raw_a);
1484         storage.add_raw_device(c"raw", raw_b);
1485         let mut gbl_ops = FakeGblOps::new(&storage);
1486         let tasks = vec![].into();
1487         let parts = gbl_ops.disks();
1488         let mut gbl_fb =
1489             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1490         let resp: TestResponder = Default::default();
1491 
1492         let boot_a = include_bytes!("../../../libstorage/test/boot_a.bin");
1493         // Flips the bits on partition "vendor_boot_a" on block device #2 to make it different from
1494         // block #1.
1495         let vendor_boot_a =
1496             flipped_bits(include_bytes!("../../../libstorage/test/vendor_boot_a.bin"));
1497         flash_part(&mut gbl_fb, "vendor_boot_a/2", &vendor_boot_a);
1498 
1499         let size = 512;
1500         let off = 512;
1501 
1502         check_var(&mut gbl_fb, "gbl-default-block", "", "None");
1503         // Sets default block to #2
1504         block_on(oem(&mut gbl_fb, "gbl-set-default-block 2", &resp)).unwrap();
1505         check_var(&mut gbl_fb, "gbl-default-block", "", "0x2");
1506         // The following fetch should succeed and fetch from "vendor_boot_a" on block 2.
1507         check_part_upload(&mut gbl_fb, "vendor_boot_a", off, size, None, &vendor_boot_a);
1508 
1509         // Sets default block to #4 (raw_b)
1510         block_on(oem(&mut gbl_fb, "gbl-set-default-block 4", &resp)).unwrap();
1511         check_var(&mut gbl_fb, "gbl-default-block", "", "0x4");
1512         // The following fetch should succeed and fetch from "raw" on block 4.
1513         check_part_upload(&mut gbl_fb, "raw", off, size, None, &raw_b);
1514 
1515         // Fetches with explicit storage ID shouldn't be affected.
1516         check_part_upload(&mut gbl_fb, "boot_a", off, size, Some(0), boot_a);
1517         check_part_upload(&mut gbl_fb, "raw", off, size, Some(3), &raw_a);
1518         check_blk_upload(&mut gbl_fb, 1, off, size, disk_dup);
1519 
1520         // Fetching without storage ID should use default ID and thus the following should fail.
1521         assert!(fetch(&mut gbl_fb, "boot_a".into(), 0, boot_a.len()).is_err());
1522 
1523         // Sets default block to #1 (unmodified `disk_dup`)
1524         block_on(oem(&mut gbl_fb, "gbl-set-default-block 1", &resp)).unwrap();
1525         check_var(&mut gbl_fb, "gbl-default-block", "", "0x1");
1526         // Fetches whole raw block but without block ID should use the default block.
1527         check_part_upload(&mut gbl_fb, "", off, size, None, disk_dup);
1528 
1529         // Unset default block
1530         block_on(oem(&mut gbl_fb, "gbl-unset-default-block", &resp)).unwrap();
1531         check_var(&mut gbl_fb, "gbl-default-block", "", "None");
1532         // Fetching non-unique partitions should now fail.
1533         assert!(fetch(&mut gbl_fb, "raw".into(), 0, raw_a.len()).is_err());
1534         assert!(fetch(&mut gbl_fb, "vendor_boot_a".into(), 0, vendor_boot_a.len()).is_err());
1535         assert!(fetch(&mut gbl_fb, "/".into(), 0, 512).is_err());
1536     }
1537 
1538     #[test]
test_set_default_block_invalid_arg()1539     fn test_set_default_block_invalid_arg() {
1540         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 2]);
1541         let storage = FakeGblOpsStorage::default();
1542         let mut gbl_ops = FakeGblOps::new(&storage);
1543         let tasks = vec![].into();
1544         let parts = gbl_ops.disks();
1545         let mut gbl_fb =
1546             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1547         let resp: TestResponder = Default::default();
1548         // Missing block device ID.
1549         assert!(block_on(oem(&mut gbl_fb, "gbl-set-default-block ", &resp)).is_err());
1550         // Invalid block device ID.
1551         assert!(block_on(oem(&mut gbl_fb, "gbl-set-default-block zzz", &resp)).is_err());
1552         // Out of range block device ID. (We've added no block device).
1553         assert!(block_on(oem(&mut gbl_fb, "gbl-set-default-block 0", &resp)).is_err());
1554     }
1555 
1556     #[test]
test_reboot_sync_all_blocks()1557     fn test_reboot_sync_all_blocks() {
1558         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 2]);
1559         let mut storage = FakeGblOpsStorage::default();
1560         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_1.bin"));
1561         let mut gbl_ops = FakeGblOps::new(&storage);
1562         let tasks = vec![].into();
1563         let parts = gbl_ops.disks();
1564         let mut gbl_fb =
1565             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1566         let tasks = gbl_fb.tasks();
1567         let resp: TestResponder = Default::default();
1568 
1569         block_on(oem(&mut gbl_fb, "gbl-enable-async-task", &resp)).unwrap();
1570 
1571         // Flashes "boot_a".
1572         let expect_boot_a = flipped_bits(include_bytes!("../../../libstorage/test/boot_a.bin"));
1573         set_download(&mut gbl_fb, expect_boot_a.as_slice());
1574         block_on(gbl_fb.flash("boot_a", &resp)).unwrap();
1575         // Checks initial state, okay_sent=false.
1576         assert!(!(*resp.okay_sent.try_lock().unwrap()));
1577         // Performs a reboot.
1578         let mut reboot_fut = pin!(gbl_fb.reboot(RebootMode::Normal, &resp));
1579         // There is a pending flash task. Reboot should wait.
1580         assert!(poll(&mut reboot_fut).is_none());
1581         assert!(!(*resp.okay_sent.try_lock().unwrap()));
1582         assert_eq!(resp.info_messages.try_lock().unwrap()[1], "Syncing storage...");
1583         // Schedules the disk IO tasks to completion.
1584         tasks.borrow_mut().run();
1585         // The reboot can now complete.
1586         assert!(poll(&mut reboot_fut).is_some());
1587         assert!((*resp.okay_sent.try_lock().unwrap()));
1588         assert_eq!(resp.info_messages.try_lock().unwrap()[2], "Rebooting...");
1589     }
1590 
1591     #[test]
test_continue_sync_all_blocks()1592     fn test_continue_sync_all_blocks() {
1593         let dl_buffers = Shared::from(vec![vec![0u8; 128 * 1024]; 2]);
1594         let mut storage = FakeGblOpsStorage::default();
1595         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_1.bin"));
1596         let mut gbl_ops = FakeGblOps::new(&storage);
1597         let tasks = vec![].into();
1598         let parts = gbl_ops.disks();
1599         let mut gbl_fb =
1600             GblFastboot::new(&mut gbl_ops, parts, Task::run, &tasks, &dl_buffers, &mut []);
1601         let tasks = gbl_fb.tasks();
1602         let resp: TestResponder = Default::default();
1603 
1604         block_on(oem(&mut gbl_fb, "gbl-enable-async-task", &resp)).unwrap();
1605 
1606         // Flashes "boot_a".
1607         let expect_boot_a = flipped_bits(include_bytes!("../../../libstorage/test/boot_a.bin"));
1608         set_download(&mut gbl_fb, expect_boot_a.as_slice());
1609         block_on(gbl_fb.flash("boot_a", &resp)).unwrap();
1610         // Performs a continue.
1611         let mut continue_fut = pin!(gbl_fb.r#continue(&resp));
1612         // There is a pending flash task. Continue should wait.
1613         assert!(poll(&mut continue_fut).is_none());
1614         assert!(!(*resp.okay_sent.try_lock().unwrap()));
1615         assert_eq!(resp.info_messages.try_lock().unwrap()[1], "Syncing storage...");
1616         // Schedules the disk IO tasks to completion.
1617         tasks.borrow_mut().run();
1618         // The continue can now complete.
1619         assert!(poll(&mut continue_fut).is_some());
1620     }
1621 
1622     /// Generates a length prefixed byte sequence.
length_prefixed(data: &[u8]) -> Vec<u8>1623     fn length_prefixed(data: &[u8]) -> Vec<u8> {
1624         [&data.len().to_be_bytes()[..], data].concat()
1625     }
1626 
1627     /// Used for a test implementation of [GblUsbTransport] and [GblTcpStream].
1628     #[derive(Default)]
1629     struct TestListener {
1630         usb_in_queue: VecDeque<Vec<u8>>,
1631         usb_out_queue: VecDeque<Vec<u8>>,
1632 
1633         tcp_in_queue: VecDeque<u8>,
1634         tcp_out_queue: VecDeque<u8>,
1635     }
1636 
1637     /// A shared [TestListener].
1638     #[derive(Default)]
1639     struct SharedTestListener(Mutex<TestListener>);
1640 
1641     impl SharedTestListener {
1642         /// Locks the listener
lock(&self) -> MutexGuard<TestListener>1643         fn lock(&self) -> MutexGuard<TestListener> {
1644             self.0.try_lock().unwrap()
1645         }
1646 
1647         /// Adds packet to USB input
add_usb_input(&self, packet: &[u8])1648         fn add_usb_input(&self, packet: &[u8]) {
1649             self.lock().usb_in_queue.push_back(packet.into());
1650         }
1651 
1652         /// Adds bytes to input stream.
add_tcp_input(&self, data: &[u8])1653         fn add_tcp_input(&self, data: &[u8]) {
1654             self.lock().tcp_in_queue.append(&mut data.to_vec().into());
1655         }
1656 
1657         /// Adds a length pre-fixed bytes stream.
add_tcp_length_prefixed_input(&self, data: &[u8])1658         fn add_tcp_length_prefixed_input(&self, data: &[u8]) {
1659             self.add_tcp_input(&length_prefixed(data));
1660         }
1661 
1662         /// Gets a copy of `Self::usb_out_queue`.
usb_out_queue(&self) -> VecDeque<Vec<u8>>1663         fn usb_out_queue(&self) -> VecDeque<Vec<u8>> {
1664             self.lock().usb_out_queue.clone()
1665         }
1666 
1667         /// Gets a copy of `Self::tcp_out_queue`.
tcp_out_queue(&self) -> VecDeque<u8>1668         fn tcp_out_queue(&self) -> VecDeque<u8> {
1669             self.lock().tcp_out_queue.clone()
1670         }
1671 
1672         /// A helper for decoding USB output packets as a string
dump_usb_out_queue(&self) -> String1673         fn dump_usb_out_queue(&self) -> String {
1674             let mut res = String::from("");
1675             for v in self.lock().usb_out_queue.iter() {
1676                 let v = String::from_utf8(v.clone()).unwrap_or(format!("{:?}", v));
1677                 res += format!("b{:?},\n", v).as_str();
1678             }
1679             res
1680         }
1681 
1682         /// A helper for decoding TCP output data as strings
dump_tcp_out_queue(&self) -> String1683         fn dump_tcp_out_queue(&self) -> String {
1684             let mut data = self.lock();
1685             let mut v;
1686             let (_, mut remains) = data.tcp_out_queue.make_contiguous().split_at(4);
1687             let mut res = String::from("");
1688             while !remains.is_empty() {
1689                 // Parses length-prefixed payload.
1690                 let (len, rest) = remains.split_first_chunk::<{ size_of::<u64>() }>().unwrap();
1691                 (v, remains) = rest.split_at(u64::from_be_bytes(*len).try_into().unwrap());
1692                 let s = String::from_utf8(v.to_vec()).unwrap_or(format!("{:?}", v));
1693                 res += format!("b{:?},\n", s).as_str();
1694             }
1695             res
1696         }
1697     }
1698 
1699     impl Transport for &SharedTestListener {
receive_packet(&mut self, out: &mut [u8]) -> Result<usize, Error>1700         async fn receive_packet(&mut self, out: &mut [u8]) -> Result<usize, Error> {
1701             match self.lock().usb_in_queue.pop_front() {
1702                 Some(v) => Ok((&v[..]).read(out).unwrap()),
1703                 _ => Err(Error::Other(Some("No more data"))),
1704             }
1705         }
1706 
send_packet(&mut self, packet: &[u8]) -> Result<(), Error>1707         async fn send_packet(&mut self, packet: &[u8]) -> Result<(), Error> {
1708             Ok(self.lock().usb_out_queue.push_back(packet.into()))
1709         }
1710     }
1711 
1712     impl GblUsbTransport for &SharedTestListener {
has_packet(&mut self) -> bool1713         fn has_packet(&mut self) -> bool {
1714             !self.lock().usb_in_queue.is_empty()
1715         }
1716     }
1717 
1718     impl TcpStream for &SharedTestListener {
read_exact(&mut self, out: &mut [u8]) -> Result<(), Error>1719         async fn read_exact(&mut self, out: &mut [u8]) -> Result<(), Error> {
1720             match self.lock().tcp_in_queue.read(out).unwrap() == out.len() {
1721                 true => Ok(()),
1722                 _ => Err(Error::Other(Some("No more data"))),
1723             }
1724         }
1725 
write_exact(&mut self, data: &[u8]) -> Result<(), Error>1726         async fn write_exact(&mut self, data: &[u8]) -> Result<(), Error> {
1727             Ok(self.lock().tcp_out_queue.append(&mut data.to_vec().into()))
1728         }
1729     }
1730 
1731     impl GblTcpStream for &SharedTestListener {
accept_new(&mut self) -> bool1732         fn accept_new(&mut self) -> bool {
1733             !self.lock().tcp_in_queue.is_empty()
1734         }
1735     }
1736 
1737     /// A helper to make an expected stream of USB output.
make_expected_usb_out(data: &[&[u8]]) -> VecDeque<Vec<u8>>1738     fn make_expected_usb_out(data: &[&[u8]]) -> VecDeque<Vec<u8>> {
1739         VecDeque::from(data.iter().map(|v| v.to_vec()).collect::<Vec<_>>())
1740     }
1741 
1742     /// A helper to make an expected stream of TCP output.
make_expected_tcp_out(data: &[&[u8]]) -> VecDeque<u8>1743     fn make_expected_tcp_out(data: &[&[u8]]) -> VecDeque<u8> {
1744         let mut res = VecDeque::<u8>::from(b"FB01".to_vec());
1745         data.iter().for_each(|v| res.append(&mut length_prefixed(v).into()));
1746         res
1747     }
1748 
1749     #[test]
test_run_gbl_fastboot()1750     fn test_run_gbl_fastboot() {
1751         let storage = FakeGblOpsStorage::default();
1752         let buffers = vec![vec![0u8; 128 * 1024]; 2];
1753         let mut gbl_ops = FakeGblOps::new(&storage);
1754         let listener: SharedTestListener = Default::default();
1755         let (usb, tcp) = (&listener, &listener);
1756 
1757         listener.add_usb_input(b"getvar:version-bootloader");
1758         listener.add_tcp_input(b"FB01");
1759         listener.add_tcp_length_prefixed_input(b"getvar:max-download-size");
1760         listener.add_tcp_length_prefixed_input(b"continue");
1761         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
1762 
1763         assert_eq!(
1764             listener.usb_out_queue(),
1765             make_expected_usb_out(&[b"OKAY1.0"]),
1766             "\nActual USB output:\n{}",
1767             listener.dump_usb_out_queue()
1768         );
1769 
1770         assert_eq!(
1771             listener.tcp_out_queue(),
1772             make_expected_tcp_out(&[b"OKAY0x20000", b"INFOSyncing storage...", b"OKAY"]),
1773             "\nActual TCP output:\n{}",
1774             listener.dump_tcp_out_queue()
1775         );
1776     }
1777 
1778     #[test]
test_run_gbl_fastboot_parallel_task()1779     fn test_run_gbl_fastboot_parallel_task() {
1780         let mut storage = FakeGblOpsStorage::default();
1781         storage.add_raw_device(c"raw_0", [0u8; 4 * 1024]);
1782         storage.add_raw_device(c"raw_1", [0u8; 8 * 1024]);
1783         let buffers = vec![vec![0u8; 128 * 1024]; 2];
1784         let mut gbl_ops = FakeGblOps::new(&storage);
1785         let listener: SharedTestListener = Default::default();
1786         let (usb, tcp) = (&listener, &listener);
1787         let mut fb_fut =
1788             pin!(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
1789 
1790         listener.add_usb_input(b"oem gbl-enable-async-task");
1791         listener.add_usb_input(format!("download:{:#x}", 4 * 1024).as_bytes());
1792         listener.add_usb_input(&[0x55u8; 4 * 1024]);
1793         listener.add_usb_input(b"flash:raw_0");
1794 
1795         listener.add_tcp_input(b"FB01");
1796         listener.add_tcp_length_prefixed_input(format!("download:{:#x}", 8 * 1024).as_bytes());
1797         listener.add_tcp_length_prefixed_input(&[0xaau8; 8 * 1024]);
1798         listener.add_tcp_length_prefixed_input(b"flash:raw_1");
1799 
1800         assert!(poll_n_times(&mut fb_fut, 100).is_none());
1801 
1802         assert_eq!(
1803             listener.usb_out_queue(),
1804             make_expected_usb_out(&[
1805                 b"OKAY",
1806                 b"DATA00001000",
1807                 b"OKAY",
1808                 b"INFOAn IO task is launched. To sync manually, run \"oem gbl-sync-tasks\".",
1809                 b"OKAY",
1810             ]),
1811             "\nActual USB output:\n{}",
1812             listener.dump_usb_out_queue()
1813         );
1814 
1815         assert_eq!(
1816             listener.tcp_out_queue(),
1817             make_expected_tcp_out(&[
1818                 b"DATA00002000",
1819                 b"OKAY",
1820                 b"INFOAn IO task is launched. To sync manually, run \"oem gbl-sync-tasks\".",
1821                 b"OKAY",
1822             ]),
1823             "\nActual TCP output:\n{}",
1824             listener.dump_tcp_out_queue()
1825         );
1826 
1827         // Verifies flashed image on raw_0.
1828         assert_eq!(storage[0].partition_io(None).unwrap().dev().io().storage, [0x55u8; 4 * 1024]);
1829 
1830         // Verifies flashed image on raw_1.
1831         assert_eq!(storage[1].partition_io(None).unwrap().dev().io().storage, [0xaau8; 8 * 1024]);
1832     }
1833 
1834     #[test]
test_oem_add_staged_bootloader_file()1835     fn test_oem_add_staged_bootloader_file() {
1836         let storage = FakeGblOpsStorage::default();
1837         let buffers = vec![vec![0u8; 128 * 1024]; 2];
1838         let mut gbl_ops = FakeGblOps::new(&storage);
1839         gbl_ops.get_zbi_bootloader_files_buffer().unwrap().fill(0);
1840         let listener: SharedTestListener = Default::default();
1841         let (usb, tcp) = (&listener, &listener);
1842 
1843         // Stages two zbi files.
1844         listener.add_usb_input(format!("download:{:#x}", 3).as_bytes());
1845         listener.add_usb_input(b"foo");
1846         listener.add_usb_input(b"oem add-staged-bootloader-file file_1");
1847         listener.add_usb_input(format!("download:{:#x}", 3).as_bytes());
1848         listener.add_usb_input(b"bar");
1849         listener.add_usb_input(b"oem add-staged-bootloader-file file_2");
1850         listener.add_usb_input(b"continue");
1851 
1852         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
1853 
1854         let buffer = gbl_ops.get_zbi_bootloader_files_buffer_aligned().unwrap();
1855         let container = ZbiContainer::parse(&buffer[..]).unwrap();
1856         let mut iter = container.iter();
1857         assert_eq!(iter.next().unwrap().payload.as_bytes(), b"\x06file_1foo");
1858         assert_eq!(iter.next().unwrap().payload.as_bytes(), b"\x06file_2bar");
1859         assert!(iter.next().is_none());
1860     }
1861 
1862     #[test]
test_oem_add_staged_bootloader_file_missing_file_name()1863     fn test_oem_add_staged_bootloader_file_missing_file_name() {
1864         let storage = FakeGblOpsStorage::default();
1865         let buffers = vec![vec![0u8; 128 * 1024]; 2];
1866         let mut gbl_ops = FakeGblOps::new(&storage);
1867         let listener: SharedTestListener = Default::default();
1868         let (usb, tcp) = (&listener, &listener);
1869 
1870         listener.add_usb_input(format!("download:{:#x}", 3).as_bytes());
1871         listener.add_usb_input(b"foo");
1872         listener.add_usb_input(b"oem add-staged-bootloader-file");
1873         listener.add_usb_input(b"continue");
1874 
1875         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
1876 
1877         assert_eq!(
1878             listener.usb_out_queue(),
1879             make_expected_usb_out(&[
1880                 b"DATA00000003",
1881                 b"OKAY",
1882                 b"FAILMissing file name",
1883                 b"INFOSyncing storage...",
1884                 b"OKAY",
1885             ]),
1886             "\nActual USB output:\n{}",
1887             listener.dump_usb_out_queue()
1888         )
1889     }
1890 
1891     #[test]
test_oem_add_staged_bootloader_file_missing_download()1892     fn test_oem_add_staged_bootloader_file_missing_download() {
1893         let storage = FakeGblOpsStorage::default();
1894         let buffers = vec![vec![0u8; 128 * 1024]; 2];
1895         let mut gbl_ops = FakeGblOps::new(&storage);
1896         let listener: SharedTestListener = Default::default();
1897         let (usb, tcp) = (&listener, &listener);
1898 
1899         listener.add_usb_input(b"oem add-staged-bootloader-file file1");
1900         listener.add_usb_input(b"continue");
1901 
1902         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
1903 
1904         assert_eq!(
1905             listener.usb_out_queue(),
1906             make_expected_usb_out(&[b"FAILNo file staged", b"INFOSyncing storage...", b"OKAY",]),
1907             "\nActual USB output:\n{}",
1908             listener.dump_usb_out_queue()
1909         );
1910     }
1911 
1912     #[test]
test_fuchsia_fastboot_mdns_packet()1913     fn test_fuchsia_fastboot_mdns_packet() {
1914         let expected = [
1915             0x00, 0x00, 0x84, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x09, 0x5f,
1916             0x66, 0x61, 0x73, 0x74, 0x62, 0x6f, 0x6f, 0x74, 0x04, 0x5f, 0x74, 0x63, 0x70, 0x05,
1917             0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x00, 0x00, 0x0c, 0x80, 0x01, 0x00, 0x00, 0x00, 0x78,
1918             0x00, 0x19, 0x16, 0x66, 0x75, 0x63, 0x68, 0x73, 0x69, 0x61, 0x2d, 0x35, 0x32, 0x35,
1919             0x34, 0x2d, 0x30, 0x30, 0x31, 0x32, 0x2d, 0x33, 0x34, 0x35, 0x36, 0xc0, 0x0c, 0xc0,
1920             0x2c, 0x00, 0x21, 0x80, 0x01, 0x00, 0x00, 0x00, 0x78, 0x00, 0x1f, 0x00, 0x00, 0x00,
1921             0x00, 0x15, 0xb2, 0x16, 0x66, 0x75, 0x63, 0x68, 0x73, 0x69, 0x61, 0x2d, 0x35, 0x32,
1922             0x35, 0x34, 0x2d, 0x30, 0x30, 0x31, 0x32, 0x2d, 0x33, 0x34, 0x35, 0x36, 0xc0, 0x1b,
1923             0xc0, 0x57, 0x00, 0x1c, 0x80, 0x01, 0x00, 0x00, 0x00, 0x78, 0x00, 0x10, 0xfe, 0x80,
1924             0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x50, 0x54, 0x00, 0xff, 0xfe, 0x12, 0x34, 0x56,
1925         ];
1926         let ip6_addr = &[
1927             0xfe, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x50, 0x54, 0x00, 0xff, 0xfe, 0x12,
1928             0x34, 0x56,
1929         ];
1930         assert_eq!(
1931             fuchsia_fastboot_mdns_packet("fuchsia-5254-0012-3456", ip6_addr).unwrap(),
1932             expected
1933         );
1934     }
1935 
1936     #[test]
test_fuchsia_fastboot_mdns_packet_invalid_node_name()1937     fn test_fuchsia_fastboot_mdns_packet_invalid_node_name() {
1938         let ip6_addr = &[
1939             0xfe, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x50, 0x54, 0x00, 0xff, 0xfe, 0x12,
1940             0x34, 0x56,
1941         ];
1942         assert!(fuchsia_fastboot_mdns_packet("fuchsia-5254-0012-345", ip6_addr).is_err());
1943         assert!(fuchsia_fastboot_mdns_packet("fuchsia-5254-0012-34567", ip6_addr).is_err());
1944     }
1945 
1946     #[test]
test_oem_update_gpt()1947     fn test_oem_update_gpt() {
1948         let disk_orig = include_bytes!("../../../libstorage/test/gpt_test_1.bin");
1949         // Erase the primary and secondary header.
1950         let mut disk = disk_orig.to_vec();
1951         disk[512..][..512].fill(0);
1952         disk.last_chunk_mut::<512>().unwrap().fill(0);
1953 
1954         let mut storage = FakeGblOpsStorage::default();
1955         storage.add_gpt_device(&disk);
1956         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_2.bin"));
1957         let buffers = vec![vec![0u8; 128 * 1024]; 2];
1958         let mut gbl_ops = FakeGblOps::new(&storage);
1959         let listener: SharedTestListener = Default::default();
1960         let (usb, tcp) = (&listener, &listener);
1961 
1962         // Checks that there is no valid partitions for block #0.
1963         listener.add_usb_input(b"getvar:partition-size:boot_a");
1964         listener.add_usb_input(b"getvar:partition-size:boot_b");
1965         // No partitions on block #0 should show up in `getvar:all` despite being a GPT device,
1966         // since the GPTs are corrupted.
1967         listener.add_usb_input(b"getvar:all");
1968         // Download a GPT
1969         let gpt = &disk_orig[..34 * 512];
1970         listener.add_usb_input(format!("download:{:#x}", gpt.len()).as_bytes());
1971         listener.add_usb_input(gpt);
1972         listener.add_usb_input(b"flash:gpt/0");
1973         // Checks that we can get partition info now.
1974         listener.add_usb_input(b"getvar:partition-size:boot_a");
1975         listener.add_usb_input(b"getvar:partition-size:boot_b");
1976         listener.add_usb_input(b"getvar:all");
1977 
1978         listener.add_usb_input(b"continue");
1979 
1980         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
1981 
1982         assert_eq!(
1983             listener.usb_out_queue(),
1984             make_expected_usb_out(&[
1985                 b"FAILNotFound",
1986                 b"FAILNotFound",
1987                 b"INFOmax-download-size: 0x20000",
1988                 b"INFOversion-bootloader: 1.0",
1989                 b"INFOmax-fetch-size: 0xffffffffffffffff",
1990                 b"INFOblock-device:0:total-blocks: 0x80",
1991                 b"INFOblock-device:0:block-size: 0x200",
1992                 b"INFOblock-device:0:status: idle",
1993                 b"INFOblock-device:1:total-blocks: 0x100",
1994                 b"INFOblock-device:1:block-size: 0x200",
1995                 b"INFOblock-device:1:status: idle",
1996                 b"INFOgbl-default-block: None",
1997                 b"INFOpartition-size:vendor_boot_a/1: 0x1000",
1998                 b"INFOpartition-type:vendor_boot_a/1: raw",
1999                 b"INFOpartition-size:vendor_boot_b/1: 0x1800",
2000                 b"INFOpartition-type:vendor_boot_b/1: raw",
2001                 format!("INFO{}:1: {}:1", FakeGblOps::GBL_TEST_VAR, FakeGblOps::GBL_TEST_VAR_VAL)
2002                     .as_bytes(),
2003                 format!("INFO{}:2: {}:2", FakeGblOps::GBL_TEST_VAR, FakeGblOps::GBL_TEST_VAR_VAL)
2004                     .as_bytes(),
2005                 b"OKAY",
2006                 b"DATA00004400",
2007                 b"OKAY",
2008                 b"INFOUpdating GPT...",
2009                 b"OKAY",
2010                 b"OKAY0x2000",
2011                 b"OKAY0x3000",
2012                 b"INFOmax-download-size: 0x20000",
2013                 b"INFOversion-bootloader: 1.0",
2014                 b"INFOmax-fetch-size: 0xffffffffffffffff",
2015                 b"INFOblock-device:0:total-blocks: 0x80",
2016                 b"INFOblock-device:0:block-size: 0x200",
2017                 b"INFOblock-device:0:status: idle",
2018                 b"INFOblock-device:1:total-blocks: 0x100",
2019                 b"INFOblock-device:1:block-size: 0x200",
2020                 b"INFOblock-device:1:status: idle",
2021                 b"INFOgbl-default-block: None",
2022                 b"INFOpartition-size:boot_a/0: 0x2000",
2023                 b"INFOpartition-type:boot_a/0: raw",
2024                 b"INFOpartition-size:boot_b/0: 0x3000",
2025                 b"INFOpartition-type:boot_b/0: raw",
2026                 b"INFOpartition-size:vendor_boot_a/1: 0x1000",
2027                 b"INFOpartition-type:vendor_boot_a/1: raw",
2028                 b"INFOpartition-size:vendor_boot_b/1: 0x1800",
2029                 b"INFOpartition-type:vendor_boot_b/1: raw",
2030                 format!("INFO{}:1: {}:1", FakeGblOps::GBL_TEST_VAR, FakeGblOps::GBL_TEST_VAR_VAL)
2031                     .as_bytes(),
2032                 format!("INFO{}:2: {}:2", FakeGblOps::GBL_TEST_VAR, FakeGblOps::GBL_TEST_VAR_VAL)
2033                     .as_bytes(),
2034                 b"OKAY",
2035                 b"INFOSyncing storage...",
2036                 b"OKAY",
2037             ]),
2038             "\nActual USB output:\n{}",
2039             listener.dump_usb_out_queue()
2040         );
2041     }
2042 
2043     #[test]
test_oem_update_gpt_resize()2044     fn test_oem_update_gpt_resize() {
2045         let disk_orig = include_bytes!("../../../libstorage/test/gpt_test_1.bin");
2046         let mut disk = disk_orig.to_vec();
2047         // Doubles the size of the disk
2048         disk.resize(disk_orig.len() * 2, 0);
2049 
2050         let mut storage = FakeGblOpsStorage::default();
2051         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_2.bin"));
2052         storage.add_gpt_device(&disk);
2053         let buffers = vec![vec![0u8; 128 * 1024]; 2];
2054         let mut gbl_ops = FakeGblOps::new(&storage);
2055         let listener: SharedTestListener = Default::default();
2056         let (usb, tcp) = (&listener, &listener);
2057 
2058         // Checks current size of last partition `boot_b`.
2059         listener.add_usb_input(b"getvar:partition-size:boot_b");
2060         // Sets a default block.
2061         listener.add_usb_input(b"oem gbl-set-default-block 1");
2062         let gpt = &disk_orig[..34 * 512];
2063         listener.add_usb_input(format!("download:{:#x}", gpt.len()).as_bytes());
2064         listener.add_usb_input(gpt);
2065         // No need to specify block device index
2066         listener.add_usb_input(b"flash:gpt//resize");
2067         // Checks updated size of last partition `boot_b`.
2068         listener.add_usb_input(b"getvar:partition-size:boot_b");
2069         listener.add_usb_input(b"continue");
2070 
2071         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
2072 
2073         assert_eq!(
2074             listener.usb_out_queue(),
2075             make_expected_usb_out(&[
2076                 b"OKAY0x3000",
2077                 b"INFODefault block device: 0x1",
2078                 b"OKAY",
2079                 b"DATA00004400",
2080                 b"OKAY",
2081                 b"INFOUpdating GPT...",
2082                 b"OKAY",
2083                 b"OKAY0x15a00",
2084                 b"INFOSyncing storage...",
2085                 b"OKAY",
2086             ]),
2087             "\nActual USB output:\n{}",
2088             listener.dump_usb_out_queue()
2089         );
2090     }
2091 
2092     #[test]
test_oem_update_gpt_no_downloaded_gpt()2093     fn test_oem_update_gpt_no_downloaded_gpt() {
2094         let disk = include_bytes!("../../../libstorage/test/gpt_test_1.bin");
2095         let mut storage = FakeGblOpsStorage::default();
2096         storage.add_gpt_device(&disk);
2097         let buffers = vec![vec![0u8; 128 * 1024]; 2];
2098         let mut gbl_ops = FakeGblOps::new(&storage);
2099         let listener: SharedTestListener = Default::default();
2100         let (usb, tcp) = (&listener, &listener);
2101 
2102         listener.add_usb_input(b"flash:gpt/0");
2103         listener.add_usb_input(b"continue");
2104 
2105         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
2106 
2107         assert_eq!(
2108             listener.usb_out_queue(),
2109             make_expected_usb_out(&[b"FAILNo GPT downloaded", b"INFOSyncing storage...", b"OKAY",]),
2110             "\nActual USB output:\n{}",
2111             listener.dump_usb_out_queue()
2112         );
2113     }
2114 
2115     #[test]
test_oem_update_gpt_bad_gpt()2116     fn test_oem_update_gpt_bad_gpt() {
2117         let disk = include_bytes!("../../../libstorage/test/gpt_test_1.bin");
2118         let mut storage = FakeGblOpsStorage::default();
2119         storage.add_gpt_device(&disk);
2120         let buffers = vec![vec![0u8; 128 * 1024]; 2];
2121         let mut gbl_ops = FakeGblOps::new(&storage);
2122         let listener: SharedTestListener = Default::default();
2123         let (usb, tcp) = (&listener, &listener);
2124         // Download a bad GPT.
2125         let mut gpt = disk[..34 * 512].to_vec();
2126         gpt[512] = !gpt[512];
2127         listener.add_usb_input(format!("download:{:#x}", gpt.len()).as_bytes());
2128         listener.add_usb_input(&gpt);
2129         listener.add_usb_input(b"flash:gpt/0");
2130         listener.add_usb_input(b"continue");
2131 
2132         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
2133 
2134         assert_eq!(
2135             listener.usb_out_queue(),
2136             make_expected_usb_out(&[
2137                 b"DATA00004400",
2138                 b"OKAY",
2139                 b"INFOUpdating GPT...",
2140                 b"FAILGptError(\n    IncorrectMagic(\n        6075990659671082682,\n    ),\n)",
2141                 b"INFOSyncing storage...",
2142                 b"OKAY",
2143             ]),
2144             "\nActual USB output:\n{}",
2145             listener.dump_usb_out_queue()
2146         );
2147     }
2148 
2149     #[test]
test_oem_update_gpt_invalid_input()2150     fn test_oem_update_gpt_invalid_input() {
2151         let disk_orig = include_bytes!("../../../libstorage/test/gpt_test_1.bin");
2152         let mut storage = FakeGblOpsStorage::default();
2153         storage.add_gpt_device(&disk_orig);
2154         storage.add_gpt_device(&disk_orig);
2155         let buffers = vec![vec![0u8; 128 * 1024]; 2];
2156         let mut gbl_ops = FakeGblOps::new(&storage);
2157         let listener: SharedTestListener = Default::default();
2158         let (usb, tcp) = (&listener, &listener);
2159 
2160         let gpt = &disk_orig[..34 * 512];
2161         listener.add_usb_input(format!("download:{:#x}", gpt.len()).as_bytes());
2162         listener.add_usb_input(gpt);
2163         // Missing block device ID.
2164         listener.add_usb_input(b"flash:gpt");
2165         // Out of range block device ID.
2166         listener.add_usb_input(b"flash:gpt/2");
2167         // Invalid option.
2168         listener.add_usb_input(b"flash:gpt/0/invalid-arg");
2169         listener.add_usb_input(b"continue");
2170         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
2171 
2172         assert_eq!(
2173             listener.usb_out_queue(),
2174             make_expected_usb_out(&[
2175                 b"DATA00004400",
2176                 b"OKAY",
2177                 b"FAILBlock ID is required for flashing GPT",
2178                 b"FAILInvalid block ID",
2179                 b"FAILUnknown argument",
2180                 b"INFOSyncing storage...",
2181                 b"OKAY",
2182             ]),
2183             "\nActual USB output:\n{}",
2184             listener.dump_usb_out_queue()
2185         );
2186     }
2187 
2188     #[test]
test_oem_update_gpt_fail_on_raw_blk()2189     fn test_oem_update_gpt_fail_on_raw_blk() {
2190         let disk_orig = include_bytes!("../../../libstorage/test/gpt_test_1.bin");
2191         let mut storage = FakeGblOpsStorage::default();
2192         storage.add_raw_device(c"raw_0", [0u8; 1024 * 1024]);
2193         let buffers = vec![vec![0u8; 128 * 1024]; 2];
2194         let mut gbl_ops = FakeGblOps::new(&storage);
2195         let listener: SharedTestListener = Default::default();
2196         let (usb, tcp) = (&listener, &listener);
2197 
2198         let gpt = &disk_orig[..34 * 512];
2199         listener.add_usb_input(format!("download:{:#x}", gpt.len()).as_bytes());
2200         listener.add_usb_input(gpt);
2201         listener.add_usb_input(b"flash:gpt/0");
2202         listener.add_usb_input(b"continue");
2203         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
2204 
2205         assert_eq!(
2206             listener.usb_out_queue(),
2207             make_expected_usb_out(&[
2208                 b"DATA00004400",
2209                 b"OKAY",
2210                 b"INFOUpdating GPT...",
2211                 b"FAILBlock device is not for GPT",
2212                 b"INFOSyncing storage...",
2213                 b"OKAY",
2214             ]),
2215             "\nActual USB output:\n{}",
2216             listener.dump_usb_out_queue()
2217         );
2218     }
2219 
2220     #[test]
test_oem_erase_gpt()2221     fn test_oem_erase_gpt() {
2222         let mut storage = FakeGblOpsStorage::default();
2223         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_1.bin"));
2224         storage.add_gpt_device(include_bytes!("../../../libstorage/test/gpt_test_2.bin"));
2225         let buffers = vec![vec![0u8; 128 * 1024]; 2];
2226         let mut gbl_ops = FakeGblOps::new(&storage);
2227         let listener: SharedTestListener = Default::default();
2228         let (usb, tcp) = (&listener, &listener);
2229 
2230         // Erases the GPT on disk #0.
2231         listener.add_usb_input(b"erase:gpt/0");
2232         // Checks that we can no longer get partition info on disk #0.
2233         listener.add_usb_input(b"getvar:partition-size:boot_a");
2234         listener.add_usb_input(b"getvar:partition-size:boot_b");
2235         // Checks that we can still get partition info on disk #1.
2236         listener.add_usb_input(b"getvar:partition-size:vendor_boot_a");
2237         listener.add_usb_input(b"getvar:partition-size:vendor_boot_b");
2238         listener.add_usb_input(b"continue");
2239 
2240         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
2241 
2242         assert_eq!(
2243             listener.usb_out_queue(),
2244             make_expected_usb_out(&[
2245                 b"OKAY",
2246                 b"FAILNotFound",
2247                 b"FAILNotFound",
2248                 b"OKAY0x1000",
2249                 b"OKAY0x1800",
2250                 b"INFOSyncing storage...",
2251                 b"OKAY",
2252             ]),
2253             "\nActual USB output:\n{}",
2254             listener.dump_usb_out_queue()
2255         );
2256     }
2257 
2258     #[test]
test_oem_erase_gpt_fail_on_raw_blk()2259     fn test_oem_erase_gpt_fail_on_raw_blk() {
2260         let mut storage = FakeGblOpsStorage::default();
2261         storage.add_raw_device(c"raw_0", [0u8; 1024 * 1024]);
2262         let buffers = vec![vec![0u8; 128 * 1024]; 2];
2263         let mut gbl_ops = FakeGblOps::new(&storage);
2264         let listener: SharedTestListener = Default::default();
2265         let (usb, tcp) = (&listener, &listener);
2266 
2267         listener.add_usb_input(b"erase:gpt/0");
2268         listener.add_usb_input(b"continue");
2269         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
2270 
2271         assert_eq!(
2272             listener.usb_out_queue(),
2273             make_expected_usb_out(&[
2274                 b"FAILBlock device is not for GPT",
2275                 b"INFOSyncing storage...",
2276                 b"OKAY",
2277             ]),
2278             "\nActual USB output:\n{}",
2279             listener.dump_usb_out_queue()
2280         );
2281     }
2282 
2283     /// Helper for testing fastboot set_active in fuchsia A/B/R mode.
test_run_gbl_fastboot_set_active_fuchsia_abr(cmd: &str, slot: SlotIndex)2284     fn test_run_gbl_fastboot_set_active_fuchsia_abr(cmd: &str, slot: SlotIndex) {
2285         let mut storage = FakeGblOpsStorage::default();
2286         storage.add_raw_device(c"durable_boot", [0x00u8; 4 * 1024]);
2287         let buffers = vec![vec![0u8; 128 * 1024]; 2];
2288         let mut gbl_ops = FakeGblOps::new(&storage);
2289         gbl_ops.os = Some(Os::Fuchsia);
2290         let listener: SharedTestListener = Default::default();
2291         let (usb, tcp) = (&listener, &listener);
2292 
2293         mark_slot_unbootable(&mut GblAbrOps(&mut gbl_ops), SlotIndex::A).unwrap();
2294         mark_slot_unbootable(&mut GblAbrOps(&mut gbl_ops), SlotIndex::B).unwrap();
2295 
2296         // Flash some data to `durable_boot` after A/B/R metadata. This is for testing that sync
2297         // storage is done first.
2298         let data = vec![0x55u8; 4 * 1024 - ABR_DATA_SIZE];
2299         listener.add_usb_input(b"oem gbl-enable-async-task");
2300         listener.add_usb_input(format!("download:{:#x}", 4 * 1024 - ABR_DATA_SIZE).as_bytes());
2301         listener.add_usb_input(&data);
2302         listener.add_usb_input(format!("flash:durable_boot//{:#x}", ABR_DATA_SIZE).as_bytes());
2303         // Issues set_active commands
2304         listener.add_usb_input(cmd.as_bytes());
2305         listener.add_usb_input(b"continue");
2306         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
2307 
2308         assert_eq!(
2309             listener.usb_out_queue(),
2310             make_expected_usb_out(&[
2311                 b"OKAY",
2312                 b"DATA00000fe0",
2313                 b"OKAY",
2314                 b"INFOAn IO task is launched. To sync manually, run \"oem gbl-sync-tasks\".",
2315                 b"OKAY",
2316                 b"OKAY",
2317                 b"INFOSyncing storage...",
2318                 b"OKAY",
2319             ]),
2320             "\nActual USB output:\n{}",
2321             listener.dump_usb_out_queue()
2322         );
2323 
2324         assert_eq!(get_boot_slot(&mut GblAbrOps(&mut gbl_ops), true), (slot, false));
2325         // Verifies storage sync
2326         assert_eq!(
2327             storage[0].partition_io(None).unwrap().dev().io().storage[ABR_DATA_SIZE..],
2328             data
2329         );
2330     }
2331 
2332     #[test]
test_run_gbl_fastboot_set_active_fuchsia_abr_a()2333     fn test_run_gbl_fastboot_set_active_fuchsia_abr_a() {
2334         test_run_gbl_fastboot_set_active_fuchsia_abr("set_active:a", SlotIndex::A);
2335     }
2336 
2337     #[test]
test_run_gbl_fastboot_set_active_fuchsia_abr_b()2338     fn test_run_gbl_fastboot_set_active_fuchsia_abr_b() {
2339         test_run_gbl_fastboot_set_active_fuchsia_abr("set_active:b", SlotIndex::B);
2340     }
2341 
2342     #[test]
test_run_gbl_fastboot_set_active_fuchsia_abr_invalid_slot()2343     fn test_run_gbl_fastboot_set_active_fuchsia_abr_invalid_slot() {
2344         let mut storage = FakeGblOpsStorage::default();
2345         storage.add_raw_device(c"durable_boot", [0x00u8; 4 * 1024]);
2346         let buffers = vec![vec![0u8; 128 * 1024]; 2];
2347         let mut gbl_ops = FakeGblOps::new(&storage);
2348         gbl_ops.os = Some(Os::Fuchsia);
2349         let listener: SharedTestListener = Default::default();
2350         let (usb, tcp) = (&listener, &listener);
2351 
2352         listener.add_usb_input(b"set_active:r");
2353         listener.add_usb_input(b"continue");
2354         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
2355 
2356         assert_eq!(
2357             listener.usb_out_queue(),
2358             make_expected_usb_out(&[
2359                 b"FAILInvalid slot index for Fuchsia A/B/R",
2360                 b"INFOSyncing storage...",
2361                 b"OKAY",
2362             ]),
2363             "\nActual USB output:\n{}",
2364             listener.dump_usb_out_queue()
2365         );
2366     }
2367 
2368     #[test]
test_run_gbl_fastboot_fuchsia_reboot_bootloader_abr()2369     fn test_run_gbl_fastboot_fuchsia_reboot_bootloader_abr() {
2370         let mut storage = FakeGblOpsStorage::default();
2371         storage.add_raw_device(c"durable_boot", [0x00u8; 4 * 1024]);
2372         let buffers = vec![vec![0u8; 128 * 1024]; 2];
2373         let mut gbl_ops = FakeGblOps::new(&storage);
2374         gbl_ops.os = Some(Os::Fuchsia);
2375         let listener: SharedTestListener = Default::default();
2376         let (usb, tcp) = (&listener, &listener);
2377 
2378         listener.add_usb_input(b"reboot-bootloader");
2379         listener.add_usb_input(b"continue");
2380         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
2381 
2382         assert_eq!(
2383             listener.usb_out_queue(),
2384             make_expected_usb_out(&[
2385                 b"INFOSyncing storage...",
2386                 b"INFORebooting to bootloader...",
2387                 b"OKAY",
2388                 b"FAILUnknown",
2389                 b"INFOSyncing storage...",
2390                 b"OKAY",
2391             ]),
2392             "\nActual USB output:\n{}",
2393             listener.dump_usb_out_queue()
2394         );
2395 
2396         assert_eq!(get_and_clear_one_shot_bootloader(&mut GblAbrOps(&mut gbl_ops)), Ok(true));
2397     }
2398 
2399     #[test]
test_run_gbl_fastboot_fuchsia_reboot_recovery_abr()2400     fn test_run_gbl_fastboot_fuchsia_reboot_recovery_abr() {
2401         let mut storage = FakeGblOpsStorage::default();
2402         storage.add_raw_device(c"durable_boot", [0x00u8; 4 * 1024]);
2403         let buffers = vec![vec![0u8; 128 * 1024]; 2];
2404         let mut gbl_ops = FakeGblOps::new(&storage);
2405         gbl_ops.os = Some(Os::Fuchsia);
2406         let listener: SharedTestListener = Default::default();
2407         let (usb, tcp) = (&listener, &listener);
2408 
2409         listener.add_usb_input(b"reboot-recovery");
2410         listener.add_usb_input(b"continue");
2411         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
2412 
2413         assert_eq!(
2414             listener.usb_out_queue(),
2415             make_expected_usb_out(&[
2416                 b"INFOSyncing storage...",
2417                 b"INFORebooting to recovery...",
2418                 b"OKAY",
2419                 b"FAILUnknown",
2420                 b"INFOSyncing storage...",
2421                 b"OKAY",
2422             ]),
2423             "\nActual USB output:\n{}",
2424             listener.dump_usb_out_queue()
2425         );
2426 
2427         // One shot recovery is set.
2428         assert_eq!(get_boot_slot(&mut GblAbrOps(&mut gbl_ops), true), (SlotIndex::R, false));
2429         assert_eq!(get_boot_slot(&mut GblAbrOps(&mut gbl_ops), true), (SlotIndex::A, false));
2430     }
2431 
2432     #[test]
test_legacy_fvm_partition_aliase()2433     fn test_legacy_fvm_partition_aliase() {
2434         let mut storage = FakeGblOpsStorage::default();
2435         storage.add_raw_device(c"fuchsia-fvm", [0x00u8; 4 * 1024]);
2436         let buffers = vec![vec![0u8; 128 * 1024]; 2];
2437         let mut gbl_ops = FakeGblOps::new(&storage);
2438         gbl_ops.os = Some(Os::Fuchsia);
2439         let listener: SharedTestListener = Default::default();
2440         let (usb, tcp) = (&listener, &listener);
2441 
2442         listener.add_usb_input(format!("download:{:#x}", 4 * 1024).as_bytes());
2443         listener.add_usb_input(&[0xaau8; 4 * 1024]);
2444         listener.add_usb_input(b"flash:fvm");
2445         listener.add_usb_input(b"continue");
2446         block_on(run_gbl_fastboot_stack::<2>(&mut gbl_ops, buffers, Some(usb), Some(tcp), &mut []));
2447 
2448         assert_eq!(
2449             listener.usb_out_queue(),
2450             make_expected_usb_out(&[
2451                 b"DATA00001000",
2452                 b"OKAY",
2453                 b"OKAY",
2454                 b"INFOSyncing storage...",
2455                 b"OKAY",
2456             ]),
2457             "\nActual USB output:\n{}",
2458             listener.dump_usb_out_queue()
2459         );
2460     }
2461 }
2462