1 // Copyright 2023, 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 //! Support for parsing GUID partition tables.
16 
17 use core::cmp::min;
18 use core::fmt;
19 use core::mem::size_of;
20 use core::ops::RangeInclusive;
21 use core::slice;
22 use static_assertions::const_assert;
23 use static_assertions::const_assert_eq;
24 use uuid::Uuid;
25 use virtio_drivers::device::blk::SECTOR_SIZE;
26 use vmbase::util::ceiling_div;
27 use vmbase::virtio::{pci, HalImpl};
28 use zerocopy::FromBytes;
29 use zerocopy::FromZeroes;
30 
31 type VirtIOBlk = pci::VirtIOBlk<HalImpl>;
32 
33 pub enum Error {
34     /// VirtIO error during read operation.
35     FailedRead(virtio_drivers::Error),
36     /// VirtIO error during write operation.
37     FailedWrite(virtio_drivers::Error),
38     /// Invalid GPT header.
39     InvalidHeader,
40     /// Invalid partition block index.
41     BlockOutsidePartition(usize),
42 }
43 
44 impl fmt::Display for Error {
fmt(&self, f: &mut fmt::Formatter) -> fmt::Result45     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
46         match self {
47             Self::FailedRead(e) => write!(f, "Failed to read from disk: {e}"),
48             Self::FailedWrite(e) => write!(f, "Failed to write to disk: {e}"),
49             Self::InvalidHeader => write!(f, "Found invalid GPT header"),
50             Self::BlockOutsidePartition(i) => write!(f, "Accessed invalid block index {i}"),
51         }
52     }
53 }
54 
55 pub type Result<T> = core::result::Result<T, Error>;
56 
57 pub struct Partition {
58     partitions: Partitions,
59     indices: RangeInclusive<usize>,
60 }
61 
62 impl Partition {
get_by_name(device: VirtIOBlk, name: &str) -> Result<Option<Self>>63     pub fn get_by_name(device: VirtIOBlk, name: &str) -> Result<Option<Self>> {
64         Partitions::new(device)?.get_partition_by_name(name)
65     }
66 
new(partitions: Partitions, entry: &Entry) -> Self67     fn new(partitions: Partitions, entry: &Entry) -> Self {
68         let first = entry.first_lba().try_into().unwrap();
69         let last = entry.last_lba().try_into().unwrap();
70 
71         Self { partitions, indices: first..=last }
72     }
73 
indices(&self) -> RangeInclusive<usize>74     pub fn indices(&self) -> RangeInclusive<usize> {
75         self.indices.clone()
76     }
77 
read_block(&mut self, index: usize, blk: &mut [u8]) -> Result<()>78     pub fn read_block(&mut self, index: usize, blk: &mut [u8]) -> Result<()> {
79         let index = self.block_index(index).ok_or(Error::BlockOutsidePartition(index))?;
80         self.partitions.read_block(index, blk)
81     }
82 
write_block(&mut self, index: usize, blk: &[u8]) -> Result<()>83     pub fn write_block(&mut self, index: usize, blk: &[u8]) -> Result<()> {
84         let index = self.block_index(index).ok_or(Error::BlockOutsidePartition(index))?;
85         self.partitions.write_block(index, blk)
86     }
87 
block_index(&self, index: usize) -> Option<usize>88     fn block_index(&self, index: usize) -> Option<usize> {
89         if self.indices.contains(&index) {
90             Some(index)
91         } else {
92             None
93         }
94     }
95 }
96 
97 pub struct Partitions {
98     device: VirtIOBlk,
99     entries_count: usize,
100 }
101 
102 impl Partitions {
103     pub const LBA_SIZE: usize = SECTOR_SIZE;
104 
new(mut device: VirtIOBlk) -> Result<Self>105     fn new(mut device: VirtIOBlk) -> Result<Self> {
106         let mut blk = [0; Self::LBA_SIZE];
107         device.read_blocks(Header::LBA, &mut blk).map_err(Error::FailedRead)?;
108         let header = Header::read_from_prefix(blk.as_slice()).unwrap();
109         if !header.is_valid() {
110             return Err(Error::InvalidHeader);
111         }
112         let entries_count = usize::try_from(header.entries_count()).unwrap();
113 
114         Ok(Self { device, entries_count })
115     }
116 
get_partition_by_name(mut self, name: &str) -> Result<Option<Partition>>117     fn get_partition_by_name(mut self, name: &str) -> Result<Option<Partition>> {
118         const_assert_eq!(Partitions::LBA_SIZE.rem_euclid(size_of::<Entry>()), 0);
119         let entries_per_blk = Partitions::LBA_SIZE.checked_div(size_of::<Entry>()).unwrap();
120 
121         // Create a UTF-16 reference against which we'll compare partition names. Note that unlike
122         // the C99 wcslen(), this comparison will cover bytes past the first L'\0' character.
123         let mut needle = [0; Entry::NAME_SIZE / size_of::<u16>()];
124         for (dest, src) in needle.iter_mut().zip(name.encode_utf16()) {
125             *dest = src;
126         }
127 
128         let mut blk = [0; Self::LBA_SIZE];
129         let mut rem = self.entries_count;
130         let num_blocks = ceiling_div(self.entries_count, entries_per_blk).unwrap();
131         for i in Header::ENTRIES_LBA..Header::ENTRIES_LBA.checked_add(num_blocks).unwrap() {
132             self.read_block(i, &mut blk)?;
133             let entries = blk.as_ptr().cast::<Entry>();
134             // SAFETY: blk is assumed to be properly aligned for Entry and its size is assert-ed
135             // above. All potential values of the slice will produce valid Entry values.
136             let entries = unsafe { slice::from_raw_parts(entries, min(rem, entries_per_blk)) };
137             for entry in entries {
138                 let entry_name = entry.name;
139                 if entry_name == needle {
140                     return Ok(Some(Partition::new(self, entry)));
141                 }
142                 rem -= 1;
143             }
144         }
145         Ok(None)
146     }
147 
read_block(&mut self, index: usize, blk: &mut [u8]) -> Result<()>148     fn read_block(&mut self, index: usize, blk: &mut [u8]) -> Result<()> {
149         self.device.read_blocks(index, blk).map_err(Error::FailedRead)
150     }
151 
write_block(&mut self, index: usize, blk: &[u8]) -> Result<()>152     fn write_block(&mut self, index: usize, blk: &[u8]) -> Result<()> {
153         self.device.write_blocks(index, blk).map_err(Error::FailedWrite)
154     }
155 }
156 
157 type Lba = u64;
158 
159 /// Structure as defined in release 2.10 of the UEFI Specification (5.3.2 GPT Header).
160 #[derive(FromZeroes, FromBytes)]
161 #[repr(C, packed)]
162 struct Header {
163     signature: u64,
164     revision: u32,
165     header_size: u32,
166     header_crc32: u32,
167     reserved0: u32,
168     current_lba: Lba,
169     backup_lba: Lba,
170     first_lba: Lba,
171     last_lba: Lba,
172     disk_guid: u128,
173     entries_lba: Lba,
174     entries_count: u32,
175     entry_size: u32,
176     entries_crc32: u32,
177 }
178 const_assert!(size_of::<Header>() < Partitions::LBA_SIZE);
179 
180 impl Header {
181     const SIGNATURE: u64 = u64::from_le_bytes(*b"EFI PART");
182     const REVISION_1_0: u32 = 1 << 16;
183     const LBA: usize = 1;
184     const ENTRIES_LBA: usize = 2;
185 
is_valid(&self) -> bool186     fn is_valid(&self) -> bool {
187         self.signature() == Self::SIGNATURE
188             && self.header_size() == size_of::<Self>().try_into().unwrap()
189             && self.revision() == Self::REVISION_1_0
190             && self.entry_size() == size_of::<Entry>().try_into().unwrap()
191             && self.current_lba() == Self::LBA.try_into().unwrap()
192             && self.entries_lba() == Self::ENTRIES_LBA.try_into().unwrap()
193     }
194 
signature(&self) -> u64195     fn signature(&self) -> u64 {
196         u64::from_le(self.signature)
197     }
198 
entries_count(&self) -> u32199     fn entries_count(&self) -> u32 {
200         u32::from_le(self.entries_count)
201     }
202 
header_size(&self) -> u32203     fn header_size(&self) -> u32 {
204         u32::from_le(self.header_size)
205     }
206 
revision(&self) -> u32207     fn revision(&self) -> u32 {
208         u32::from_le(self.revision)
209     }
210 
entry_size(&self) -> u32211     fn entry_size(&self) -> u32 {
212         u32::from_le(self.entry_size)
213     }
214 
entries_lba(&self) -> Lba215     fn entries_lba(&self) -> Lba {
216         Lba::from_le(self.entries_lba)
217     }
218 
current_lba(&self) -> Lba219     fn current_lba(&self) -> Lba {
220         Lba::from_le(self.current_lba)
221     }
222 }
223 
224 /// Structure as defined in release 2.10 of the UEFI Specification (5.3.3 GPT Partition Entry
225 /// Array).
226 #[repr(C, packed)]
227 struct Entry {
228     type_guid: Uuid,
229     guid: Uuid,
230     first_lba: Lba,
231     last_lba: Lba,
232     flags: u64,
233     name: [u16; Entry::NAME_SIZE / size_of::<u16>()], // UTF-16
234 }
235 
236 impl Entry {
237     const NAME_SIZE: usize = 72;
238 
first_lba(&self) -> Lba239     fn first_lba(&self) -> Lba {
240         Lba::from_le(self.first_lba)
241     }
242 
last_lba(&self) -> Lba243     fn last_lba(&self) -> Lba {
244         Lba::from_le(self.last_lba)
245     }
246 }
247