Skip to main content

gpt/
lib.rs

1// Copyright 2024 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use anyhow::{Context as _, Error, anyhow};
6use block_client::{BlockClient, BufferSlice, MutableBufferSlice, RemoteBlockClient};
7use fuchsia_sync::Mutex;
8use std::collections::BTreeMap;
9use std::sync::Arc;
10use zerocopy::{FromBytes as _, IntoBytes as _};
11
12pub mod format;
13
14/// GPT GUIDs are stored in mixed-endian format (see Appendix A of the EFI spec).  To ensure this is
15/// correctly handled, wrap the Uuid type to hide methods that use the UUIDs inappropriately.
16#[derive(Clone, Default, Debug)]
17pub struct Guid(uuid::Uuid);
18
19impl From<uuid::Uuid> for Guid {
20    fn from(uuid: uuid::Uuid) -> Self {
21        Self(uuid)
22    }
23}
24
25impl Guid {
26    pub fn from_bytes(bytes: [u8; 16]) -> Self {
27        Self(uuid::Uuid::from_bytes_le(bytes))
28    }
29
30    pub fn to_bytes(&self) -> [u8; 16] {
31        self.0.to_bytes_le()
32    }
33
34    pub fn to_string(&self) -> String {
35        self.0.to_string()
36    }
37
38    pub fn nil() -> Self {
39        Self(uuid::Uuid::nil())
40    }
41
42    pub fn generate() -> Self {
43        Self(uuid::Uuid::new_v4())
44    }
45}
46
47#[derive(Clone, Debug)]
48pub struct PartitionInfo {
49    pub label: String,
50    pub type_guid: Guid,
51    pub instance_guid: Guid,
52    pub start_block: u64,
53    pub num_blocks: u64,
54    pub flags: u64,
55}
56
57impl PartitionInfo {
58    pub fn from_entry(entry: &format::PartitionTableEntry) -> Result<Self, Error> {
59        let label = String::from_utf16(entry.name.split(|v| *v == 0u16).next().unwrap())?;
60        Ok(Self {
61            label,
62            type_guid: Guid::from_bytes(entry.type_guid),
63            instance_guid: Guid::from_bytes(entry.instance_guid),
64            start_block: entry.first_lba,
65            num_blocks: entry
66                .last_lba
67                .checked_add(1)
68                .unwrap()
69                .checked_sub(entry.first_lba)
70                .unwrap(),
71            flags: entry.flags,
72        })
73    }
74
75    pub fn as_entry(&self) -> format::PartitionTableEntry {
76        let mut name = [0u16; 36];
77        let raw = self.label.encode_utf16().collect::<Vec<_>>();
78        assert!(raw.len() <= name.len());
79        name[..raw.len()].copy_from_slice(&raw[..]);
80        format::PartitionTableEntry {
81            type_guid: self.type_guid.to_bytes(),
82            instance_guid: self.instance_guid.to_bytes(),
83            first_lba: self.start_block,
84            last_lba: self.start_block + self.num_blocks.saturating_sub(1),
85            flags: self.flags,
86            name,
87        }
88    }
89
90    pub fn nil() -> Self {
91        Self {
92            label: String::default(),
93            type_guid: Guid::default(),
94            instance_guid: Guid::default(),
95            start_block: 0,
96            num_blocks: 0,
97            flags: 0,
98        }
99    }
100
101    pub fn is_nil(&self) -> bool {
102        self.label == ""
103            && self.type_guid.0.is_nil()
104            && self.instance_guid.0.is_nil()
105            && self.start_block == 0
106            && self.num_blocks == 0
107            && self.flags == 0
108    }
109}
110
111enum WhichHeader {
112    Primary,
113    Backup,
114}
115
116impl WhichHeader {
117    fn offset(&self, block_size: u64, block_count: u64) -> u64 {
118        match self {
119            Self::Primary => block_size,
120            Self::Backup => (block_count - 1) * block_size,
121        }
122    }
123}
124
125async fn load_metadata(
126    client: &RemoteBlockClient,
127    which: WhichHeader,
128) -> Result<(format::Header, BTreeMap<u32, PartitionInfo>), Error> {
129    let bs = client.block_size() as usize;
130    let mut header_block = vec![0u8; client.block_size() as usize];
131    client
132        .read_at(
133            MutableBufferSlice::Memory(&mut header_block[..]),
134            which.offset(bs as u64, client.block_count() as u64),
135        )
136        .await
137        .context("Read header")?;
138    let (header, _) = format::Header::ref_from_prefix(&header_block[..])
139        .map_err(|_| anyhow!("Header has invalid size"))?;
140    header.ensure_integrity(client.block_count(), client.block_size() as u64)?;
141    let partition_table_offset = header.part_start * bs as u64;
142    let partition_table_size = (header.num_parts * header.part_size) as usize;
143    let partition_table_size_rounded = partition_table_size
144        .checked_next_multiple_of(bs)
145        .ok_or_else(|| anyhow!("Overflow when rounding up partition table size "))?;
146    let mut partition_table = BTreeMap::new();
147    if header.num_parts > 0 {
148        let mut partition_table_blocks = vec![0u8; partition_table_size_rounded];
149        client
150            .read_at(
151                MutableBufferSlice::Memory(&mut partition_table_blocks[..]),
152                partition_table_offset,
153            )
154            .await
155            .with_context(|| {
156                format!(
157                    "Failed to read partition table (sz {}) from offset {}",
158                    partition_table_size, partition_table_offset
159                )
160            })?;
161        let crc = crc::Crc::<u32>::new(&crc::CRC_32_ISO_HDLC)
162            .checksum(&partition_table_blocks[..partition_table_size]);
163        anyhow::ensure!(header.crc32_parts == crc, "Invalid partition table checksum");
164
165        for i in 0..header.num_parts as usize {
166            let entry_raw = &partition_table_blocks
167                [i * header.part_size as usize..(i + 1) * header.part_size as usize];
168            let (entry, _) = format::PartitionTableEntry::ref_from_prefix(entry_raw)
169                .map_err(|_| anyhow!("Failed to parse partition {i}"))?;
170            if entry.is_empty() {
171                continue;
172            }
173            entry.ensure_integrity().context("GPT partition table entry invalid!")?;
174
175            partition_table.insert(i as u32, PartitionInfo::from_entry(entry)?);
176        }
177    }
178    Ok((header.clone(), partition_table))
179}
180
181struct TransactionState {
182    pending_id: u64,
183    next_id: u64,
184}
185
186impl Default for TransactionState {
187    fn default() -> Self {
188        Self { pending_id: u64::MAX, next_id: 0 }
189    }
190}
191
192/// Manages a connection to a GPT-formatted block device.
193pub struct Gpt {
194    client: Arc<RemoteBlockClient>,
195    header: format::Header,
196    partitions: BTreeMap<u32, PartitionInfo>,
197    transaction_state: Arc<Mutex<TransactionState>>,
198}
199
200impl std::fmt::Debug for Gpt {
201    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
202        f.debug_struct("Gpt")
203            .field("header", &self.header)
204            .field("partitions", &self.partitions)
205            .finish()
206    }
207}
208
209#[derive(Eq, thiserror::Error, Clone, Debug, PartialEq)]
210pub enum TransactionCommitError {
211    #[error("I/O error")]
212    Io,
213    #[error("Invalid arguments")]
214    InvalidArguments,
215    #[error("No space")]
216    NoSpace,
217}
218
219impl From<format::FormatError> for TransactionCommitError {
220    fn from(error: format::FormatError) -> Self {
221        match error {
222            format::FormatError::InvalidArguments => Self::InvalidArguments,
223            format::FormatError::NoSpace => Self::NoSpace,
224        }
225    }
226}
227
228impl From<TransactionCommitError> for zx::Status {
229    fn from(error: TransactionCommitError) -> zx::Status {
230        match error {
231            TransactionCommitError::Io => zx::Status::IO,
232            TransactionCommitError::InvalidArguments => zx::Status::INVALID_ARGS,
233            TransactionCommitError::NoSpace => zx::Status::NO_SPACE,
234        }
235    }
236}
237
238#[derive(Eq, thiserror::Error, Clone, Debug, PartialEq)]
239pub enum AddPartitionError {
240    #[error("Invalid arguments")]
241    InvalidArguments,
242    #[error("No space")]
243    NoSpace,
244}
245
246impl From<AddPartitionError> for zx::Status {
247    fn from(error: AddPartitionError) -> zx::Status {
248        match error {
249            AddPartitionError::InvalidArguments => zx::Status::INVALID_ARGS,
250            AddPartitionError::NoSpace => zx::Status::NO_SPACE,
251        }
252    }
253}
254
255impl Gpt {
256    /// Loads and validates a GPT-formatted block device.
257    pub async fn open(client: Arc<RemoteBlockClient>) -> Result<Self, Error> {
258        let mut restore_primary = false;
259        let (header, partitions) = match load_metadata(&client, WhichHeader::Primary).await {
260            Ok(v) => v,
261            Err(error) => {
262                log::warn!(error:?; "Failed to load primary metadata; falling back to backup");
263                restore_primary = true;
264                load_metadata(&client, WhichHeader::Backup)
265                    .await
266                    .context("Failed to load backup metadata")?
267            }
268        };
269        let mut this = Self {
270            client,
271            header,
272            partitions,
273            transaction_state: Arc::new(Mutex::new(TransactionState::default())),
274        };
275        if restore_primary {
276            log::info!("Restoring primary metadata from backup!");
277            this.header.backup_lba = this.header.current_lba;
278            this.header.current_lba = 1;
279            this.header.part_start = 2;
280            this.header.crc32 = this.header.compute_checksum();
281            let partition_table =
282                this.flattened_partitions().into_iter().map(|v| v.as_entry()).collect::<Vec<_>>();
283            let partition_table_raw = format::serialize_partition_table(
284                &mut this.header,
285                this.client.block_size() as usize,
286                this.client.block_count(),
287                &partition_table[..],
288            )
289            .context("Failed to serialize existing partition table")?;
290            this.write_metadata(&this.header, &partition_table_raw[..])
291                .await
292                .context("Failed to restore primary metadata")?;
293        }
294        Ok(this)
295    }
296
297    /// Formats `client` as a new GPT with `partitions`.  Overwrites any existing GPT on the block
298    /// device.
299    pub async fn format(
300        client: Arc<RemoteBlockClient>,
301        partitions: Vec<PartitionInfo>,
302    ) -> Result<Self, Error> {
303        let header = format::Header::new(
304            client.block_count(),
305            client.block_size(),
306            partitions.len() as u32,
307        )?;
308        let mut this = Self {
309            client,
310            header,
311            partitions: BTreeMap::new(),
312            transaction_state: Arc::new(Mutex::new(TransactionState::default())),
313        };
314        let mut transaction = this.create_transaction().unwrap();
315        transaction.partitions = partitions;
316        this.commit_transaction(transaction).await?;
317        Ok(this)
318    }
319
320    pub fn client(&self) -> &Arc<RemoteBlockClient> {
321        &self.client
322    }
323
324    #[cfg(test)]
325    fn take_client(self) -> Arc<RemoteBlockClient> {
326        self.client
327    }
328
329    pub fn header(&self) -> &format::Header {
330        &self.header
331    }
332
333    pub fn partitions(&self) -> &BTreeMap<u32, PartitionInfo> {
334        &self.partitions
335    }
336
337    // We only store valid partitions in memory.  This function allows us to flatten this back out
338    // to a non-sparse array for serialization.
339    fn flattened_partitions(&self) -> Vec<PartitionInfo> {
340        let mut partitions = vec![PartitionInfo::nil(); self.header.num_parts as usize];
341        for (idx, partition) in &self.partitions {
342            partitions[*idx as usize] = partition.clone();
343        }
344        partitions
345    }
346
347    /// Returns None if there's already a pending transaction.
348    pub fn create_transaction(&self) -> Option<Transaction> {
349        {
350            let mut state = self.transaction_state.lock();
351            if state.pending_id != u64::MAX {
352                return None;
353            } else {
354                state.pending_id = state.next_id;
355                state.next_id += 1;
356            }
357        }
358        Some(Transaction {
359            partitions: self.flattened_partitions(),
360            transaction_state: self.transaction_state.clone(),
361        })
362    }
363
364    pub async fn commit_transaction(
365        &mut self,
366        mut transaction: Transaction,
367    ) -> Result<(), TransactionCommitError> {
368        let mut new_header = self.header.clone();
369        let entries =
370            transaction.partitions.iter().map(|entry| entry.as_entry()).collect::<Vec<_>>();
371        let partition_table_raw = format::serialize_partition_table(
372            &mut new_header,
373            self.client.block_size() as usize,
374            self.client.block_count(),
375            &entries[..],
376        )?;
377
378        let mut backup_header = new_header.clone();
379        backup_header.current_lba = backup_header.backup_lba;
380        backup_header.backup_lba = 1;
381        backup_header.part_start = backup_header.last_usable + 1;
382        backup_header.crc32 = backup_header.compute_checksum();
383
384        // Per section 5.3.2 of the UEFI spec, the backup metadata must be written first.  The spec
385        // permits the partition table entries and header to be written in either order.
386        self.write_metadata(&backup_header, &partition_table_raw[..]).await.map_err(|err| {
387            log::warn!(err:?; "Failed to write metadata");
388            TransactionCommitError::Io
389        })?;
390        // NB: It would be preferable to use a barrier here, but not all drivers support barriers at
391        // this time.
392        // TODO(https://fxbug.dev/416348380): Use a barrier between writing secondary/primary.
393        self.client.flush().await.map_err(|err| {
394            log::warn!(err:?; "Failed to flush metadata writes");
395            TransactionCommitError::Io
396        })?;
397        self.write_metadata(&new_header, &partition_table_raw[..]).await.map_err(|err| {
398            log::warn!(err:?; "Failed to write metadata");
399            TransactionCommitError::Io
400        })?;
401        self.client.flush().await.map_err(|err| {
402            log::warn!(err:?; "Failed to flush metadata writes");
403            TransactionCommitError::Io
404        })?;
405
406        self.header = new_header;
407        self.partitions = BTreeMap::new();
408        let mut idx = 0;
409        for partition in std::mem::take(&mut transaction.partitions) {
410            if !partition.is_nil() {
411                self.partitions.insert(idx, partition);
412            }
413            idx += 1;
414        }
415        Ok(())
416    }
417
418    /// Adds a partition in `transaction`.  `info.start_block` must be unset and will be dynamically
419    /// chosen in a first-fit manner.
420    /// The indedx of the partition in the table is returned on success.
421    pub fn add_partition(
422        &mut self,
423        transaction: &mut Transaction,
424        mut info: PartitionInfo,
425    ) -> Result<usize, AddPartitionError> {
426        assert_eq!(info.start_block, 0);
427
428        if info.label.is_empty()
429            || info.type_guid.0.is_nil()
430            || info.instance_guid.0.is_nil()
431            || info.num_blocks == 0
432        {
433            return Err(AddPartitionError::InvalidArguments);
434        }
435
436        let mut allocated_ranges = vec![
437            0..self.header.first_usable,
438            self.header.last_usable + 1..self.client.block_count(),
439        ];
440        let mut slot_idx = None;
441        for i in 0..transaction.partitions.len() {
442            let partition = &transaction.partitions[i];
443            if slot_idx.is_none() && partition.is_nil() {
444                slot_idx = Some(i);
445            }
446            if !partition.is_nil() {
447                allocated_ranges
448                    .push(partition.start_block..partition.start_block + partition.num_blocks);
449            }
450        }
451        let slot_idx = slot_idx.ok_or(AddPartitionError::NoSpace)?;
452        allocated_ranges.sort_by_key(|range| range.start);
453
454        let mut start_block = None;
455        for [a, b] in allocated_ranges.array_windows() {
456            if b.start - a.end >= info.num_blocks {
457                start_block = Some(a.end);
458                break;
459            }
460        }
461        info.start_block = start_block.ok_or(AddPartitionError::NoSpace)?;
462
463        transaction.partitions[slot_idx] = info;
464        Ok(slot_idx)
465    }
466
467    async fn write_metadata(
468        &self,
469        header: &format::Header,
470        partition_table: &[u8],
471    ) -> Result<(), Error> {
472        let bs = self.client.block_size() as usize;
473        let mut header_block = vec![0u8; bs];
474        header.write_to_prefix(&mut header_block[..]).unwrap();
475        self.client
476            .write_at(BufferSlice::Memory(&header_block[..]), header.current_lba * bs as u64)
477            .await
478            .context("Failed to write header")?;
479        if !partition_table.is_empty() {
480            self.client
481                .write_at(BufferSlice::Memory(partition_table), header.part_start * bs as u64)
482                .await
483                .context("Failed to write partition table")?;
484        }
485        Ok(())
486    }
487}
488
489/// Pending changes to the GPT.
490pub struct Transaction {
491    pub partitions: Vec<PartitionInfo>,
492    transaction_state: Arc<Mutex<TransactionState>>,
493}
494
495impl std::fmt::Debug for Transaction {
496    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
497        f.debug_struct("Transaction").field("partitions", &self.partitions).finish()
498    }
499}
500
501impl Drop for Transaction {
502    fn drop(&mut self) {
503        let mut state = self.transaction_state.lock();
504        debug_assert!(state.pending_id != u64::MAX);
505        state.pending_id = u64::MAX;
506    }
507}
508
509#[cfg(test)]
510mod tests {
511    use crate::{AddPartitionError, Gpt, Guid, PartitionInfo};
512    use block_client::{BlockClient as _, BufferSlice, MutableBufferSlice, RemoteBlockClient};
513    use fidl_fuchsia_storage_block as fblock;
514    use fuchsia_async as fasync;
515    use std::ops::Range;
516    use std::sync::Arc;
517    use std::sync::atomic::{AtomicBool, Ordering};
518    use test_vmo_backed_block_server::{
519        InitialContents, Observer, VmoBackedServer, VmoBackedServerOptions, WriteAction, WriteCache,
520    };
521
522    async fn connect_to_server(
523        server: VmoBackedServer,
524    ) -> (Arc<RemoteBlockClient>, fasync::Task<()>) {
525        let (client, server_end) = fidl::endpoints::create_proxy::<fblock::BlockMarker>();
526        let task =
527            fasync::Task::spawn(
528                async move { server.serve(server_end.into_stream()).await.unwrap() },
529            );
530        let client = Arc::new(RemoteBlockClient::new(client).await.unwrap());
531        (client, task)
532    }
533
534    #[fuchsia::test]
535    async fn load_unformatted_gpt() {
536        let server = VmoBackedServer::new(8, 512, &[]).expect("Failed to create VmoBackedServer");
537        let (client, _task) = connect_to_server(server).await;
538        Gpt::open(client).await.expect_err("load should fail");
539    }
540
541    #[fuchsia::test]
542    async fn load_formatted_empty_gpt() {
543        let server = VmoBackedServer::new(8, 512, &[]).expect("Failed to create VmoBackedServer");
544        let (client, _task) = connect_to_server(server).await;
545        Gpt::format(client.clone(), vec![]).await.expect("format failed");
546        Gpt::open(client).await.expect("load should succeed");
547    }
548
549    #[fuchsia::test]
550    async fn load_formatted_gpt_with_minimal_size() {
551        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
552        const PART_INSTANCE_GUID: [u8; 16] = [2u8; 16];
553        const PART_NAME: &str = "part";
554
555        let server = VmoBackedServer::new(6, 4096, &[]).expect("Failed to create VmoBackedServer");
556        let (client, _task) = connect_to_server(server).await;
557        Gpt::format(
558            client.clone(),
559            vec![PartitionInfo {
560                label: PART_NAME.to_string(),
561                type_guid: Guid::from_bytes(PART_TYPE_GUID),
562                instance_guid: Guid::from_bytes(PART_INSTANCE_GUID),
563                start_block: 3,
564                num_blocks: 1,
565                flags: 0,
566            }],
567        )
568        .await
569        .expect("format failed");
570        let manager = Gpt::open(client).await.expect("load should succeed");
571        assert_eq!(manager.header.first_usable, 3);
572        assert_eq!(manager.header.last_usable, 3);
573        let partition = manager.partitions().get(&0).expect("No entry found");
574        assert_eq!(partition.start_block, 3);
575        assert_eq!(partition.num_blocks, 1);
576        assert!(manager.partitions().get(&1).is_none());
577    }
578
579    #[fuchsia::test]
580    async fn load_formatted_gpt_with_one_partition() {
581        const PART_TYPE_GUID: [u8; 16] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15];
582        const PART_INSTANCE_GUID: [u8; 16] =
583            [16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31];
584        const PART_NAME: &str = "part";
585
586        let server = VmoBackedServer::new(8, 512, &[]).expect("Failed to create VmoBackedServer");
587        let (client, _task) = connect_to_server(server).await;
588        Gpt::format(
589            client.clone(),
590            vec![PartitionInfo {
591                label: PART_NAME.to_string(),
592                type_guid: Guid::from_bytes(PART_TYPE_GUID),
593                instance_guid: Guid::from_bytes(PART_INSTANCE_GUID),
594                start_block: 4,
595                num_blocks: 1,
596                flags: 0,
597            }],
598        )
599        .await
600        .expect("format failed");
601        let manager = Gpt::open(client).await.expect("load should succeed");
602        let partition = manager.partitions().get(&0).expect("No entry found");
603        assert_eq!(partition.label, "part");
604        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
605        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_GUID);
606        assert_eq!(partition.start_block, 4);
607        assert_eq!(partition.num_blocks, 1);
608        assert!(manager.partitions().get(&1).is_none());
609    }
610
611    #[fuchsia::test]
612    async fn load_formatted_gpt_with_two_partitions() {
613        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
614        const PART_INSTANCE_1_GUID: [u8; 16] = [2u8; 16];
615        const PART_INSTANCE_2_GUID: [u8; 16] = [3u8; 16];
616        const PART_1_NAME: &str = "part1";
617        const PART_2_NAME: &str = "part2";
618
619        let server = VmoBackedServer::new(16, 512, &[]).expect("Failed to create VmoBackedServer");
620        let (client, _task) = connect_to_server(server).await;
621        Gpt::format(
622            client.clone(),
623            vec![
624                PartitionInfo {
625                    label: PART_1_NAME.to_string(),
626                    type_guid: Guid::from_bytes(PART_TYPE_GUID),
627                    instance_guid: Guid::from_bytes(PART_INSTANCE_1_GUID),
628                    start_block: 4,
629                    num_blocks: 1,
630                    flags: 0,
631                },
632                PartitionInfo {
633                    label: PART_2_NAME.to_string(),
634                    type_guid: Guid::from_bytes(PART_TYPE_GUID),
635                    instance_guid: Guid::from_bytes(PART_INSTANCE_2_GUID),
636                    start_block: 7,
637                    num_blocks: 1,
638                    flags: 0,
639                },
640            ],
641        )
642        .await
643        .expect("format failed");
644        let manager = Gpt::open(client).await.expect("load should succeed");
645        let partition = manager.partitions().get(&0).expect("No entry found");
646        assert_eq!(partition.label, PART_1_NAME);
647        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
648        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_1_GUID);
649        assert_eq!(partition.start_block, 4);
650        assert_eq!(partition.num_blocks, 1);
651        let partition = manager.partitions().get(&1).expect("No entry found");
652        assert_eq!(partition.label, PART_2_NAME);
653        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
654        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_2_GUID);
655        assert_eq!(partition.start_block, 7);
656        assert_eq!(partition.num_blocks, 1);
657        assert!(manager.partitions().get(&2).is_none());
658    }
659
660    #[fuchsia::test]
661    async fn load_formatted_gpt_with_extra_bytes_in_partition_name() {
662        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
663        const PART_INSTANCE_GUID: [u8; 16] = [2u8; 16];
664        const PART_NAME: &str = "part\0extrastuff";
665
666        let server = VmoBackedServer::new(8, 512, &[]).expect("Failed to create VmoBackedServer");
667        let (client, _task) = connect_to_server(server).await;
668        Gpt::format(
669            client.clone(),
670            vec![PartitionInfo {
671                label: PART_NAME.to_string(),
672                type_guid: Guid::from_bytes(PART_TYPE_GUID),
673                instance_guid: Guid::from_bytes(PART_INSTANCE_GUID),
674                start_block: 4,
675                num_blocks: 1,
676                flags: 0,
677            }],
678        )
679        .await
680        .expect("format failed");
681        let manager = Gpt::open(client).await.expect("load should succeed");
682        let partition = manager.partitions().get(&0).expect("No entry found");
683        // The name should have everything after the first nul byte stripped.
684        assert_eq!(partition.label, "part");
685    }
686
687    #[fuchsia::test]
688    async fn load_formatted_gpt_with_empty_partition_name() {
689        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
690        const PART_INSTANCE_GUID: [u8; 16] = [2u8; 16];
691        const PART_NAME: &str = "";
692
693        let server = VmoBackedServer::new(8, 512, &[]).expect("Failed to create VmoBackedServer");
694        let (client, _task) = connect_to_server(server).await;
695        Gpt::format(
696            client.clone(),
697            vec![PartitionInfo {
698                label: PART_NAME.to_string(),
699                type_guid: Guid::from_bytes(PART_TYPE_GUID),
700                instance_guid: Guid::from_bytes(PART_INSTANCE_GUID),
701                start_block: 4,
702                num_blocks: 1,
703                flags: 0,
704            }],
705        )
706        .await
707        .expect("format failed");
708        let manager = Gpt::open(client).await.expect("load should succeed");
709        let partition = manager.partitions().get(&0).expect("No entry found");
710        assert_eq!(partition.label, "");
711    }
712
713    #[fuchsia::test]
714    async fn load_formatted_gpt_with_invalid_primary_header() {
715        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
716        const PART_INSTANCE_1_GUID: [u8; 16] = [2u8; 16];
717        const PART_INSTANCE_2_GUID: [u8; 16] = [3u8; 16];
718        const PART_1_NAME: &str = "part1";
719        const PART_2_NAME: &str = "part2";
720
721        let server = VmoBackedServer::new(16, 512, &[]).expect("Failed to create VmoBackedServer");
722        let (client, _task) = connect_to_server(server).await;
723        Gpt::format(
724            client.clone(),
725            vec![
726                PartitionInfo {
727                    label: PART_1_NAME.to_string(),
728                    type_guid: Guid::from_bytes(PART_TYPE_GUID),
729                    instance_guid: Guid::from_bytes(PART_INSTANCE_1_GUID),
730                    start_block: 4,
731                    num_blocks: 1,
732                    flags: 0,
733                },
734                PartitionInfo {
735                    label: PART_2_NAME.to_string(),
736                    type_guid: Guid::from_bytes(PART_TYPE_GUID),
737                    instance_guid: Guid::from_bytes(PART_INSTANCE_2_GUID),
738                    start_block: 7,
739                    num_blocks: 1,
740                    flags: 0,
741                },
742            ],
743        )
744        .await
745        .expect("format failed");
746        // Clobber the primary header.  The backup should allow the GPT to be used.
747        client.write_at(BufferSlice::Memory(&[0xffu8; 512]), 512).await.unwrap();
748        let manager = Gpt::open(client).await.expect("load should succeed");
749        let partition = manager.partitions().get(&0).expect("No entry found");
750        assert_eq!(partition.label, PART_1_NAME);
751        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
752        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_1_GUID);
753        assert_eq!(partition.start_block, 4);
754        assert_eq!(partition.num_blocks, 1);
755        let partition = manager.partitions().get(&1).expect("No entry found");
756        assert_eq!(partition.label, PART_2_NAME);
757        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
758        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_2_GUID);
759        assert_eq!(partition.start_block, 7);
760        assert_eq!(partition.num_blocks, 1);
761        assert!(manager.partitions().get(&2).is_none());
762    }
763
764    #[fuchsia::test]
765    async fn load_formatted_gpt_with_invalid_primary_partition_table() {
766        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
767        const PART_INSTANCE_1_GUID: [u8; 16] = [2u8; 16];
768        const PART_INSTANCE_2_GUID: [u8; 16] = [3u8; 16];
769        const PART_1_NAME: &str = "part1";
770        const PART_2_NAME: &str = "part2";
771
772        let server = VmoBackedServer::new(16, 512, &[]).expect("Failed to create VmoBackedServer");
773        let (client, _task) = connect_to_server(server).await;
774        Gpt::format(
775            client.clone(),
776            vec![
777                PartitionInfo {
778                    label: PART_1_NAME.to_string(),
779                    type_guid: Guid::from_bytes(PART_TYPE_GUID),
780                    instance_guid: Guid::from_bytes(PART_INSTANCE_1_GUID),
781                    start_block: 4,
782                    num_blocks: 1,
783                    flags: 0,
784                },
785                PartitionInfo {
786                    label: PART_2_NAME.to_string(),
787                    type_guid: Guid::from_bytes(PART_TYPE_GUID),
788                    instance_guid: Guid::from_bytes(PART_INSTANCE_2_GUID),
789                    start_block: 7,
790                    num_blocks: 1,
791                    flags: 0,
792                },
793            ],
794        )
795        .await
796        .expect("format failed");
797        // Clobber the primary partition table.  The backup should allow the GPT to be used.
798        client.write_at(BufferSlice::Memory(&[0xffu8; 512]), 1024).await.unwrap();
799        let manager = Gpt::open(client).await.expect("load should succeed");
800        let partition = manager.partitions().get(&0).expect("No entry found");
801        assert_eq!(partition.label, PART_1_NAME);
802        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
803        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_1_GUID);
804        assert_eq!(partition.start_block, 4);
805        assert_eq!(partition.num_blocks, 1);
806        let partition = manager.partitions().get(&1).expect("No entry found");
807        assert_eq!(partition.label, PART_2_NAME);
808        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
809        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_2_GUID);
810        assert_eq!(partition.start_block, 7);
811        assert_eq!(partition.num_blocks, 1);
812        assert!(manager.partitions().get(&2).is_none());
813    }
814
815    #[fuchsia::test]
816    async fn drop_transaction() {
817        let server = VmoBackedServer::new(16, 512, &[]).expect("Failed to create VmoBackedServer");
818        let (client, _task) = connect_to_server(server).await;
819        Gpt::format(client.clone(), vec![]).await.expect("format failed");
820        let manager = Gpt::open(client).await.expect("load should succeed");
821        {
822            let _transaction = manager.create_transaction().unwrap();
823            assert!(manager.create_transaction().is_none());
824        }
825        let _transaction =
826            manager.create_transaction().expect("Transaction dropped but not available");
827    }
828
829    #[fuchsia::test]
830    async fn commit_empty_transaction() {
831        let server = VmoBackedServer::new(16, 512, &[]).expect("Failed to create VmoBackedServer");
832        let (client, _task) = connect_to_server(server).await;
833        Gpt::format(client.clone(), vec![]).await.expect("format failed");
834        let mut manager = Gpt::open(client).await.expect("load should succeed");
835        let transaction = manager.create_transaction().unwrap();
836        manager.commit_transaction(transaction).await.expect("Commit failed");
837
838        // Check state before and after a reload, to ensure both the in-memory and on-disk
839        // representation match.
840        assert_eq!(manager.header().num_parts, 0);
841        assert!(manager.partitions().is_empty());
842        let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
843        assert_eq!(manager.header().num_parts, 0);
844        assert!(manager.partitions().is_empty());
845    }
846
847    #[fuchsia::test]
848    async fn add_partition_in_transaction() {
849        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
850        const PART_INSTANCE_1_GUID: [u8; 16] = [2u8; 16];
851        const PART_INSTANCE_2_GUID: [u8; 16] = [3u8; 16];
852        const PART_1_NAME: &str = "part1";
853        const PART_2_NAME: &str = "part2";
854
855        let server = VmoBackedServer::new(16, 512, &[]).expect("Failed to create VmoBackedServer");
856        let (client, _task) = connect_to_server(server).await;
857        Gpt::format(
858            client.clone(),
859            vec![PartitionInfo {
860                label: PART_1_NAME.to_string(),
861                type_guid: Guid::from_bytes(PART_TYPE_GUID),
862                instance_guid: Guid::from_bytes(PART_INSTANCE_1_GUID),
863                start_block: 4,
864                num_blocks: 1,
865                flags: 0,
866            }],
867        )
868        .await
869        .expect("format failed");
870        let mut manager = Gpt::open(client).await.expect("load should succeed");
871        let mut transaction = manager.create_transaction().unwrap();
872        assert_eq!(transaction.partitions.len(), 1);
873        transaction.partitions.push(crate::PartitionInfo {
874            label: PART_2_NAME.to_string(),
875            type_guid: crate::Guid::from_bytes(PART_TYPE_GUID),
876            instance_guid: crate::Guid::from_bytes(PART_INSTANCE_2_GUID),
877            start_block: 7,
878            num_blocks: 1,
879            flags: 0,
880        });
881        manager.commit_transaction(transaction).await.expect("Commit failed");
882
883        // Check state before and after a reload, to ensure both the in-memory and on-disk
884        // representation match.
885        assert_eq!(manager.header().num_parts, 2);
886        assert!(manager.partitions().get(&2).is_none());
887        let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
888        assert_eq!(manager.header().num_parts, 2);
889        let partition = manager.partitions().get(&0).expect("No entry found");
890        assert_eq!(partition.label, PART_1_NAME);
891        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
892        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_1_GUID);
893        assert_eq!(partition.start_block, 4);
894        assert_eq!(partition.num_blocks, 1);
895        let partition = manager.partitions().get(&1).expect("No entry found");
896        assert_eq!(partition.label, PART_2_NAME);
897        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
898        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_2_GUID);
899        assert_eq!(partition.start_block, 7);
900        assert_eq!(partition.num_blocks, 1);
901        assert!(manager.partitions().get(&2).is_none());
902    }
903
904    #[fuchsia::test]
905    async fn remove_partition_in_transaction() {
906        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
907        const PART_INSTANCE_GUID: [u8; 16] = [2u8; 16];
908        const PART_NAME: &str = "part1";
909
910        let server = VmoBackedServer::new(16, 512, &[]).expect("Failed to create VmoBackedServer");
911        let (client, _task) = connect_to_server(server).await;
912        Gpt::format(
913            client.clone(),
914            vec![PartitionInfo {
915                label: PART_NAME.to_string(),
916                type_guid: Guid::from_bytes(PART_TYPE_GUID),
917                instance_guid: Guid::from_bytes(PART_INSTANCE_GUID),
918                start_block: 4,
919                num_blocks: 1,
920                flags: 0,
921            }],
922        )
923        .await
924        .expect("format failed");
925        let mut manager = Gpt::open(client).await.expect("load should succeed");
926        let mut transaction = manager.create_transaction().unwrap();
927        assert_eq!(transaction.partitions.len(), 1);
928        transaction.partitions.clear();
929        manager.commit_transaction(transaction).await.expect("Commit failed");
930
931        // Check state before and after a reload, to ensure both the in-memory and on-disk
932        // representation match.
933        assert_eq!(manager.header().num_parts, 0);
934        assert!(manager.partitions().get(&0).is_none());
935        let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
936        assert_eq!(manager.header().num_parts, 0);
937        assert!(manager.partitions().get(&0).is_none());
938    }
939
940    #[fuchsia::test]
941    async fn modify_partition_in_transaction() {
942        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
943        const PART_INSTANCE_1_GUID: [u8; 16] = [2u8; 16];
944        const PART_INSTANCE_2_GUID: [u8; 16] = [3u8; 16];
945        const PART_1_NAME: &str = "part1";
946        const PART_2_NAME: &str = "part2";
947
948        let server = VmoBackedServer::new(16, 512, &[]).expect("Failed to create VmoBackedServer");
949        let (client, _task) = connect_to_server(server).await;
950        Gpt::format(
951            client.clone(),
952            vec![PartitionInfo {
953                label: PART_1_NAME.to_string(),
954                type_guid: Guid::from_bytes(PART_TYPE_GUID),
955                instance_guid: Guid::from_bytes(PART_INSTANCE_1_GUID),
956                start_block: 4,
957                num_blocks: 1,
958                flags: 0,
959            }],
960        )
961        .await
962        .expect("format failed");
963        let mut manager = Gpt::open(client).await.expect("load should succeed");
964        let mut transaction = manager.create_transaction().unwrap();
965        assert_eq!(transaction.partitions.len(), 1);
966        transaction.partitions[0] = crate::PartitionInfo {
967            label: PART_2_NAME.to_string(),
968            type_guid: crate::Guid::from_bytes(PART_TYPE_GUID),
969            instance_guid: crate::Guid::from_bytes(PART_INSTANCE_2_GUID),
970            start_block: 7,
971            num_blocks: 1,
972            flags: 0,
973        };
974        manager.commit_transaction(transaction).await.expect("Commit failed");
975
976        // Check state before and after a reload, to ensure both the in-memory and on-disk
977        // representation match.
978        assert_eq!(manager.header().num_parts, 1);
979        let partition = manager.partitions().get(&0).expect("No entry found");
980        assert_eq!(partition.label, PART_2_NAME);
981        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
982        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_2_GUID);
983        assert_eq!(partition.start_block, 7);
984        assert_eq!(partition.num_blocks, 1);
985        let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
986        assert_eq!(manager.header().num_parts, 1);
987        let partition = manager.partitions().get(&0).expect("No entry found");
988        assert_eq!(partition.label, PART_2_NAME);
989        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
990        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_2_GUID);
991        assert_eq!(partition.start_block, 7);
992        assert_eq!(partition.num_blocks, 1);
993        assert!(manager.partitions().get(&1).is_none());
994    }
995
996    #[fuchsia::test]
997    async fn grow_partition_table_in_transaction() {
998        let server =
999            VmoBackedServer::new(2048, 512, &[]).expect("Failed to create VmoBackedServer");
1000        let (client, _task) = connect_to_server(server).await;
1001        Gpt::format(
1002            client.clone(),
1003            vec![PartitionInfo {
1004                label: "part".to_string(),
1005                type_guid: Guid::from_bytes([1u8; 16]),
1006                instance_guid: Guid::from_bytes([1u8; 16]),
1007                start_block: 34,
1008                num_blocks: 1,
1009                flags: 0,
1010            }],
1011        )
1012        .await
1013        .expect("format failed");
1014        let mut manager = Gpt::open(client).await.expect("load should succeed");
1015        assert_eq!(manager.header().num_parts, 1);
1016        assert_eq!(manager.header().first_usable, 3);
1017        let mut transaction = manager.create_transaction().unwrap();
1018        transaction.partitions.resize(128, crate::PartitionInfo::nil());
1019        manager.commit_transaction(transaction).await.expect("Commit failed");
1020
1021        // Check state before and after a reload, to ensure both the in-memory and on-disk
1022        // representation match.
1023        assert_eq!(manager.header().num_parts, 128);
1024        assert_eq!(manager.header().first_usable, 34);
1025        let partition = manager.partitions().get(&0).expect("No entry found");
1026        assert_eq!(partition.label, "part");
1027        assert_eq!(partition.type_guid.to_bytes(), [1u8; 16]);
1028        assert_eq!(partition.instance_guid.to_bytes(), [1u8; 16]);
1029        assert_eq!(partition.start_block, 34);
1030        assert_eq!(partition.num_blocks, 1);
1031        assert!(manager.partitions().get(&1).is_none());
1032        let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
1033        assert_eq!(manager.header().num_parts, 128);
1034        assert_eq!(manager.header().first_usable, 34);
1035        let partition = manager.partitions().get(&0).expect("No entry found");
1036        assert_eq!(partition.label, "part");
1037        assert_eq!(partition.type_guid.to_bytes(), [1u8; 16]);
1038        assert_eq!(partition.instance_guid.to_bytes(), [1u8; 16]);
1039        assert_eq!(partition.start_block, 34);
1040        assert_eq!(partition.num_blocks, 1);
1041        assert!(manager.partitions().get(&1).is_none());
1042    }
1043
1044    #[fuchsia::test]
1045    async fn shrink_partition_table_in_transaction() {
1046        let mut partitions = vec![];
1047        for i in 0..128 {
1048            partitions.push(PartitionInfo {
1049                label: format!("part-{i}"),
1050                type_guid: Guid::from_bytes([i as u8 + 1; 16]),
1051                instance_guid: Guid::from_bytes([i as u8 + 1; 16]),
1052                start_block: 34 + i,
1053                num_blocks: 1,
1054                flags: 0,
1055            });
1056        }
1057        let server =
1058            VmoBackedServer::new(2048, 512, &[]).expect("Failed to create VmoBackedServer");
1059        let (client, _task) = connect_to_server(server).await;
1060        Gpt::format(client.clone(), partitions).await.expect("format failed");
1061        let mut manager = Gpt::open(client).await.expect("load should succeed");
1062        assert_eq!(manager.header().num_parts, 128);
1063        assert_eq!(manager.header().first_usable, 34);
1064        let mut transaction = manager.create_transaction().unwrap();
1065        transaction.partitions.clear();
1066        manager.commit_transaction(transaction).await.expect("Commit failed");
1067
1068        // Check state before and after a reload, to ensure both the in-memory and on-disk
1069        // representation match.
1070        assert_eq!(manager.header().num_parts, 0);
1071        assert_eq!(manager.header().first_usable, 2);
1072        assert!(manager.partitions().get(&0).is_none());
1073        let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
1074        assert_eq!(manager.header().num_parts, 0);
1075        assert_eq!(manager.header().first_usable, 2);
1076        assert!(manager.partitions().get(&0).is_none());
1077    }
1078
1079    #[fuchsia::test]
1080    async fn invalid_transaction_rejected() {
1081        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
1082        const PART_INSTANCE_GUID: [u8; 16] = [2u8; 16];
1083        const PART_NAME: &str = "part1";
1084
1085        let server = VmoBackedServer::new(16, 512, &[]).expect("Failed to create VmoBackedServer");
1086        let (client, _task) = connect_to_server(server).await;
1087        Gpt::format(
1088            client.clone(),
1089            vec![PartitionInfo {
1090                label: PART_NAME.to_string(),
1091                type_guid: Guid::from_bytes(PART_TYPE_GUID),
1092                instance_guid: Guid::from_bytes(PART_INSTANCE_GUID),
1093                start_block: 4,
1094                num_blocks: 1,
1095                flags: 0,
1096            }],
1097        )
1098        .await
1099        .expect("format failed");
1100        let mut manager = Gpt::open(client).await.expect("load should succeed");
1101        let mut transaction = manager.create_transaction().unwrap();
1102        assert_eq!(transaction.partitions.len(), 1);
1103        // This overlaps with the GPT metadata, so is invalid.
1104        transaction.partitions[0].start_block = 0;
1105        manager.commit_transaction(transaction).await.expect_err("Commit should have failed");
1106
1107        // Ensure nothing changed. Check state before and after a reload, to ensure both the
1108        // in-memory and on-disk representation match.
1109        assert_eq!(manager.header().num_parts, 1);
1110        let partition = manager.partitions().get(&0).expect("No entry found");
1111        assert_eq!(partition.label, PART_NAME);
1112        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
1113        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_GUID);
1114        assert_eq!(partition.start_block, 4);
1115        assert_eq!(partition.num_blocks, 1);
1116        let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
1117        assert_eq!(manager.header().num_parts, 1);
1118        let partition = manager.partitions().get(&0).expect("No entry found");
1119        assert_eq!(partition.label, PART_NAME);
1120        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
1121        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_GUID);
1122        assert_eq!(partition.start_block, 4);
1123        assert_eq!(partition.num_blocks, 1);
1124    }
1125
1126    /// An Observer that discards all writes overlapping its range (specified in bytes, not blocks).
1127    struct DiscardingObserver {
1128        block_size: u64,
1129        discard_range: Range<u64>,
1130    }
1131
1132    impl Observer for DiscardingObserver {
1133        fn write(
1134            &self,
1135            device_block_offset: u64,
1136            block_count: u32,
1137            _vmo: &Arc<zx::Vmo>,
1138            _vmo_offset: u64,
1139            _opts: block_server::WriteOptions,
1140        ) -> WriteAction {
1141            let write_range = (device_block_offset * self.block_size)
1142                ..(device_block_offset + block_count as u64) * self.block_size;
1143            if write_range.end <= self.discard_range.start
1144                || write_range.start >= self.discard_range.end
1145            {
1146                WriteAction::Write
1147            } else {
1148                WriteAction::Discard
1149            }
1150        }
1151    }
1152
1153    #[fuchsia::test]
1154    async fn transaction_applied_if_primary_metadata_partially_written() {
1155        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
1156        const PART_INSTANCE_1_GUID: [u8; 16] = [2u8; 16];
1157        const PART_INSTANCE_2_GUID: [u8; 16] = [3u8; 16];
1158        const PART_1_NAME: &str = "part1";
1159        const PART_2_NAME: &str = "part2";
1160
1161        let vmo = zx::Vmo::create(8192).unwrap();
1162        let server = VmoBackedServerOptions {
1163            initial_contents: InitialContents::FromVmo(vmo),
1164            block_size: 512,
1165            observer: Some(Box::new(DiscardingObserver {
1166                discard_range: 1024..1536,
1167                block_size: 512,
1168            })),
1169            ..Default::default()
1170        }
1171        .build()
1172        .unwrap();
1173        let (client, _task) = connect_to_server(server).await;
1174        Gpt::format(
1175            client.clone(),
1176            vec![PartitionInfo {
1177                label: PART_1_NAME.to_string(),
1178                type_guid: Guid::from_bytes(PART_TYPE_GUID),
1179                instance_guid: Guid::from_bytes(PART_INSTANCE_1_GUID),
1180                start_block: 4,
1181                num_blocks: 1,
1182                flags: 0,
1183            }],
1184        )
1185        .await
1186        .expect("format failed");
1187        let mut manager = Gpt::open(client).await.expect("load should succeed");
1188        let mut transaction = manager.create_transaction().unwrap();
1189        transaction.partitions.push(crate::PartitionInfo {
1190            label: PART_2_NAME.to_string(),
1191            type_guid: crate::Guid::from_bytes(PART_TYPE_GUID),
1192            instance_guid: crate::Guid::from_bytes(PART_INSTANCE_2_GUID),
1193            start_block: 7,
1194            num_blocks: 1,
1195            flags: 0,
1196        });
1197        manager.commit_transaction(transaction).await.expect("Commit failed");
1198
1199        let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
1200        assert_eq!(manager.header().num_parts, 2);
1201        let partition = manager.partitions().get(&0).expect("No entry found");
1202        assert_eq!(partition.label, PART_1_NAME);
1203        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
1204        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_1_GUID);
1205        assert_eq!(partition.start_block, 4);
1206        assert_eq!(partition.num_blocks, 1);
1207        let partition = manager.partitions().get(&1).expect("No entry found");
1208        assert_eq!(partition.label, PART_2_NAME);
1209        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
1210        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_2_GUID);
1211        assert_eq!(partition.start_block, 7);
1212        assert_eq!(partition.num_blocks, 1);
1213    }
1214
1215    #[fuchsia::test]
1216    async fn transaction_not_applied_if_primary_metadata_not_written() {
1217        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
1218        const PART_INSTANCE_1_GUID: [u8; 16] = [2u8; 16];
1219        const PART_INSTANCE_2_GUID: [u8; 16] = [3u8; 16];
1220        const PART_1_NAME: &str = "part1";
1221        const PART_2_NAME: &str = "part2";
1222
1223        let vmo = zx::Vmo::create(8192).unwrap();
1224        let vmo_dup = vmo.duplicate_handle(zx::Rights::SAME_RIGHTS).unwrap();
1225        {
1226            let server =
1227                VmoBackedServer::from_vmo(512, vmo_dup).expect("Failed to create VmoBackedServer");
1228            let (client, _task) = connect_to_server(server).await;
1229            Gpt::format(
1230                client.clone(),
1231                vec![PartitionInfo {
1232                    label: PART_1_NAME.to_string(),
1233                    type_guid: Guid::from_bytes(PART_TYPE_GUID),
1234                    instance_guid: Guid::from_bytes(PART_INSTANCE_1_GUID),
1235                    start_block: 4,
1236                    num_blocks: 1,
1237                    flags: 0,
1238                }],
1239            )
1240            .await
1241            .expect("format failed");
1242        }
1243        let server = VmoBackedServerOptions {
1244            initial_contents: InitialContents::FromVmo(vmo),
1245            block_size: 512,
1246            observer: Some(Box::new(DiscardingObserver {
1247                discard_range: 0..2048,
1248                block_size: 512,
1249            })),
1250            ..Default::default()
1251        }
1252        .build()
1253        .unwrap();
1254        let (client, _task) = connect_to_server(server).await;
1255
1256        let mut manager = Gpt::open(client).await.expect("load should succeed");
1257        let mut transaction = manager.create_transaction().unwrap();
1258        transaction.partitions.push(crate::PartitionInfo {
1259            label: PART_2_NAME.to_string(),
1260            type_guid: crate::Guid::from_bytes(PART_TYPE_GUID),
1261            instance_guid: crate::Guid::from_bytes(PART_INSTANCE_2_GUID),
1262            start_block: 7,
1263            num_blocks: 1,
1264            flags: 0,
1265        });
1266        manager.commit_transaction(transaction).await.expect("Commit failed");
1267
1268        let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
1269        assert_eq!(manager.header().num_parts, 1);
1270        let partition = manager.partitions().get(&0).expect("No entry found");
1271        assert_eq!(partition.label, PART_1_NAME);
1272        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
1273        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_1_GUID);
1274        assert_eq!(partition.start_block, 4);
1275        assert_eq!(partition.num_blocks, 1);
1276        assert!(manager.partitions().get(&1).is_none());
1277    }
1278
1279    #[fuchsia::test]
1280    async fn transaction_not_applied_if_backup_metadata_partially_written() {
1281        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
1282        const PART_INSTANCE_1_GUID: [u8; 16] = [2u8; 16];
1283        const PART_INSTANCE_2_GUID: [u8; 16] = [3u8; 16];
1284        const PART_1_NAME: &str = "part1";
1285        const PART_2_NAME: &str = "part2";
1286
1287        let vmo = zx::Vmo::create(8192).unwrap();
1288        let vmo_dup = vmo.duplicate_handle(zx::Rights::SAME_RIGHTS).unwrap();
1289        {
1290            let server =
1291                VmoBackedServer::from_vmo(512, vmo_dup).expect("Failed to create VmoBackedServer");
1292            let (client, _task) = connect_to_server(server).await;
1293            Gpt::format(
1294                client.clone(),
1295                vec![PartitionInfo {
1296                    label: PART_1_NAME.to_string(),
1297                    type_guid: Guid::from_bytes(PART_TYPE_GUID),
1298                    instance_guid: Guid::from_bytes(PART_INSTANCE_1_GUID),
1299                    start_block: 4,
1300                    num_blocks: 1,
1301                    flags: 0,
1302                }],
1303            )
1304            .await
1305            .expect("format failed");
1306        }
1307        let server = VmoBackedServerOptions {
1308            initial_contents: InitialContents::FromVmo(vmo),
1309            block_size: 512,
1310            observer: Some(Box::new(DiscardingObserver {
1311                discard_range: 0..7680,
1312                block_size: 512,
1313            })),
1314            ..Default::default()
1315        }
1316        .build()
1317        .unwrap();
1318        let (client, _task) = connect_to_server(server).await;
1319
1320        let mut manager = Gpt::open(client).await.expect("load should succeed");
1321        let mut transaction = manager.create_transaction().unwrap();
1322        transaction.partitions.push(crate::PartitionInfo {
1323            label: PART_2_NAME.to_string(),
1324            type_guid: crate::Guid::from_bytes(PART_TYPE_GUID),
1325            instance_guid: crate::Guid::from_bytes(PART_INSTANCE_2_GUID),
1326            start_block: 7,
1327            num_blocks: 1,
1328            flags: 0,
1329        });
1330        manager.commit_transaction(transaction).await.expect("Commit failed");
1331
1332        let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
1333        assert_eq!(manager.header().num_parts, 1);
1334        let partition = manager.partitions().get(&0).expect("No entry found");
1335        assert_eq!(partition.label, PART_1_NAME);
1336        assert_eq!(partition.type_guid.to_bytes(), PART_TYPE_GUID);
1337        assert_eq!(partition.instance_guid.to_bytes(), PART_INSTANCE_1_GUID);
1338        assert_eq!(partition.start_block, 4);
1339        assert_eq!(partition.num_blocks, 1);
1340        assert!(manager.partitions().get(&1).is_none());
1341    }
1342
1343    #[fuchsia::test]
1344    async fn restore_primary_from_backup() {
1345        const PART_TYPE_GUID: [u8; 16] = [2u8; 16];
1346        const PART_INSTANCE_GUID: [u8; 16] = [2u8; 16];
1347        const PART_NAME: &str = "part1";
1348
1349        let server = VmoBackedServer::new(16, 512, &[]).expect("Failed to create VmoBackedServer");
1350        let (client, _task) = connect_to_server(server).await;
1351        Gpt::format(
1352            client.clone(),
1353            vec![PartitionInfo {
1354                label: PART_NAME.to_string(),
1355                type_guid: Guid::from_bytes(PART_TYPE_GUID),
1356                instance_guid: Guid::from_bytes(PART_INSTANCE_GUID),
1357                start_block: 4,
1358                num_blocks: 1,
1359                flags: 0,
1360            }],
1361        )
1362        .await
1363        .expect("format failed");
1364        let mut old_metadata = vec![0u8; 2048];
1365        client.read_at(MutableBufferSlice::Memory(&mut old_metadata[..]), 0).await.unwrap();
1366        let mut buffer = vec![0u8; 2048];
1367        client.write_at(BufferSlice::Memory(&buffer[..]), 0).await.unwrap();
1368
1369        let manager = Gpt::open(client).await.expect("load should succeed");
1370        let client = manager.take_client();
1371
1372        client.read_at(MutableBufferSlice::Memory(&mut buffer[..]), 0).await.unwrap();
1373        assert_eq!(old_metadata, buffer);
1374    }
1375
1376    #[fuchsia::test]
1377    async fn load_golden_gpt_linux() {
1378        let server = VmoBackedServer::from_file(512, "/pkg/data/gpt_golden/gpt.linux.blk");
1379        let (client, _task) = connect_to_server(server).await;
1380        let manager = Gpt::open(client).await.expect("load should succeed");
1381        let partition = manager.partitions().get(&0).expect("No entry found");
1382        assert_eq!(partition.label, "ext");
1383        assert_eq!(partition.type_guid.to_string(), "0fc63daf-8483-4772-8e79-3d69d8477de4");
1384        assert_eq!(partition.start_block, 8);
1385        assert_eq!(partition.num_blocks, 1);
1386        assert!(manager.partitions().get(&1).is_none());
1387    }
1388
1389    #[fuchsia::test]
1390    async fn load_golden_gpt_fuchsia() {
1391        let server = VmoBackedServer::from_file(512, "/pkg/data/gpt_golden/gpt.fuchsia.blk");
1392        let (client, _task) = connect_to_server(server).await;
1393
1394        struct ExpectedPartition {
1395            label: &'static str,
1396            type_guid: &'static str,
1397            blocks: Range<u64>,
1398        }
1399        const EXPECTED_PARTITIONS: [ExpectedPartition; 8] = [
1400            ExpectedPartition {
1401                label: "bootloader",
1402                type_guid: "5ece94fe-4c86-11e8-a15b-480fcf35f8e6",
1403                blocks: 11..12,
1404            },
1405            ExpectedPartition {
1406                label: "zircon_a",
1407                type_guid: "9b37fff6-2e58-466a-983a-f7926d0b04e0",
1408                blocks: 12..13,
1409            },
1410            ExpectedPartition {
1411                label: "zircon_b",
1412                type_guid: "9b37fff6-2e58-466a-983a-f7926d0b04e0",
1413                blocks: 13..14,
1414            },
1415            ExpectedPartition {
1416                label: "zircon_r",
1417                type_guid: "9b37fff6-2e58-466a-983a-f7926d0b04e0",
1418                blocks: 14..15,
1419            },
1420            ExpectedPartition {
1421                label: "vbmeta_a",
1422                type_guid: "421a8bfc-85d9-4d85-acda-b64eec0133e9",
1423                blocks: 15..16,
1424            },
1425            ExpectedPartition {
1426                label: "vbmeta_b",
1427                type_guid: "421a8bfc-85d9-4d85-acda-b64eec0133e9",
1428                blocks: 16..17,
1429            },
1430            ExpectedPartition {
1431                label: "vbmeta_r",
1432                type_guid: "421a8bfc-85d9-4d85-acda-b64eec0133e9",
1433                blocks: 17..18,
1434            },
1435            ExpectedPartition {
1436                label: "durable_boot",
1437                type_guid: "a409e16b-78aa-4acc-995c-302352621a41",
1438                blocks: 18..19,
1439            },
1440        ];
1441
1442        let manager = Gpt::open(client).await.expect("load should succeed");
1443        for i in 0..EXPECTED_PARTITIONS.len() as u32 {
1444            let partition = manager.partitions().get(&i).expect("No entry found");
1445            let expected = &EXPECTED_PARTITIONS[i as usize];
1446            assert_eq!(partition.label, expected.label);
1447            assert_eq!(partition.type_guid.to_string(), expected.type_guid);
1448            assert_eq!(partition.start_block, expected.blocks.start);
1449            assert_eq!(partition.num_blocks, expected.blocks.end - expected.blocks.start);
1450        }
1451    }
1452
1453    #[fuchsia::test]
1454    async fn add_partitions_till_no_blocks_left() {
1455        let server = VmoBackedServer::new(128, 512, &[]).expect("Failed to create VmoBackedServer");
1456        let (client, _task) = connect_to_server(server).await;
1457        Gpt::format(client.clone(), vec![PartitionInfo::nil(); 32]).await.expect("format failed");
1458        let mut manager = Gpt::open(client).await.expect("load should succeed");
1459        let mut transaction = manager.create_transaction().unwrap();
1460        assert_eq!(transaction.partitions.len(), 32);
1461        let mut num = 0;
1462        loop {
1463            match manager.add_partition(
1464                &mut transaction,
1465                crate::PartitionInfo {
1466                    label: format!("part-{num}"),
1467                    type_guid: crate::Guid::generate(),
1468                    instance_guid: crate::Guid::generate(),
1469                    start_block: 0,
1470                    num_blocks: 1,
1471                    flags: 0,
1472                },
1473            ) {
1474                Ok(_) => {
1475                    num += 1;
1476                }
1477                Err(AddPartitionError::InvalidArguments) => panic!("Unexpected error"),
1478                Err(AddPartitionError::NoSpace) => break,
1479            };
1480        }
1481        assert!(num <= 32);
1482        manager.commit_transaction(transaction).await.expect("Commit failed");
1483
1484        // Check state before and after a reload, to ensure both the in-memory and on-disk
1485        // representation match.
1486        assert_eq!(manager.header().num_parts, 32);
1487        assert_eq!(manager.partitions().len(), num);
1488
1489        let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
1490        assert_eq!(manager.header().num_parts, 32);
1491        assert_eq!(manager.partitions().len(), num);
1492    }
1493
1494    #[fuchsia::test]
1495    async fn add_partitions_till_no_slots_left() {
1496        let server = VmoBackedServer::new(128, 512, &[]).expect("Failed to create VmoBackedServer");
1497        let (client, _task) = connect_to_server(server).await;
1498        Gpt::format(client.clone(), vec![PartitionInfo::nil(); 4]).await.expect("format failed");
1499        let mut manager = Gpt::open(client).await.expect("load should succeed");
1500        let mut transaction = manager.create_transaction().unwrap();
1501        assert_eq!(transaction.partitions.len(), 4);
1502        let mut num = 0;
1503        loop {
1504            match manager.add_partition(
1505                &mut transaction,
1506                crate::PartitionInfo {
1507                    label: format!("part-{num}"),
1508                    type_guid: crate::Guid::generate(),
1509                    instance_guid: crate::Guid::generate(),
1510                    start_block: 0,
1511                    num_blocks: 1,
1512                    flags: 0,
1513                },
1514            ) {
1515                Ok(_) => {
1516                    num += 1;
1517                }
1518                Err(AddPartitionError::InvalidArguments) => panic!("Unexpected error"),
1519                Err(AddPartitionError::NoSpace) => break,
1520            };
1521        }
1522        assert!(num <= 4);
1523        manager.commit_transaction(transaction).await.expect("Commit failed");
1524
1525        // Check state before and after a reload, to ensure both the in-memory and on-disk
1526        // representation match.
1527        assert_eq!(manager.header().num_parts, 4);
1528        assert_eq!(manager.partitions().len(), num);
1529
1530        let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
1531        assert_eq!(manager.header().num_parts, 4);
1532        assert_eq!(manager.partitions().len(), num);
1533    }
1534
1535    /// An Observer that shuffles writes and discards some of the tail since last flush.
1536    struct ShufflingObserver {
1537        // Only start shuffling once this is set.
1538        start: Arc<AtomicBool>,
1539        // Only shuffle if there is a write to this offset.
1540        shuffle_if_contains_offset: u64,
1541    }
1542
1543    impl Observer for ShufflingObserver {
1544        fn flush(&self, writes: Option<&mut WriteCache>) {
1545            if self.start.load(Ordering::Relaxed) {
1546                let Some(writes) = writes else { unreachable!() };
1547                if writes
1548                    .iter()
1549                    .filter(|(offset, _)| **offset == self.shuffle_if_contains_offset)
1550                    .next()
1551                    .is_some()
1552                {
1553                    writes.shuffle();
1554                    writes.discard_some();
1555                }
1556            }
1557        }
1558
1559        fn close(&self, writes: Option<&mut WriteCache>) {
1560            // Always shuffle every write which had yet to be flushed when the client closed.
1561            if self.start.load(Ordering::Relaxed) {
1562                let Some(writes) = writes else { unreachable!() };
1563                writes.shuffle();
1564            }
1565        }
1566    }
1567
1568    #[fuchsia::test]
1569    async fn metadata_update_is_atomic() {
1570        const BLOCK_SIZE: u64 = 512;
1571        const BLOCK_COUNT: u64 = 128;
1572        // Test once where we shuffle any set of writes which contains the primary superblock, and
1573        // once where we shuffle any set of writes which contains the secondary superblock.
1574        // The goal is to ensure that writes are correctly sequenced with some sort of flush or
1575        // barrier (secondary, <barrier>, primary), so metadata updates are atomic.
1576        for shuffle_if_contains_offset in [1, BLOCK_COUNT - 1] {
1577            let vmo = zx::Vmo::create(BLOCK_SIZE * BLOCK_COUNT).unwrap();
1578            let start_shuffling = Arc::new(AtomicBool::new(false));
1579            let server = VmoBackedServerOptions {
1580                initial_contents: InitialContents::FromVmo(vmo),
1581                block_size: BLOCK_SIZE as u32,
1582                observer: Some(Box::new(ShufflingObserver {
1583                    start: start_shuffling.clone(),
1584                    shuffle_if_contains_offset,
1585                })),
1586                write_tracking: true,
1587                ..Default::default()
1588            }
1589            .build()
1590            .unwrap();
1591            let (client, _task) = connect_to_server(server).await;
1592            Gpt::format(client.clone(), vec![PartitionInfo::nil(); 80])
1593                .await
1594                .expect("format failed");
1595
1596            start_shuffling.store(true, Ordering::Relaxed);
1597
1598            let mut manager = Gpt::open(client).await.expect("load should succeed");
1599            let mut transaction = manager.create_transaction().unwrap();
1600            transaction.partitions.truncate(40);
1601            let mut num = 0;
1602            loop {
1603                match manager.add_partition(
1604                    &mut transaction,
1605                    crate::PartitionInfo {
1606                        label: format!("part-{num}"),
1607                        type_guid: crate::Guid::generate(),
1608                        instance_guid: crate::Guid::generate(),
1609                        start_block: 0,
1610                        num_blocks: 1,
1611                        flags: 0,
1612                    },
1613                ) {
1614                    Ok(_) => {
1615                        num += 1;
1616                    }
1617                    Err(AddPartitionError::InvalidArguments) => panic!("Unexpected error"),
1618                    Err(AddPartitionError::NoSpace) => break,
1619                };
1620            }
1621            assert!(num <= 40);
1622            manager.commit_transaction(transaction).await.expect("Commit failed");
1623
1624            // Check state before and after a reload.
1625            assert_eq!(manager.header().num_parts, 40);
1626            assert_eq!(manager.partitions().len(), num);
1627
1628            // If the GPT implementation has appropriate barriers/flushes between secondary and
1629            // primary metadata updates, then we will end up in either the old state or the new
1630            // state.  Otherwise, both copies might become corrupt and the GPT would be unreadable.
1631            let manager = Gpt::open(manager.take_client()).await.expect("reload should succeed");
1632            let len = manager.partitions().len();
1633            assert!(len == 0 || len == num);
1634        }
1635    }
1636}