1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
// Copyright 2022 Alberto Ruiz <aruiz@redhat.com>
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

#![cfg_attr(feature = "no_std", no_std)]

use crc::Crc;

#[cfg(feature = "no_std")]
use core::{mem, ops};

#[cfg(not(feature = "no_std"))]
use std::{mem, ops};

use ops::{Index, Range};

use mem::size_of;

mod endian;

use endian::u16le;
use endian::u32le;
use endian::u64le;

#[repr(C, align(1))]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Uuid {
    time_low: u32le,
    time_mid: u16le,
    time_hi_and_version: u16le,
    clock_seq_hi_and_reseved: u8,
    clock_seq_loq: u8,
    node: [u8; 6],
}

impl Uuid {
    pub fn nil() -> Self {
        Self {
            time_low: 0.into(),
            time_mid: 0.into(),
            time_hi_and_version: 0.into(),
            clock_seq_hi_and_reseved: 0,
            clock_seq_loq: 0,
            node: [0, 0, 0, 0, 0, 0],
        }
    }
}

#[repr(C, align(1))]
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub struct GPTPartition {
    part_type: Uuid,
    id: Uuid,
    first_lba: u64le,
    last_lba: u64le,
    attr: [u8; 8],
    name: [u16le; 36],
}

#[repr(C, align(1))]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct GPTHeader {
    /// GPT header magic signature, hardcoded to "EFI PART".
    pub signature: [u8; 8], // Offset  0. "EFI PART", 45h 46h 49h 20h 50h 41h 52h 54h
    /// 00 00 01 00
    pub revision: u32le, // Offset  8
    /// little endian
    pub header_size_le: u32le, // Offset 12
    /// CRC32 of the header with crc32 section zeroed
    pub header_crc32: u32le, // Offset 16
    /// must be 0
    pub reserved: u32le, // Offset 20
    /// For main header, 1
    pub current_lba: u64le, // Offset 24
    /// LBA for backup header
    pub backup_lba: u64le, // Offset 32
    /// First usable LBA for partitions (primary table last LBA + 1)
    pub first_usable: u64le, // Offset 40
    /// Last usable LBA (secondary partition table first LBA - 1)
    pub last_usable: u64le, // Offset 48
    /// UUID of the disk
    pub disk_guid: Uuid, // Offset 56
    /// Starting LBA of partition entries
    pub part_start_lba: u64le, // Offset 72
    /// Number of partition entries
    pub num_parts: u32le, // Offset 80
    /// Size of a partition entry, usually 128
    pub part_size: u32le, // Offset 84
    /// CRC32 of the partition table
    pub part_table_crc32: u32le, // Offset 88
                                 // RESERVED: [u8; 96 - 512],
}

impl GPTPartition {
    pub unsafe fn partition_table_checksum(table: &[GPTPartition]) -> u32 {
        let table: &[u8] = mem::transmute(table);
        let hasher = Crc::<u32>::new(&crc::CRC_32_JAMCRC);
        hasher.checksum(table) ^ 0xffffffff
    }
}

#[repr(C)]
#[derive(Debug, Eq, PartialEq)]
pub struct LBA<T: Sized + LBAIndex>(T);

#[repr(C)]
pub struct LBA512([u8; 512]);

#[repr(C)]
pub struct LBA4K([u8; 4096]);

pub trait LBAIndex {
    fn index<'a>(&'a self, slice: Range<usize>) -> &'a [u8];
}

impl LBAIndex for LBA512 {
    fn index<'a>(&'a self, slice: Range<usize>) -> &'a [u8] {
        &self.0[slice]
    }
}

impl LBAIndex for LBA4K {
    fn index<'a>(&'a self, slice: Range<usize>) -> &'a [u8] {
        &self.0[slice]
    }
}

impl<T: LBAIndex> Index<Range<usize>> for LBA<T> {
    type Output = [u8];
    fn index<'a>(&'a self, slice: Range<usize>) -> &'a [u8] {
        &self.0.index(slice)
    }
}

impl GPTHeader {
    pub fn verify_signature(&self) -> bool {
        const SIGNATURE: [u8; 8] = *b"EFI PART";
        self.signature == SIGNATURE
    }

    pub unsafe fn parse_gpt_header<'a, T: Sized + LBAIndex>(
        gpt_header_lba: &'a LBA<T>,
    ) -> Result<&'a Self, &'static str> {
        if size_of::<LBA<T>>() > isize::MAX as usize {
            return Err("Size of LBA block larger than the 32bit address space");
        }
        if size_of::<LBA<T>>() < 512 {
            return Err("Size of LBA smaller than 512 bytes");
        }

        let header = (gpt_header_lba as *const LBA<T>) as *const GPTHeader;
        let header = &(*header) as &GPTHeader;
        if !header.verify_signature() {
            return Err("Could not verify GPT header signature `EFI PART`");
        }

        let part_header_size = u32::from(header.part_size) as usize;

        // Partition header size has to be 2^x * 128
        if part_header_size % 128 != 0 {
            return Err("Error in GPT table: Partition size not a multiple of 128");
        }
        if (part_header_size / 128).count_ones() != 1 {
            return Err("Error in GPT table: Partition size not a power of 2");
        }

        // Verify header
        assert_eq!(
            Self::header_checksum::<T>(&(*header)),
            u32::from(header.header_crc32),
            "GPT header CRC32 verification failed"
        );

        //FIXME: Check if there are enough LBAs to fit all partitions

        Ok(header) // Figure out if we can make this a reference instead of a copy
    }

    pub fn parse_partitions<F: Fn(&GPTPartition, &mut T), T>(
        part_array: &[GPTPartition],
        size: usize,
        foreach: F,
        data: &mut T,
    ) {
        //FIXME: Verify that part_header_size * num_parts does not overfloy beyond first_usable LBA
        let nil_uuid = Uuid::nil();

        // Size comes from the GPT Header and it is always a multiple of 128 bytes
        // which is the size of GPTPartition, we use it as a stride in our array
        let stride = size / size_of::<GPTPartition>();

        for i in 0..part_array.len() {
            if i % stride != 0 {
                continue;
            }
            if part_array[i].part_type == nil_uuid {
                continue;
            }
            foreach(&part_array[i], data);
        }
    }

    fn header_checksum<T: LBAIndex>(header: &GPTHeader) -> u32 {
        let verify_size = u32::from(header.header_size_le) as usize;
        let mut verify_header: GPTHeader = *header;
        let v_header_block = (&verify_header as *const GPTHeader) as *const LBA<T>;
        verify_header.header_crc32 = 0.into();
        let hasher = Crc::<u32>::new(&crc::CRC_32_JAMCRC);
        (unsafe { hasher.checksum(&(*v_header_block)[0..verify_size]) }) ^ 0xffffffff
    }
}

#[cfg(test)]
mod tests {
    extern crate base64;
    extern crate flate2;
    use crate::{GPTHeader, GPTPartition};
    use crate::{LBAIndex, LBA, LBA4K, LBA512};

    use std::io::Cursor;
    use std::io::Read;
    use std::vec::Vec;

    // This is a base64 encoded gzipped 100K disk image formatted for 512 bytes block size and two partitions
    const GPT_EXAMPLE1: &[u8] = b"H4sIAAAAAAAAA+3cPUgjQRQA4E1AOLAQezFi5d9ZWQgSwSKGEAsJCRwStBTFyoBgIQQMlnZqqVUOQtA+XRolVpJKjlME7YSzsEkjuQib9gqPHCd8Hwxv583beQNT7wYBn1k0+NVutyOdp6sPvJ2rJBZTI8sLmWwQRIJ8J/Pl549v7yuRsKK762gYy2GcfDx6fsilosX+6+HXfDoWDfPFcIyvr7x84ED8Yxfxy4H90nbycDe+cZvce8pOn9xUWom30/O72Xpiqtq999Ue9R8rTNRjreZgLR2cNeYbczsHhWzfZnroeGareZ9Zqq6FdeU/7gIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACfy0X8cmC/tJ083I1v3Cb3nrLTJzeVVuLt9Pxutp6Yqo6Gdas96j9WmKjHWs3BWjo4a8w35nYOCtm+zfTQ8cxW8z6zVF0L68o96g8AAAAAAAAAAAAAAAAAAAD/g8RiamR5IZMNgkiQ78wzX2ul9/xVuB4JY/c/AN3v8Ccfj54fcqlosf96+DWfjn0P88VwjK+vvPT+9Pyt3wAlEswAkAEA";

    // This is a base64 encoded gzipped 100K disk image formatted for 4Kbytes block size and two partitions
    const GPT_EXAMPLE_LBA4K: &[u8] = b"H4sICFR+fWIAA2xiYTRrLmltZwDt3L9LI0EUAOBdixSJaJErrA57aysJomBiCvEHURCsBRsRGwmiaLS4SrCwCoiV4DVaCEGwE2z8B0RMozYipFYQ1giTVkQu4sH3wTA7b97se+0Uu1HE/6wjaiRJEjefer5wevrvv+4HAACAn2kkX+ydGJoqRVEczTXXS5U/pfd4HPZbt8pUmLNhPuq8WV3pH62dZV6S58b+cEeIb4RxUdvda3/3AAAAwGcc5y67t7aXCzvl3MJ1Ye3hZOBqMb0+/npe3ZypHk7ete796W+qP1jPz9/3jc2edlXq5dTjQSbkZT98CwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAx7nL7q3t5cJOObdwXVh7OBm4Wkyvj7+eVzdnqoeTd6mQl/6m+oP1/Px939jsaVelXk49HmRCXrZN9QEAAAAAAAAAAAAAAAAAAOAnGMkXeyeGpkpRFEdzzfVt8en3e7wn7Mdhbv0HoPUd/lHnzepK/2jtLPOSPDf2h3+F+EYYF7XdvfZ3DwAAAHzGGwzc9x8AkAEA";

    fn test_disk_image<T: LBAIndex>(image: &[u8], uuid_time_low_1: u32, uuid_time_low_2: u32) {
        let mut base_decoded: Vec<u8> = Vec::new();
        let mut cur = Cursor::new(image);
        assert!(
            base64::read::DecoderReader::new(&mut cur, base64::STANDARD)
                .read_to_end(&mut base_decoded)
                .is_ok(),
            "Could not decode base64 blob"
        );
        let mut disk_data: Vec<u8> = Vec::new();
        assert!(
            flate2::read::GzDecoder::new(&base_decoded[..])
                .read_to_end(&mut disk_data)
                .is_ok(),
            "flat2 could not gunzip disk blob"
        );

        // We hardcode LBA512 for the disk, typically you'd get this from the device
        let mbr_lba = disk_data.as_ptr() as *const LBA<T>;

        let gpt_header = unsafe { GPTHeader::parse_gpt_header(&*mbr_lba.add(1)) };
        assert!(
            gpt_header.is_ok(),
            "Could not parse GPT Header: {}",
            gpt_header.unwrap_err()
        );
        let gpt_header = gpt_header.unwrap();

        // This is the offset of LBA blocks starting from the GPT Header LBA
        let partition_lba_offset = u64::from(gpt_header.part_start_lba) as usize;
        let partition_lba_ptr = unsafe { mbr_lba.add(partition_lba_offset) as *const GPTPartition };

        let part_size = u32::from(gpt_header.part_size) as usize;
        let num_parts = u32::from(gpt_header.num_parts) as usize;

        assert!(
            disk_data.len() > (std::mem::size_of::<LBA<T>>() * 2 + (part_size * num_parts)),
            "Partition table is bigger than buffer size"
        );

        let stride = part_size / std::mem::size_of::<GPTPartition>();
        let length = stride * num_parts;
        let part_array_ptr = core::ptr::slice_from_raw_parts(partition_lba_ptr, length);
        let part_array = unsafe { &*part_array_ptr };

        let checksum_array =
            core::ptr::slice_from_raw_parts(partition_lba_ptr, num_parts * part_size);
        let sum = unsafe { GPTPartition::partition_table_checksum(&*checksum_array) };

        assert_eq!(
            sum,
            u32::from(gpt_header.part_table_crc32),
            "Bad CRC for partition array"
        );

        let mut partitions: Vec<GPTPartition> = Vec::new();
        GPTHeader::parse_partitions(
            part_array,
            u32::from(gpt_header.part_size) as usize,
            |part, parts| {
                parts.push(part.clone());
            },
            &mut partitions,
        );

        assert_eq!(partitions.len(), 2, "List of partitions should have been 2");
        let guid_time_low = u32::from(partitions[0].part_type.time_low);
        assert_eq!(
            guid_time_low, uuid_time_low_1,
            "First partition type did not match (time_low) expected value {:x}: {:x}",
            uuid_time_low_1, guid_time_low
        );
        let guid_time_low = u32::from(partitions[1].part_type.time_low);
        assert_eq!(
            guid_time_low, uuid_time_low_2,
            "Second partition type did not match (time_low) expected value {:x}: {:x}",
            uuid_time_low_2, guid_time_low
        );
    }

    #[test]
    fn parse_512_header() {
        test_disk_image::<LBA512>(GPT_EXAMPLE1, 0x0fc63daf, 0xc12a7328);
    }

    #[test]
    fn parse_4k_header() {
        test_disk_image::<LBA4K>(GPT_EXAMPLE_LBA4K, 0x0fc63daf, 0x0fc63daf);
    }
}