xref: /aosp_15_r20/bootable/libbootloader/gbl/libefi/src/protocol/gbl_efi_image_loading.rs (revision 5225e6b173e52d2efc6bcf950c27374fd72adabc)
1 // Copyright 2024, The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 //! Rust wrapper for `EFI_IMAGE_LOADING_PROTOCOL`.
16 
17 use crate::efi_call;
18 use crate::protocol::{Protocol, ProtocolInfo};
19 use arrayvec::ArrayVec;
20 use core::mem::{size_of, MaybeUninit};
21 use efi_types::{
22     EfiGuid, GblEfiImageBuffer, GblEfiImageInfo, GblEfiImageLoadingProtocol, GblEfiPartitionName,
23     PARTITION_NAME_LEN_U16,
24 };
25 use liberror::{Error, Result};
26 use spin::Mutex;
27 
28 /// GBL_IMAGE_LOADING_PROTOCOL
29 pub struct GblImageLoadingProtocol;
30 
31 impl ProtocolInfo for GblImageLoadingProtocol {
32     type InterfaceType = GblEfiImageLoadingProtocol;
33 
34     const GUID: EfiGuid =
35         EfiGuid::new(0xdb84b4fa, 0x53bd, 0x4436, [0x98, 0xa7, 0x4e, 0x02, 0x71, 0x42, 0x8b, 0xa8]);
36 }
37 
38 /// Max length of partition name in UTF8 in bytes.
39 pub const PARTITION_NAME_LEN_U8: usize = size_of::<char>() * PARTITION_NAME_LEN_U16;
40 
41 const MAX_ARRAY_SIZE: usize = 100;
42 static RETURNED_BUFFERS: Mutex<ArrayVec<usize, MAX_ARRAY_SIZE>> = Mutex::new(ArrayVec::new_const());
43 
44 /// Wrapper class for buffer received with [get_buffer] function.
45 ///
46 /// Helps to keep track of allocated memory and avoid getting same buffer more than once.
47 #[derive(Debug)]
48 pub struct EfiImageBuffer {
49     buffer: Option<&'static mut [MaybeUninit<u8>]>,
50 }
51 
52 impl EfiImageBuffer {
53     // # Safety
54     //
55     // `gbl_buffer` must represent valid buffer.
56     //
57     // # Return
58     //
59     // Err(EFI_STATUS_INVALID_PARAMETER) - If `gbl_buffer.Memory` == NULL
60     // Err(EFI_STATUS_ALREADY_STARTED) - Requested buffer was already returned and is still in use.
61     // Err(err) - on error
62     // Ok(_) - on success
new(gbl_buffer: GblEfiImageBuffer) -> Result<EfiImageBuffer>63     unsafe fn new(gbl_buffer: GblEfiImageBuffer) -> Result<EfiImageBuffer> {
64         if gbl_buffer.Memory.is_null() {
65             return Err(Error::InvalidInput);
66         }
67 
68         let addr = gbl_buffer.Memory as usize;
69         let mut returned_buffers = RETURNED_BUFFERS.lock();
70         if returned_buffers.contains(&addr) {
71             return Err(Error::AlreadyStarted);
72         }
73         returned_buffers.push(addr);
74 
75         // SAFETY:
76         // `gbl_buffer.Memory` is guarantied to be not null
77         // This code is relying on EFI protocol implementation to provide valid buffer pointer
78         // to memory region of size `gbl_buffer.SizeBytes`.
79         Ok(EfiImageBuffer {
80             buffer: Some(unsafe {
81                 core::slice::from_raw_parts_mut(
82                     gbl_buffer.Memory as *mut MaybeUninit<u8>,
83                     gbl_buffer.SizeBytes,
84                 )
85             }),
86         })
87     }
88 
89     /// Move buffer ownership out of EfiImageBuffer, and consume it.
take(mut self) -> &'static mut [MaybeUninit<u8>]90     pub fn take(mut self) -> &'static mut [MaybeUninit<u8>] {
91         self.buffer.take().unwrap()
92     }
93 
94     // Removes address from `RETURNED_BUFFERS`.
95     //
96     // # Safety
97     //
98     // Caller must guarantee that address is not referenced anymore.
release(address: usize)99     unsafe fn release(address: usize) {
100         let mut returned_buffers = RETURNED_BUFFERS.lock();
101         let res = returned_buffers.iter().position(|&val| val == address);
102         debug_assert!(
103             res.is_some(),
104             "EfiImageBuffer::release trying to release address ({address}) that is not tracked"
105         );
106         if let Some(pos) = res {
107             returned_buffers.swap_remove(pos);
108         }
109     }
110 }
111 
112 impl Drop for EfiImageBuffer {
drop(&mut self)113     fn drop(&mut self) {
114         if self.buffer.is_none() {
115             return;
116         }
117 
118         // SAFETY:
119         // EfiIMageBuffer is the only owner of the buffer. The only way to get address for it is to
120         // call `take()` which consumes `self.buffer`, which we check above.
121         unsafe { EfiImageBuffer::release((*self.buffer.as_ref().unwrap()).as_ptr() as usize) };
122     }
123 }
124 
125 // Protocol interface wrappers.
126 impl Protocol<'_, GblImageLoadingProtocol> {
127     /// Wrapper of `GBL_IMAGE_LOADING_PROTOCOL.get_buffer()`
128     ///
129     /// # Return
130     /// Ok(Some(EfiImageBuffer)) if buffer was successfully provided,
131     /// Ok(None) if buffer was not provided
132     /// Err(Error::EFI_STATUS_BUFFER_TOO_SMALL) if provided buffer is too small
133     /// Err(Error::EFI_STATUS_INVALID_PARAMETER) if received buffer is NULL
134     /// Err(Error::EFI_STATUS_ALREADY_STARTED) buffer was already returned and is still in use.
135     /// Err(err) if `err` occurred
get_buffer(&self, gbl_image_info: &GblEfiImageInfo) -> Result<EfiImageBuffer>136     pub fn get_buffer(&self, gbl_image_info: &GblEfiImageInfo) -> Result<EfiImageBuffer> {
137         let mut gbl_buffer: GblEfiImageBuffer = Default::default();
138         // SAFETY:
139         // `self.interface()?` guarantees self.interface is non-null and points to a valid object
140         // established by `Protocol::new()`.
141         // `self.interface` and `gbl_buffer` are input/output parameters, outlive the call and
142         // will not be retained.
143         // `gbl_buffer` returned by this call must not overlap, and will be checked by
144         // `EfiImageBuffer`
145         unsafe {
146             efi_call!(
147                 @bufsize gbl_image_info.SizeBytes,
148                 self.interface()?.get_buffer,
149                 self.interface,
150                 gbl_image_info,
151                 &mut gbl_buffer
152             )?;
153         }
154 
155         if gbl_buffer.SizeBytes < gbl_image_info.SizeBytes {
156             return Err(Error::BufferTooSmall(Some(gbl_image_info.SizeBytes)));
157         }
158 
159         // SAFETY:
160         // `gbl_buffer.Memory` must be not null. This checked in `new()` call
161         // `gbl_buffer.Size` must be valid size of the buffer.
162         // This protocol is relying on EFI protocol implementation to provide valid buffer pointer
163         // to memory region of size `gbl_buffer.SizeBytes`.
164         let image_buffer = unsafe { EfiImageBuffer::new(gbl_buffer)? };
165 
166         Ok(image_buffer)
167     }
168 
169     /// Wrapper of `GBL_IMAGE_LOADING_PROTOCOL.get_verify_partitions()`
170     ///
171     /// # Result
172     /// Err(BufferTooSmall(Some(size))) - when provided `partition_names` is less than expected
173     /// `size`
174     /// Err(err) - if error occurred.
175     /// Ok(len) - will return number of `GblEfiPartitionName`s copied to `partition_names` slice.
get_verify_partitions( &self, partition_names: &mut [GblEfiPartitionName], ) -> Result<usize>176     pub fn get_verify_partitions(
177         &self,
178         partition_names: &mut [GblEfiPartitionName],
179     ) -> Result<usize> {
180         let partition_count_in: usize = partition_names.len();
181         let mut partition_count: usize = partition_count_in;
182 
183         // SAFETY:
184         // `self.interface()?` guarantees self.interface is non-null and points to a valid object
185         // established by `Protocol::new()`.
186         // `self.interface` is input parameter, outlives the call, and will not be retained.
187         // `partition_count` must be set to valid length of `partition_names` array after the call.
188         // `partition_names` must be valid array of length `partition_count` after the call.
189         unsafe {
190             efi_call!(
191                 @bufsize partition_count,
192                 self.interface()?.get_verify_partitions,
193                 self.interface,
194                 &mut partition_count,
195                 partition_names.as_mut_ptr(),
196             )?;
197         }
198 
199         Ok(partition_count)
200     }
201 }
202 
203 #[cfg(test)]
204 mod test {
205     use super::*;
206     use crate::{
207         protocol::gbl_efi_image_loading::GblImageLoadingProtocol, test::run_test, DeviceHandle,
208         EfiEntry,
209     };
210     use core::{ffi::c_void, iter::zip, ptr::null_mut};
211     use efi_types::{
212         EfiStatus, EFI_STATUS_BAD_BUFFER_SIZE, EFI_STATUS_BUFFER_TOO_SMALL,
213         EFI_STATUS_INVALID_PARAMETER, EFI_STATUS_SUCCESS,
214     };
215     use spin::MutexGuard;
216     use std::cell::RefCell;
217     use std::collections::HashSet;
218 
219     const UCS2_STR: [u16; 8] = [0x2603, 0x0073, 0x006e, 0x006f, 0x0077, 0x006d, 0x0061, 0x006e];
220     const UTF8_STR: &str = "☃snowman";
221     const PARTITIONS_MAX: usize = 128;
222 
get_buffer_utf8() -> [[u8; PARTITION_NAME_LEN_U8]; PARTITIONS_MAX]223     fn get_buffer_utf8() -> [[u8; PARTITION_NAME_LEN_U8]; PARTITIONS_MAX] {
224         [[0; PARTITION_NAME_LEN_U8]; PARTITIONS_MAX]
225     }
226 
get_printable_utf16() -> Vec<u16>227     fn get_printable_utf16() -> Vec<u16> {
228         (0x0021..0x007e).collect::<Vec<u16>>()
229     }
230 
get_printable_string() -> String231     fn get_printable_string() -> String {
232         String::from_utf8((0x21..0x7e).collect::<Vec<u8>>()).unwrap()
233     }
234 
235     #[test]
test_partition_name_get_str()236     fn test_partition_name_get_str() {
237         let mut buffer = [0u8; 100];
238         // empty string
239         assert_eq!(GblEfiPartitionName::from([0u16]).get_str(&mut buffer).unwrap(), "");
240         assert_eq!(GblEfiPartitionName::from([0u16]).get_str(&mut buffer).unwrap(), "");
241         assert_eq!(GblEfiPartitionName::from([0x0000]).get_str(&mut buffer[..0]).unwrap(), "");
242 
243         // Special characters
244         assert_eq!(GblEfiPartitionName::from(UCS2_STR).get_str(&mut buffer).unwrap(), UTF8_STR);
245 
246         // Null character in the middle
247         assert_eq!(
248             GblEfiPartitionName::from([0x006d, 0x0075, 0x0000, 0x0073, 0x0069, 0x0063])
249                 .get_str(&mut buffer),
250             Ok("mu")
251         );
252 
253         // Null character at the end
254         assert_eq!(
255             GblEfiPartitionName::from([0x006d, 0x0075, 0x0073, 0x0069, 0x0063, 0x0000])
256                 .get_str(&mut buffer),
257             Ok("music")
258         );
259 
260         // exact buffer size
261         assert_eq!(
262             GblEfiPartitionName::from([0x006d, 0x0075, 0x0073, 0x0069, 0x0063])
263                 .get_str(&mut buffer[..5]),
264             Ok("music")
265         );
266         assert_eq!(
267             GblEfiPartitionName::from([0x006d, 0x0075, 0x0000, 0x0073, 0x0069, 0x0063])
268                 .get_str(&mut buffer[..2]),
269             Ok("mu")
270         );
271     }
272 
273     #[test]
test_partition_name_get_str_small_buffer()274     fn test_partition_name_get_str_small_buffer() {
275         let mut buffer = [0u8; 8];
276         let partition_name: GblEfiPartitionName = UCS2_STR.into();
277         assert_eq!(partition_name.get_str(&mut buffer), Err(10usize));
278     }
279 
generate_protocol<'a, P: ProtocolInfo>( efi_entry: &'a EfiEntry, proto: &'a mut P::InterfaceType, ) -> Protocol<'a, P>280     fn generate_protocol<'a, P: ProtocolInfo>(
281         efi_entry: &'a EfiEntry,
282         proto: &'a mut P::InterfaceType,
283     ) -> Protocol<'a, P> {
284         // SAFETY:
285         // proto is a valid pointer and lasts at least as long as efi_entry.
286         unsafe { Protocol::<'a, P>::new(DeviceHandle::new(null_mut()), proto, efi_entry) }
287     }
288 
289     #[test]
test_proto_get_partitions_count()290     fn test_proto_get_partitions_count() {
291         const EXPECTED_PARTITIONS_NUM: usize = 2;
292         unsafe extern "C" fn get_verify_partitions(
293             _: *mut GblEfiImageLoadingProtocol,
294             number_of_partitions: *mut usize,
295             _: *mut GblEfiPartitionName,
296         ) -> EfiStatus {
297             assert!(!number_of_partitions.is_null());
298             // SAFETY
299             // `number_of_partitions` must be valid pointer to usize
300             let number_of_partitions = unsafe { number_of_partitions.as_mut() }.unwrap();
301 
302             *number_of_partitions = EXPECTED_PARTITIONS_NUM;
303             EFI_STATUS_BUFFER_TOO_SMALL
304         }
305 
306         run_test(|image_handle, systab_ptr| {
307             let mut image_loading = GblEfiImageLoadingProtocol {
308                 get_verify_partitions: Some(get_verify_partitions),
309                 ..Default::default()
310             };
311             let efi_entry = EfiEntry { image_handle, systab_ptr };
312             let protocol =
313                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
314 
315             let mut partitions: [GblEfiPartitionName; 0] = Default::default();
316             assert_eq!(
317                 protocol.get_verify_partitions(&mut partitions).unwrap_err(),
318                 Error::BufferTooSmall(Some(EXPECTED_PARTITIONS_NUM))
319             );
320         });
321     }
322 
323     #[test]
test_proto_get_partitions_count_error()324     fn test_proto_get_partitions_count_error() {
325         unsafe extern "C" fn get_verify_partitions(
326             _: *mut GblEfiImageLoadingProtocol,
327             _: *mut usize,
328             _: *mut GblEfiPartitionName,
329         ) -> EfiStatus {
330             EFI_STATUS_INVALID_PARAMETER
331         }
332 
333         run_test(|image_handle, systab_ptr| {
334             let mut image_loading = GblEfiImageLoadingProtocol {
335                 get_verify_partitions: Some(get_verify_partitions),
336                 ..Default::default()
337             };
338             let efi_entry = EfiEntry { image_handle, systab_ptr };
339             let protocol =
340                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
341 
342             let mut partitions: [GblEfiPartitionName; 0] = Default::default();
343             assert_eq!(
344                 protocol.get_verify_partitions(&mut partitions).unwrap_err(),
345                 Error::InvalidInput
346             );
347         });
348     }
349 
350     #[test]
test_proto_get_partitions_len_and_value()351     fn test_proto_get_partitions_len_and_value() {
352         const EXPECTED_PARTITIONS_NUM: usize = 1;
353         unsafe extern "C" fn get_verify_partitions(
354             _: *mut GblEfiImageLoadingProtocol,
355             number_of_partitions: *mut usize,
356             partitions: *mut GblEfiPartitionName,
357         ) -> EfiStatus {
358             // SAFETY
359             // `number_of_partitions` must be valid pointer to usize
360             let number_of_partitions = unsafe { number_of_partitions.as_mut() }.unwrap();
361 
362             match *number_of_partitions {
363                 n if n < EXPECTED_PARTITIONS_NUM => {
364                     *number_of_partitions = EXPECTED_PARTITIONS_NUM;
365                     EFI_STATUS_BUFFER_TOO_SMALL
366                 }
367                 _ => {
368                     // SAFETY
369                     // `partitions` must be valid array of size `number_of_partitions`
370                     let partitions = unsafe {
371                         core::slice::from_raw_parts_mut(partitions, *number_of_partitions)
372                     };
373                     *number_of_partitions = 1;
374                     partitions[0].StrUtf16[..UCS2_STR.len()].copy_from_slice(&UCS2_STR);
375                     EFI_STATUS_SUCCESS
376                 }
377             }
378         }
379 
380         run_test(|image_handle, systab_ptr| {
381             let mut buffer_utf8 = get_buffer_utf8();
382             let mut image_loading = GblEfiImageLoadingProtocol {
383                 get_verify_partitions: Some(get_verify_partitions),
384                 ..Default::default()
385             };
386             let efi_entry = EfiEntry { image_handle, systab_ptr };
387             let protocol =
388                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
389             let mut partitions: [GblEfiPartitionName; 2] = Default::default();
390 
391             assert_eq!(
392                 protocol.get_verify_partitions(&mut partitions[..0]).unwrap_err(),
393                 Error::BufferTooSmall(Some(EXPECTED_PARTITIONS_NUM))
394             );
395             let verify_partitions_len = protocol.get_verify_partitions(&mut partitions).unwrap();
396             assert_eq!(verify_partitions_len, 1);
397             assert_eq!(partitions[0].get_str(&mut buffer_utf8[0]), Ok(UTF8_STR));
398         });
399     }
400 
401     #[test]
test_proto_get_partitions_zero_len()402     fn test_proto_get_partitions_zero_len() {
403         const EXPECTED_PARTITIONS_NUM: usize = 1;
404         unsafe extern "C" fn get_verify_partitions(
405             _: *mut GblEfiImageLoadingProtocol,
406             number_of_partitions: *mut usize,
407             _: *mut GblEfiPartitionName,
408         ) -> EfiStatus {
409             // SAFETY
410             // `number_of_partitions` must be valid pointer to usize
411             let number_of_partitions = unsafe { number_of_partitions.as_mut() }.unwrap();
412             assert_eq!(*number_of_partitions, 0);
413             *number_of_partitions = EXPECTED_PARTITIONS_NUM;
414             EFI_STATUS_BUFFER_TOO_SMALL
415         }
416 
417         run_test(|image_handle, systab_ptr| {
418             let mut image_loading = GblEfiImageLoadingProtocol {
419                 get_verify_partitions: Some(get_verify_partitions),
420                 ..Default::default()
421             };
422             let efi_entry = EfiEntry { image_handle, systab_ptr };
423             let protocol =
424                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
425             let mut partitions: [GblEfiPartitionName; 0] = Default::default();
426 
427             let verify_partitions_res = protocol.get_verify_partitions(&mut partitions);
428             assert_eq!(
429                 verify_partitions_res.unwrap_err(),
430                 Error::BufferTooSmall(Some(EXPECTED_PARTITIONS_NUM))
431             );
432         });
433     }
434 
435     #[test]
test_proto_get_partitions_less_than_buffer()436     fn test_proto_get_partitions_less_than_buffer() {
437         const EXPECTED_PARTITIONS_NUM: usize = 1;
438         unsafe extern "C" fn get_verify_partitions(
439             _: *mut GblEfiImageLoadingProtocol,
440             number_of_partitions: *mut usize,
441             partitions: *mut GblEfiPartitionName,
442         ) -> EfiStatus {
443             // SAFETY
444             // `number_of_partitions` must be valid pointer to usize
445             let number_of_partitions = unsafe { number_of_partitions.as_mut() }.unwrap();
446 
447             assert!(!partitions.is_null());
448             assert!(*number_of_partitions > 0);
449 
450             // SAFETY
451             // `partitions` must be valid array of size `number_of_partitions`
452             let partitions =
453                 unsafe { core::slice::from_raw_parts_mut(partitions, *number_of_partitions) };
454             partitions[0].StrUtf16[..UCS2_STR.len()].copy_from_slice(&UCS2_STR);
455             *number_of_partitions = EXPECTED_PARTITIONS_NUM;
456             EFI_STATUS_SUCCESS
457         }
458 
459         run_test(|image_handle, systab_ptr| {
460             let mut buffer_utf8 = get_buffer_utf8();
461             let mut image_loading = GblEfiImageLoadingProtocol {
462                 get_verify_partitions: Some(get_verify_partitions),
463                 ..Default::default()
464             };
465             let efi_entry = EfiEntry { image_handle, systab_ptr };
466             let protocol =
467                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
468             let mut partitions: [GblEfiPartitionName; 2] = Default::default();
469 
470             let verify_partitions_len = protocol.get_verify_partitions(&mut partitions).unwrap();
471             assert_eq!(verify_partitions_len, EXPECTED_PARTITIONS_NUM);
472             assert_eq!(partitions[0].get_str(&mut buffer_utf8[0]), Ok(UTF8_STR));
473         });
474     }
475 
476     #[test]
test_proto_get_partitions_name_max()477     fn test_proto_get_partitions_name_max() {
478         unsafe extern "C" fn get_verify_partitions(
479             _: *mut GblEfiImageLoadingProtocol,
480             number_of_partitions: *mut usize,
481             partitions: *mut GblEfiPartitionName,
482         ) -> EfiStatus {
483             let printable_utf16 = get_printable_utf16();
484             // SAFETY
485             // `number_of_partitions` must be valid pointer to usize
486             let number_of_partitions = unsafe { number_of_partitions.as_mut() }.unwrap();
487 
488             assert!(!partitions.is_null());
489             assert!(*number_of_partitions > 0);
490 
491             // SAFETY
492             // `partitions` must be valid array of size `number_of_partitions`
493             let partitions =
494                 unsafe { core::slice::from_raw_parts_mut(partitions, *number_of_partitions) };
495 
496             let partition_names: [GblEfiPartitionName; PARTITIONS_MAX] = (0
497                 ..PARTITION_NAME_LEN_U16)
498                 .cycle()
499                 .take(PARTITIONS_MAX)
500                 .map(|i| printable_utf16[i..i + PARTITION_NAME_LEN_U16].into())
501                 .collect::<Vec<_>>()
502                 .try_into()
503                 .unwrap();
504 
505             *number_of_partitions = partition_names.len();
506 
507             for (p_out, p_gen) in zip(partitions.iter_mut(), partition_names.iter()) {
508                 *p_out = *p_gen;
509             }
510 
511             EFI_STATUS_SUCCESS
512         }
513 
514         run_test(|image_handle, systab_ptr| {
515             let mut buffer_utf8 = get_buffer_utf8();
516             let printable_str = get_printable_string();
517             let mut image_loading = GblEfiImageLoadingProtocol {
518                 get_verify_partitions: Some(get_verify_partitions),
519                 ..Default::default()
520             };
521             let efi_entry = EfiEntry { image_handle, systab_ptr };
522             let protocol =
523                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
524             let mut partitions = [GblEfiPartitionName::default(); PARTITIONS_MAX];
525             let expected_strs: Vec<&str> = (0..PARTITION_NAME_LEN_U16)
526                 .cycle()
527                 .take(PARTITIONS_MAX)
528                 .map(|i| &printable_str[i..i + PARTITION_NAME_LEN_U16])
529                 .collect();
530 
531             let verify_partitions_len = protocol.get_verify_partitions(&mut partitions).unwrap();
532             assert_eq!(verify_partitions_len, PARTITIONS_MAX);
533 
534             assert!(zip(partitions.iter(), expected_strs.iter())
535                 .all(|(p, expected_str)| { p.get_str(&mut buffer_utf8[0]) == Ok(*expected_str) }));
536         });
537     }
538 
539     #[test]
test_proto_get_partitions()540     fn test_proto_get_partitions() {
541         const EXPECTED_PARTITIONS_NUM: usize = 2;
542         unsafe extern "C" fn get_verify_partitions(
543             _: *mut GblEfiImageLoadingProtocol,
544             number_of_partitions: *mut usize,
545             partitions: *mut GblEfiPartitionName,
546         ) -> EfiStatus {
547             // SAFETY
548             // `number_of_partitions` must be valid pointer to usize
549             let number_of_partitions = unsafe { number_of_partitions.as_mut() }.unwrap();
550 
551             assert!(!partitions.is_null());
552             assert!(*number_of_partitions > 0);
553 
554             // SAFETY
555             // `partitions` must be valid array of size `number_of_partitions`
556             let partitions =
557                 unsafe { core::slice::from_raw_parts_mut(partitions, *number_of_partitions) };
558             partitions[0].StrUtf16[..UCS2_STR.len()].copy_from_slice(&UCS2_STR);
559             partitions[1].StrUtf16[..UCS2_STR.len() - 1].copy_from_slice(&UCS2_STR[1..]);
560             *number_of_partitions = EXPECTED_PARTITIONS_NUM;
561             EFI_STATUS_SUCCESS
562         }
563 
564         run_test(|image_handle, systab_ptr| {
565             let mut buffer_utf8 = get_buffer_utf8();
566             let mut image_loading = GblEfiImageLoadingProtocol {
567                 get_verify_partitions: Some(get_verify_partitions),
568                 ..Default::default()
569             };
570             let efi_entry = EfiEntry { image_handle, systab_ptr };
571             let protocol =
572                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
573             let mut partitions: [GblEfiPartitionName; EXPECTED_PARTITIONS_NUM] = Default::default();
574 
575             let verify_partitions_len = protocol.get_verify_partitions(&mut partitions).unwrap();
576             assert_eq!(verify_partitions_len, EXPECTED_PARTITIONS_NUM);
577 
578             let mut char_idx = UTF8_STR.char_indices();
579             char_idx.next();
580             let (next_char_pos, _) = char_idx.next().unwrap();
581 
582             assert_eq!(partitions[0].get_str(&mut buffer_utf8[0]), Ok(UTF8_STR));
583             assert_eq!(partitions[1].get_str(&mut buffer_utf8[1]), Ok(&UTF8_STR[next_char_pos..]));
584         });
585     }
586 
587     #[test]
test_proto_get_partitions_empty()588     fn test_proto_get_partitions_empty() {
589         const EXPECTED_PARTITIONS_NUM: usize = 0;
590         unsafe extern "C" fn get_verify_partitions(
591             _: *mut GblEfiImageLoadingProtocol,
592             number_of_partitions: *mut usize,
593             partitions: *mut GblEfiPartitionName,
594         ) -> EfiStatus {
595             // SAFETY
596             // `number_of_partitions` must be valid pointer to usize
597             let number_of_partitions = unsafe { number_of_partitions.as_mut() }.unwrap();
598             assert!(!partitions.is_null());
599             *number_of_partitions = EXPECTED_PARTITIONS_NUM;
600             EFI_STATUS_SUCCESS
601         }
602 
603         run_test(|image_handle, systab_ptr| {
604             let mut image_loading = GblEfiImageLoadingProtocol {
605                 get_verify_partitions: Some(get_verify_partitions),
606                 ..Default::default()
607             };
608             let efi_entry = EfiEntry { image_handle, systab_ptr };
609             let protocol =
610                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
611             let mut partitions: [GblEfiPartitionName; 2] = Default::default();
612 
613             let verify_partitions_len = protocol.get_verify_partitions(&mut partitions).unwrap();
614             assert_eq!(verify_partitions_len, EXPECTED_PARTITIONS_NUM);
615         });
616     }
617 
618     #[test]
test_proto_get_partitions_error()619     fn test_proto_get_partitions_error() {
620         unsafe extern "C" fn get_verify_partitions(
621             _: *mut GblEfiImageLoadingProtocol,
622             _: *mut usize,
623             _: *mut GblEfiPartitionName,
624         ) -> EfiStatus {
625             EFI_STATUS_BAD_BUFFER_SIZE
626         }
627 
628         run_test(|image_handle, systab_ptr| {
629             let mut image_loading = GblEfiImageLoadingProtocol {
630                 get_verify_partitions: Some(get_verify_partitions),
631                 ..Default::default()
632             };
633             let efi_entry = EfiEntry { image_handle, systab_ptr };
634             let protocol =
635                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
636             let mut partitions: [GblEfiPartitionName; 1] = Default::default();
637 
638             assert!(protocol.get_verify_partitions(&mut partitions).is_err());
639         });
640     }
641 
642     // Mutex to make sure tests that use `static RETURNED_BUFFERS` do not run in parallel to avoid
643     // unexpected results since this is global static that would be shared between tests. And can
644     // overflow due to amount of tests.
645     //
646     // See MEMORY_TEST thread local variable that should be used for convenience.
647     static GET_BUFFER_MUTEX: Mutex<()> = Mutex::new(());
648 
649     // Size of MEMORY_TEST buffers
650     const MEMORY_TEST_BUF_SIZE: usize = 100;
651 
652     // Helper struct for safe acquisition of the memory and releasing it on exit
653     struct MemoryTest<'a> {
654         // Tracking if test guard was acquired with `start()`
655         init: bool,
656         // Keep track of all buffers returned
657         returned_buffers: HashSet<*mut [u8; MEMORY_TEST_BUF_SIZE]>,
658         // Store same buffer value for `get_memory_same()` calls.
659         same_buffer: Option<*mut c_void>,
660         // It is necessary to run 1 test at a time that uses UEFI `get_buffer()`.
661         // Because it is uses static size array to track returned values to prevent reusing same
662         // buffer. With current number of test if they run simultaneously there are situations when
663         // array limit is reached and unlucky test will fail. To prevent this flakiness this guard
664         // is used.
665         _get_buffer_guard: MutexGuard<'a, ()>,
666     }
667 
668     thread_local! {
669         static MEMORY_TEST: RefCell<MemoryTest<'static>> = RefCell::new(MemoryTest::new());
670     }
671     struct MemoryTestInitGuard {}
672 
673     impl Drop for MemoryTestInitGuard {
drop(&mut self)674         fn drop(&mut self) {
675             MEMORY_TEST.with_borrow_mut(|v| v.stop());
676         }
677     }
678 
679     // Helper implementation for getting raw buffers for `get_buffer()` calls.
680     // And cleanly releasing buffers at the end of the test to prevent memory leaks.
681     //
682     // Use `thread_local` static MEMORY_TEST variable.
683     //
684     // ```
685     // let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
686     // ...
687     // buffer.Memory = MEMORY_TEST.with_borrow_mut(|v| v.get_memory());
688     // ...
689     // ```
690     // _memory_guard will make sure to cleanup all memory that was retrieved by `get_memory()` call
691     //
692     // Note:
693     // If using raw EfiImageBuffer, there is no need for this helper. Since the structure does
694     // cleaning on its own.
695     // Except when using `EfiImageBuffer::take()` then manual `EfiImageBuffer::release()` must be
696     // used.
697     impl MemoryTest<'_> {
new() -> Self698         fn new() -> Self {
699             MemoryTest {
700                 init: false,
701                 returned_buffers: HashSet::new(),
702                 same_buffer: None,
703                 _get_buffer_guard: GET_BUFFER_MUTEX.lock(),
704             }
705         }
706 
start(&mut self) -> MemoryTestInitGuard707         fn start(&mut self) -> MemoryTestInitGuard {
708             assert!(!self.init);
709             self.init = true;
710             MemoryTestInitGuard {}
711         }
712 
713         // Return heap allocated buffer, and keep track of its address
714         // To verify it was properly released
715         //
716         // # Safety
717         //
718         // Returned pointers must not be used after guard returned by `start()`
719         // is destroyed.
get_memory(&mut self) -> *mut c_void720         unsafe fn get_memory(&mut self) -> *mut c_void {
721             assert!(self.init);
722             let ptr = Box::into_raw(Box::new([0u8; MEMORY_TEST_BUF_SIZE]));
723             assert!(self.returned_buffers.insert(ptr));
724             ptr as *mut c_void
725         }
726 
727         // Return same buffer for all calls, allocating and tracking it only for first call.
728         //
729         // # Safety
730         //
731         // Returned pointers must not be used after guard returned by `start()`
732         // is destroyed.
get_memory_same(&mut self) -> *mut c_void733         unsafe fn get_memory_same(&mut self) -> *mut c_void {
734             if self.same_buffer.is_none() {
735                 // SAFETY:
736                 // This function has same requirements as `get_memory()`
737                 let address = unsafe { self.get_memory() };
738 
739                 self.same_buffer = Some(address);
740             }
741 
742             *self.same_buffer.as_mut().unwrap()
743         }
744 
745         // Clear address from buffers returned list
746         // Which allows to reuse it in other tests.
stop(&mut self)747         fn stop(&mut self) {
748             assert!(self.init);
749             self.init = false;
750             self.same_buffer = None;
751             for ptr in self.returned_buffers.drain() {
752                 // SAFETY:
753                 // `ptr` is valid since was created by `Box::into_raw()`.
754                 // Double free is covered by safety requirements for this function. (`release_memory()`
755                 // must be called only on buffer holding the only reference to buffer.)
756                 // As well as tracking `returned_buffers` and asserting remove in the line above.
757                 unsafe {
758                     let _restore_box = Box::from_raw(ptr);
759                 }
760             }
761         }
762     }
763 
764     #[test]
test_proto_get_buffer_error()765     fn test_proto_get_buffer_error() {
766         unsafe extern "C" fn get_buffer(
767             _: *mut GblEfiImageLoadingProtocol,
768             _: *const GblEfiImageInfo,
769             _: *mut GblEfiImageBuffer,
770         ) -> EfiStatus {
771             EFI_STATUS_INVALID_PARAMETER
772         }
773 
774         run_test(|image_handle, systab_ptr| {
775             let gbl_image_info: GblEfiImageInfo = Default::default();
776             let mut image_loading =
777                 GblEfiImageLoadingProtocol { get_buffer: Some(get_buffer), ..Default::default() };
778             let efi_entry = EfiEntry { image_handle, systab_ptr };
779             let protocol =
780                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
781 
782             let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
783             assert!(protocol.get_buffer(&gbl_image_info).is_err());
784         });
785     }
786 
787     #[test]
test_proto_get_buffer_not_provided()788     fn test_proto_get_buffer_not_provided() {
789         unsafe extern "C" fn get_buffer(
790             _: *mut GblEfiImageLoadingProtocol,
791             image_info: *const GblEfiImageInfo,
792             buffer: *mut GblEfiImageBuffer,
793         ) -> EfiStatus {
794             assert!(!image_info.is_null());
795             assert!(!buffer.is_null());
796             // SAFETY
797             // `buffer` must be valid pointer to `GblEfiImageBuffer`
798             let buffer = unsafe { buffer.as_mut() }.unwrap();
799 
800             buffer.Memory = null_mut();
801             buffer.SizeBytes = 10;
802 
803             EFI_STATUS_SUCCESS
804         }
805 
806         run_test(|image_handle, systab_ptr| {
807             let gbl_image_info: GblEfiImageInfo = Default::default();
808             let mut image_loading =
809                 GblEfiImageLoadingProtocol { get_buffer: Some(get_buffer), ..Default::default() };
810             let efi_entry = EfiEntry { image_handle, systab_ptr };
811             let protocol =
812                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
813 
814             let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
815             let res = protocol.get_buffer(&gbl_image_info);
816             assert_eq!(res.unwrap_err(), Error::InvalidInput);
817         });
818     }
819 
820     #[test]
test_proto_get_buffer_zero_size()821     fn test_proto_get_buffer_zero_size() {
822         unsafe extern "C" fn get_buffer(
823             _: *mut GblEfiImageLoadingProtocol,
824             image_info: *const GblEfiImageInfo,
825             buffer: *mut GblEfiImageBuffer,
826         ) -> EfiStatus {
827             assert!(!image_info.is_null());
828             assert!(!buffer.is_null());
829             // SAFETY
830             // `buffer` must be valid pointer to `GblEfiImageBuffer`
831             let buffer = unsafe { buffer.as_mut() }.unwrap();
832 
833             // SAFETY:
834             // `get_memory()` results are returned in `buffer` in `get_buffer()` function.
835             // All usage of `get_buffer()` results are not leaving `run_test()` scope.
836             // Same function where `start()` guard is acquired, so it will not outlive guard.
837             unsafe {
838                 buffer.Memory = MEMORY_TEST.with_borrow_mut(|v| v.get_memory());
839             }
840             buffer.SizeBytes = 0;
841 
842             EFI_STATUS_SUCCESS
843         }
844 
845         run_test(|image_handle, systab_ptr| {
846             let gbl_image_info: GblEfiImageInfo = Default::default();
847             let mut image_loading =
848                 GblEfiImageLoadingProtocol { get_buffer: Some(get_buffer), ..Default::default() };
849             let efi_entry = EfiEntry { image_handle, systab_ptr };
850             let protocol =
851                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
852 
853             let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
854             let res = protocol.get_buffer(&gbl_image_info).unwrap();
855             assert!(res.buffer.as_ref().unwrap().is_empty());
856         });
857     }
858 
859     #[test]
test_proto_get_buffer_small()860     fn test_proto_get_buffer_small() {
861         unsafe extern "C" fn get_buffer(
862             _: *mut GblEfiImageLoadingProtocol,
863             image_info: *const GblEfiImageInfo,
864             buffer: *mut GblEfiImageBuffer,
865         ) -> EfiStatus {
866             assert!(!image_info.is_null());
867             // SAFETY
868             // `image_info` must be valid pointer to `GblEfiImageInfo`
869             let image_info = unsafe { image_info.as_ref() }.unwrap();
870             assert!(!buffer.is_null());
871             // SAFETY
872             // `buffer` must be valid pointer to `GblEfiImageBuffer`
873             let buffer = unsafe { buffer.as_mut() }.unwrap();
874 
875             // SAFETY:
876             // `get_memory()` results are returned in `buffer` in `get_buffer()` function.
877             // All usage of `get_buffer()` results are not leaving `run_test()` scope.
878             // Same function where `start()` guard is acquired, so it will not outlive guard.
879             unsafe {
880                 buffer.Memory = MEMORY_TEST.with_borrow_mut(|v| v.get_memory());
881             }
882             buffer.SizeBytes = image_info.SizeBytes - 1;
883 
884             EFI_STATUS_SUCCESS
885         }
886 
887         run_test(|image_handle, systab_ptr| {
888             let gbl_image_info: GblEfiImageInfo =
889                 GblEfiImageInfo { ImageType: [0; PARTITION_NAME_LEN_U16], SizeBytes: 10 };
890             let mut image_loading =
891                 GblEfiImageLoadingProtocol { get_buffer: Some(get_buffer), ..Default::default() };
892             let efi_entry = EfiEntry { image_handle, systab_ptr };
893             let protocol =
894                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
895 
896             let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
897             let res = protocol.get_buffer(&gbl_image_info);
898             assert_eq!(res.unwrap_err(), Error::BufferTooSmall(Some(10)));
899         });
900     }
901 
902     #[test]
test_proto_get_buffer()903     fn test_proto_get_buffer() {
904         unsafe extern "C" fn get_buffer(
905             _: *mut GblEfiImageLoadingProtocol,
906             image_info: *const GblEfiImageInfo,
907             buffer: *mut GblEfiImageBuffer,
908         ) -> EfiStatus {
909             assert!(!image_info.is_null());
910             assert!(!buffer.is_null());
911             // SAFETY
912             // `buffer` must be valid pointer to `GblEfiImageBuffer`
913             let buffer = unsafe { buffer.as_mut() }.unwrap();
914 
915             // SAFETY:
916             // `get_memory()` results are returned in `buffer` in `get_buffer()` function.
917             // All usage of `get_buffer()` results are not leaving `run_test()` scope.
918             // Same function where `start()` guard is acquired, so it will not outlive guard.
919             unsafe {
920                 buffer.Memory = MEMORY_TEST.with_borrow_mut(|v| v.get_memory());
921             }
922             buffer.SizeBytes = MEMORY_TEST_BUF_SIZE;
923 
924             EFI_STATUS_SUCCESS
925         }
926 
927         run_test(|image_handle, systab_ptr| {
928             let gbl_image_info: GblEfiImageInfo =
929                 GblEfiImageInfo { ImageType: [0; PARTITION_NAME_LEN_U16], SizeBytes: 100 };
930             let mut image_loading =
931                 GblEfiImageLoadingProtocol { get_buffer: Some(get_buffer), ..Default::default() };
932             let efi_entry = EfiEntry { image_handle, systab_ptr };
933             let protocol =
934                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
935 
936             let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
937             let buf = protocol.get_buffer(&gbl_image_info).unwrap();
938             assert_ne!(buf.buffer.as_ref().unwrap().as_ptr(), null_mut());
939             assert_eq!(buf.buffer.as_ref().unwrap().len(), 100);
940         });
941     }
942 
943     #[test]
test_proto_get_buffer_image_type()944     fn test_proto_get_buffer_image_type() {
945         const IMAGE_TYPE_STR: &'static str = "test";
946         unsafe extern "C" fn get_buffer(
947             _: *mut GblEfiImageLoadingProtocol,
948             image_info: *const GblEfiImageInfo,
949             buffer: *mut GblEfiImageBuffer,
950         ) -> EfiStatus {
951             assert!(!image_info.is_null());
952             // SAFETY
953             // `image_info` must be valid pointer to `GblEfiImageInfo`
954             let image_info = unsafe { image_info.as_ref() }.unwrap();
955             assert!(!buffer.is_null());
956             // SAFETY
957             // `buffer` must be valid pointer to `GblEfiImageBuffer`
958             let buffer = unsafe { buffer.as_mut() }.unwrap();
959 
960             let mut buffer_utf8 = [0u8; 100];
961             assert_eq!(
962                 GblEfiPartitionName::from(image_info.ImageType).get_str(&mut buffer_utf8).unwrap(),
963                 IMAGE_TYPE_STR
964             );
965 
966             // SAFETY:
967             // `get_memory()` results are returned in `buffer` in `get_buffer()` function.
968             // All usage of `get_buffer()` results are not leaving `run_test()` scope.
969             // Same function where `start()` guard is acquired, so it will not outlive guard.
970             unsafe {
971                 buffer.Memory = MEMORY_TEST.with_borrow_mut(|v| v.get_memory());
972             }
973             buffer.SizeBytes = MEMORY_TEST_BUF_SIZE;
974 
975             EFI_STATUS_SUCCESS
976         }
977 
978         run_test(|image_handle, systab_ptr| {
979             let mut image_type = [0u16; PARTITION_NAME_LEN_U16];
980             image_type[..4].copy_from_slice(&IMAGE_TYPE_STR.encode_utf16().collect::<Vec<u16>>());
981             let gbl_image_info: GblEfiImageInfo =
982                 GblEfiImageInfo { ImageType: image_type, SizeBytes: 100 };
983             let mut image_loading =
984                 GblEfiImageLoadingProtocol { get_buffer: Some(get_buffer), ..Default::default() };
985             let efi_entry = EfiEntry { image_handle, systab_ptr };
986             let protocol =
987                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
988 
989             let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
990             assert!(protocol.get_buffer(&gbl_image_info).is_ok());
991         });
992     }
993 
994     #[test]
test_proto_get_buffer_double_call()995     fn test_proto_get_buffer_double_call() {
996         unsafe extern "C" fn get_buffer(
997             _: *mut GblEfiImageLoadingProtocol,
998             image_info: *const GblEfiImageInfo,
999             buffer: *mut GblEfiImageBuffer,
1000         ) -> EfiStatus {
1001             assert!(!image_info.is_null());
1002             assert!(!buffer.is_null());
1003             // SAFETY
1004             // `buffer` must be valid pointer to `GblEfiImageBuffer`
1005             let buffer = unsafe { buffer.as_mut() }.unwrap();
1006 
1007             // SAFETY:
1008             // `get_memory_same()` results are returned in `buffer` in `get_buffer()` function.
1009             // All usage of `get_buffer()` results are not leaving `run_test()` scope.
1010             // Same function where `start()` guard is acquired, so it will not outlive guard.
1011             unsafe {
1012                 buffer.Memory = MEMORY_TEST.with_borrow_mut(|v| v.get_memory_same());
1013             }
1014             buffer.SizeBytes = MEMORY_TEST_BUF_SIZE;
1015 
1016             EFI_STATUS_SUCCESS
1017         }
1018 
1019         run_test(|image_handle, systab_ptr| {
1020             let gbl_image_info: GblEfiImageInfo =
1021                 GblEfiImageInfo { ImageType: [0; PARTITION_NAME_LEN_U16], SizeBytes: 100 };
1022             let mut image_loading =
1023                 GblEfiImageLoadingProtocol { get_buffer: Some(get_buffer), ..Default::default() };
1024             let efi_entry = EfiEntry { image_handle, systab_ptr };
1025             let protocol =
1026                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
1027 
1028             let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
1029             let _buf = protocol.get_buffer(&gbl_image_info).unwrap();
1030             assert_eq!(protocol.get_buffer(&gbl_image_info).unwrap_err(), Error::AlreadyStarted);
1031         });
1032     }
1033 
1034     #[test]
test_proto_get_buffer_double_call_after_drop()1035     fn test_proto_get_buffer_double_call_after_drop() {
1036         unsafe extern "C" fn get_buffer(
1037             _: *mut GblEfiImageLoadingProtocol,
1038             image_info: *const GblEfiImageInfo,
1039             buffer: *mut GblEfiImageBuffer,
1040         ) -> EfiStatus {
1041             assert!(!image_info.is_null());
1042             assert!(!buffer.is_null());
1043             // SAFETY
1044             // `buffer` must be valid pointer to `GblEfiImageBuffer`
1045             let buffer = unsafe { buffer.as_mut() }.unwrap();
1046 
1047             // SAFETY:
1048             // `get_memory()` results are returned in `buffer` in `get_buffer()` function.
1049             // All usage of `get_buffer()` results are not leaving `run_test()` scope.
1050             // Same function where `start()` guard is acquired, so it will not outlive guard.
1051             unsafe {
1052                 buffer.Memory = MEMORY_TEST.with_borrow_mut(|v| v.get_memory_same());
1053             }
1054             buffer.SizeBytes = MEMORY_TEST_BUF_SIZE;
1055 
1056             EFI_STATUS_SUCCESS
1057         }
1058 
1059         run_test(|image_handle, systab_ptr| {
1060             let gbl_image_info: GblEfiImageInfo =
1061                 GblEfiImageInfo { ImageType: [0; PARTITION_NAME_LEN_U16], SizeBytes: 100 };
1062             let mut image_loading =
1063                 GblEfiImageLoadingProtocol { get_buffer: Some(get_buffer), ..Default::default() };
1064             let efi_entry = EfiEntry { image_handle, systab_ptr };
1065             let protocol =
1066                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
1067 
1068             let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
1069             protocol.get_buffer(&gbl_image_info).unwrap();
1070             protocol.get_buffer(&gbl_image_info).unwrap();
1071         });
1072     }
1073 
1074     #[test]
1075     #[should_panic]
test_proto_get_buffer_too_many_times()1076     fn test_proto_get_buffer_too_many_times() {
1077         unsafe extern "C" fn get_buffer(
1078             _: *mut GblEfiImageLoadingProtocol,
1079             image_info: *const GblEfiImageInfo,
1080             buffer: *mut GblEfiImageBuffer,
1081         ) -> EfiStatus {
1082             assert!(!image_info.is_null());
1083             assert!(!buffer.is_null());
1084             // SAFETY
1085             // `buffer` must be valid pointer to `GblEfiImageBuffer`
1086             let buffer = unsafe { buffer.as_mut() }.unwrap();
1087 
1088             // SAFETY:
1089             // `get_memory()` results are returned in `buffer` in `get_buffer()` function.
1090             // All usage of `get_buffer()` results are not leaving `run_test()` scope.
1091             // Same function where `start()` guard is acquired, so it will not outlive guard.
1092             unsafe {
1093                 buffer.Memory = MEMORY_TEST.with_borrow_mut(|v| v.get_memory());
1094             }
1095             buffer.SizeBytes = MEMORY_TEST_BUF_SIZE;
1096 
1097             EFI_STATUS_SUCCESS
1098         }
1099 
1100         run_test(|image_handle, systab_ptr| {
1101             let gbl_image_info: GblEfiImageInfo =
1102                 GblEfiImageInfo { ImageType: [0; PARTITION_NAME_LEN_U16], SizeBytes: 100 };
1103             let mut image_loading =
1104                 GblEfiImageLoadingProtocol { get_buffer: Some(get_buffer), ..Default::default() };
1105             let efi_entry = EfiEntry { image_handle, systab_ptr };
1106             let protocol =
1107                 generate_protocol::<GblImageLoadingProtocol>(&efi_entry, &mut image_loading);
1108 
1109             let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
1110             let mut keep_alive: Vec<EfiImageBuffer> = vec![];
1111             for _ in 1..=MAX_ARRAY_SIZE + 1 {
1112                 keep_alive.push(protocol.get_buffer(&gbl_image_info).unwrap());
1113             }
1114         });
1115     }
1116 
1117     #[test]
test_efi_image_buffer()1118     fn test_efi_image_buffer() {
1119         let mut v = vec![0u8; 1];
1120         let gbl_buffer =
1121             GblEfiImageBuffer { Memory: v.as_mut_ptr() as *mut c_void, SizeBytes: v.len() };
1122 
1123         let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
1124         // SAFETY:
1125         // 'gbl_buffer` represents valid buffer created by vector.
1126         let res = unsafe { EfiImageBuffer::new(gbl_buffer) };
1127         assert!(res.is_ok());
1128     }
1129 
1130     #[test]
test_efi_image_buffer_null()1131     fn test_efi_image_buffer_null() {
1132         let gbl_buffer = GblEfiImageBuffer { Memory: null_mut(), SizeBytes: 1 };
1133 
1134         let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
1135         // SAFETY:
1136         // 'gbl_buffer` contains Memory == NULL, which is valid input value. And we expect Error as
1137         // a result
1138         let res = unsafe { EfiImageBuffer::new(gbl_buffer) };
1139         assert_eq!(res.unwrap_err(), Error::InvalidInput);
1140     }
1141 
1142     #[test]
test_efi_image_buffer_same_buffer()1143     fn test_efi_image_buffer_same_buffer() {
1144         let mut v = vec![0u8; 1];
1145         let gbl_buffer =
1146             GblEfiImageBuffer { Memory: v.as_mut_ptr() as *mut c_void, SizeBytes: v.len() };
1147 
1148         let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
1149         // SAFETY:
1150         // 'gbl_buffer` represents valid buffer created by vector.
1151         let res1 = unsafe { EfiImageBuffer::new(gbl_buffer) };
1152         assert!(res1.is_ok());
1153 
1154         // Since we keep `res1`, second return of same buffer should fail
1155         // SAFETY:
1156         // 'gbl_buffer` represents valid buffer created by vector.
1157         let res2 = unsafe { EfiImageBuffer::new(gbl_buffer) };
1158         assert_eq!(res2.unwrap_err(), Error::AlreadyStarted);
1159     }
1160 
1161     #[test]
test_efi_image_buffer_same_buffer_after_drop()1162     fn test_efi_image_buffer_same_buffer_after_drop() {
1163         let mut v = vec![0u8; 1];
1164         let gbl_buffer =
1165             GblEfiImageBuffer { Memory: v.as_mut_ptr() as *mut c_void, SizeBytes: v.len() };
1166 
1167         let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
1168         // SAFETY:
1169         // 'gbl_buffer` represents valid buffer created by vector.
1170         let res1 = unsafe { EfiImageBuffer::new(gbl_buffer) };
1171         drop(res1);
1172 
1173         // Since `res1` was dropped same buffer can be returned.
1174         // SAFETY:
1175         // 'gbl_buffer` represents valid buffer created by vector.
1176         let res2 = unsafe { EfiImageBuffer::new(gbl_buffer) };
1177         assert!(res2.is_ok());
1178     }
1179 
1180     #[test]
test_efi_image_buffer_take()1181     fn test_efi_image_buffer_take() {
1182         let mut v = vec![0u8; 1];
1183         let gbl_buffer =
1184             GblEfiImageBuffer { Memory: v.as_mut_ptr() as *mut c_void, SizeBytes: v.len() };
1185 
1186         let _memory_guard = MEMORY_TEST.with_borrow_mut(|v| v.start());
1187         // SAFETY:
1188         // 'gbl_buffer` represents valid buffer created by vector.
1189         let res1 = unsafe { EfiImageBuffer::new(gbl_buffer) }.unwrap();
1190         let buf_no_owner = res1.take();
1191 
1192         // Since `res1` was taken, we can't reuse same buffer.
1193         // SAFETY:
1194         // 'gbl_buffer` represents valid buffer created by vector.
1195         let res2 = unsafe { EfiImageBuffer::new(gbl_buffer) };
1196         assert_eq!(res2.unwrap_err(), Error::AlreadyStarted);
1197 
1198         // Make sure to clean tracking
1199         // SAFETY:
1200         // `buf_no_owner` is the only reference to buffer
1201         unsafe {
1202             EfiImageBuffer::release(buf_no_owner.as_ptr() as usize);
1203         }
1204     }
1205 }
1206