xref: /aosp_15_r20/bootable/libbootloader/gbl/efi/src/fastboot.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 // This EFI application implements a demo for booting Android/Fuchsia from disk. See
16 // bootable/libbootloader/gbl/README.md for how to run the demo. See comments of
17 // `android_boot:android_boot_demo()` and `fuchsia_boot:fuchsia_boot_demo()` for
18 // supported/unsupported features at the moment.
19 
20 use crate::{
21     net::{EfiGblNetwork, EfiTcpSocket},
22     ops::Ops,
23 };
24 use alloc::{boxed::Box, vec::Vec};
25 use core::{cmp::min, fmt::Write, future::Future, mem::take, pin::Pin, sync::atomic::AtomicU64};
26 use efi::{
27     efi_print, efi_println,
28     protocol::{gbl_efi_fastboot_usb::GblFastbootUsbProtocol, Protocol},
29     EfiEntry,
30 };
31 use fastboot::{TcpStream, Transport};
32 use gbl_async::{block_on, YieldCounter};
33 use liberror::{Error, Result};
34 use libgbl::fastboot::{run_gbl_fastboot, GblTcpStream, GblUsbTransport, PinFutContainer};
35 
36 const DEFAULT_TIMEOUT_MS: u64 = 5_000;
37 const FASTBOOT_TCP_PORT: u16 = 5554;
38 
39 struct EfiFastbootTcpTransport<'a, 'b, 'c> {
40     socket: &'c mut EfiTcpSocket<'a, 'b>,
41 }
42 
43 impl<'a, 'b, 'c> EfiFastbootTcpTransport<'a, 'b, 'c> {
new(socket: &'c mut EfiTcpSocket<'a, 'b>) -> Self44     fn new(socket: &'c mut EfiTcpSocket<'a, 'b>) -> Self {
45         Self { socket: socket }
46     }
47 }
48 
49 impl TcpStream for EfiFastbootTcpTransport<'_, '_, '_> {
50     /// Reads to `out` for exactly `out.len()` number bytes from the TCP connection.
read_exact(&mut self, out: &mut [u8]) -> Result<()>51     async fn read_exact(&mut self, out: &mut [u8]) -> Result<()> {
52         self.socket.receive_exact(out, DEFAULT_TIMEOUT_MS).await
53     }
54 
55     /// Sends exactly `data.len()` number bytes from `data` to the TCP connection.
write_exact(&mut self, data: &[u8]) -> Result<()>56     async fn write_exact(&mut self, data: &[u8]) -> Result<()> {
57         self.socket.send_exact(data, DEFAULT_TIMEOUT_MS).await
58     }
59 }
60 
61 impl GblTcpStream for EfiFastbootTcpTransport<'_, '_, '_> {
accept_new(&mut self) -> bool62     fn accept_new(&mut self) -> bool {
63         let efi_entry = self.socket.efi_entry;
64         self.socket.poll();
65         // If not listenining, start listening.
66         // If not connected but it's been `DEFAULT_TIMEOUT_MS`, restart listening in case the remote
67         // client disconnects in the middle of TCP handshake and leaves the socket in a half open
68         // state.
69         if !self.socket.is_listening_or_handshaking()
70             || (!self.socket.check_active()
71                 && self.socket.time_since_last_listen() > DEFAULT_TIMEOUT_MS)
72         {
73             let _ = self
74                 .socket
75                 .listen(FASTBOOT_TCP_PORT)
76                 .inspect_err(|e| efi_println!(efi_entry, "TCP listen error: {:?}", e));
77 
78             // TODO(b/368647237): Enable only in Fuchsia context.
79             self.socket.broadcast_fuchsia_fastboot_mdns();
80         } else if self.socket.check_active() {
81             self.socket.set_io_yield_threshold(1024 * 1024); // 1MB
82             let remote = self.socket.get_socket().remote_endpoint().unwrap();
83             efi_println!(efi_entry, "TCP connection from {}", remote);
84             return true;
85         }
86         false
87     }
88 }
89 
90 /// `UsbTransport` implements the `fastboot::Transport` trait using USB interfaces from
91 /// GBL_EFI_FASTBOOT_USB_PROTOCOL.
92 pub struct UsbTransport<'a> {
93     max_packet_size: usize,
94     protocol: Protocol<'a, GblFastbootUsbProtocol>,
95     io_yield_counter: YieldCounter,
96     // Buffer for prefetching an incoming packet in `wait_for_packet()`.
97     // Alternatively we can also consider adding an EFI event for packet arrive. But UEFI firmware
98     // may be more complicated.
99     prefetched: (Vec<u8>, usize),
100 }
101 
102 impl<'a> UsbTransport<'a> {
new(max_packet_size: usize, protocol: Protocol<'a, GblFastbootUsbProtocol>) -> Self103     fn new(max_packet_size: usize, protocol: Protocol<'a, GblFastbootUsbProtocol>) -> Self {
104         Self {
105             max_packet_size,
106             protocol,
107             io_yield_counter: YieldCounter::new(1024 * 1024),
108             prefetched: (vec![0u8; max_packet_size], 0),
109         }
110     }
111 
112     /// Polls and cache the next USB packet.
113     ///
114     /// Returns Ok(true) if there is a new packet. Ok(false) if there is no incoming packet. Err()
115     /// otherwise.
poll_next_packet(&mut self) -> Result<bool>116     fn poll_next_packet(&mut self) -> Result<bool> {
117         match &mut self.prefetched {
118             (pkt, len) if *len == 0 => match self.protocol.fastboot_usb_receive(pkt) {
119                 Ok(out_size) => {
120                     *len = out_size;
121                     return Ok(true);
122                 }
123                 Err(Error::NotReady) => return Ok(false),
124                 Err(e) => return Err(e),
125             },
126             _ => Ok(true),
127         }
128     }
129 }
130 
131 impl Transport for UsbTransport<'_> {
receive_packet(&mut self, out: &mut [u8]) -> Result<usize>132     async fn receive_packet(&mut self, out: &mut [u8]) -> Result<usize> {
133         let len = match &mut self.prefetched {
134             (pkt, len) if *len > 0 => {
135                 let out = out.get_mut(..*len).ok_or(Error::BufferTooSmall(Some(*len)))?;
136                 let src = pkt.get(..*len).ok_or(Error::Other(Some("Invalid USB read size")))?;
137                 out.clone_from_slice(src);
138                 take(len)
139             }
140             _ => self.protocol.receive_packet(out).await?,
141         };
142         // Forces a yield to the executor if the data received/sent reaches a certain
143         // threshold. This is to prevent the async code from holding up the CPU for too long
144         // in case IO speed is high and the executor uses cooperative scheduling.
145         self.io_yield_counter.increment(len.try_into().unwrap()).await;
146         Ok(len)
147     }
148 
send_packet(&mut self, packet: &[u8]) -> Result<()>149     async fn send_packet(&mut self, packet: &[u8]) -> Result<()> {
150         let mut curr = &packet[..];
151         while !curr.is_empty() {
152             let to_send = min(curr.len(), self.max_packet_size);
153             self.protocol.send_packet(&curr[..to_send], DEFAULT_TIMEOUT_MS).await?;
154             // Forces a yield to the executor if the data received/sent reaches a certain
155             // threshold. This is to prevent the async code from holding up the CPU for too long
156             // in case IO speed is high and the executor uses cooperative scheduling.
157             self.io_yield_counter.increment(to_send.try_into().unwrap()).await;
158             curr = &curr[to_send..];
159         }
160         Ok(())
161     }
162 }
163 
164 impl GblUsbTransport for UsbTransport<'_> {
has_packet(&mut self) -> bool165     fn has_packet(&mut self) -> bool {
166         let efi_entry = self.protocol.efi_entry();
167         self.poll_next_packet()
168             .inspect_err(|e| efi_println!(efi_entry, "Error while polling next packet: {:?}", e))
169             .unwrap_or(false)
170     }
171 }
172 
173 /// Initializes the Fastboot USB interface and returns a `UsbTransport`.
init_usb(efi_entry: &EfiEntry) -> Result<UsbTransport>174 fn init_usb(efi_entry: &EfiEntry) -> Result<UsbTransport> {
175     let protocol =
176         efi_entry.system_table().boot_services().find_first_and_open::<GblFastbootUsbProtocol>()?;
177     match protocol.fastboot_usb_interface_stop() {
178         Err(e) if e != Error::NotStarted => Err(e),
179         _ => Ok(UsbTransport::new(protocol.fastboot_usb_interface_start()?, protocol)),
180     }
181 }
182 
183 // Wrapper of vector of pinned futures.
184 #[derive(Default)]
185 struct VecPinFut<'a>(Vec<Pin<Box<dyn Future<Output = ()> + 'a>>>);
186 
187 impl<'a> PinFutContainer<'a> for VecPinFut<'a> {
add_with<F: Future<Output = ()> + 'a>(&mut self, f: impl FnOnce() -> F)188     fn add_with<F: Future<Output = ()> + 'a>(&mut self, f: impl FnOnce() -> F) {
189         self.0.push(Box::pin(f()));
190     }
191 
for_each_remove_if( &mut self, mut cb: impl FnMut(&mut Pin<&mut (dyn Future<Output = ()> + 'a)>) -> bool, )192     fn for_each_remove_if(
193         &mut self,
194         mut cb: impl FnMut(&mut Pin<&mut (dyn Future<Output = ()> + 'a)>) -> bool,
195     ) {
196         for idx in (0..self.0.len()).rev() {
197             cb(&mut self.0[idx].as_mut()).then(|| self.0.swap_remove(idx));
198         }
199     }
200 }
201 
fastboot(efi_gbl_ops: &mut Ops, bootimg_buf: &mut [u8]) -> Result<()>202 pub fn fastboot(efi_gbl_ops: &mut Ops, bootimg_buf: &mut [u8]) -> Result<()> {
203     let efi_entry = efi_gbl_ops.efi_entry;
204     efi_println!(efi_entry, "Entering fastboot mode...");
205 
206     let usb = init_usb(efi_entry)
207         .inspect(|_| efi_println!(efi_entry, "Started Fastboot over USB."))
208         .inspect_err(|e| efi_println!(efi_entry, "Failed to start Fastboot over USB. {:?}.", e))
209         .ok();
210 
211     let ts = AtomicU64::new(0);
212     let mut net: EfiGblNetwork = Default::default();
213     let mut tcp = net
214         .init(efi_entry, &ts)
215         .inspect(|v| {
216             efi_println!(efi_entry, "Started Fastboot over TCP");
217             efi_println!(efi_entry, "IP address:");
218             v.interface().ip_addrs().iter().for_each(|v| {
219                 efi_println!(efi_entry, "\t{}", v.address());
220             });
221         })
222         .inspect_err(|e| efi_println!(efi_entry, "Failed to start EFI network. {:?}.", e))
223         .ok();
224     let tcp = tcp.as_mut().map(|v| EfiFastbootTcpTransport::new(v));
225 
226     let download_buffers = vec![vec![0u8; 512 * 1024 * 1024]; 2].into();
227     block_on(run_gbl_fastboot(
228         efi_gbl_ops,
229         &download_buffers,
230         VecPinFut::default(),
231         usb,
232         tcp,
233         bootimg_buf,
234     ));
235 
236     efi_println!(efi_entry, "Leaving fastboot mode...");
237 
238     Ok(())
239 }
240