Merge pull request #7 from csssuf/user-choice

Add initial boot selection

src/picker.rs

@@ -25,8 +25,72 @@ pub mod util;
 
 pub mod uefi_entry;
 
-pub fn efi_main() -> Status {
-    let cons = uefi::get_system_table().console();
+const BOOTPATH_1: &'static str = "\\efi\\boot\\shim_a.efi";
+const BOOTPATH_2: &'static str = "\\efi\\boot\\shim_b.efi";
+
+pub fn efi_main(image_handle: Handle) -> Status {
+    let sys_table = uefi::get_system_table();
+    let cons = sys_table.console();
+
     cons.write("picker v0.0.1\r\n");
+
+    loop {
+        cons.write("Option 1: ");
+        cons.write(BOOTPATH_1);
+        cons.write("\r\n");
+
+        cons.write("Option 2: ");
+        cons.write(BOOTPATH_2);
+        cons.write("\r\n");
+
+        cons.write("Option (taking default in 5 seconds...): ");
+        match util::read_key_timeout(5000) {
+            Ok(Some(key)) => {
+                let output: [u16; 2] = [key.unicode_char, 0];
+                cons.write_raw(&output as *const u16);
+                cons.write("\r\n");
+
+                match key.unicode_char as u8 as char {
+                    '1' => {
+                        if let Err(e) = util::boot_image(BOOTPATH_1, image_handle) {
+                            cons.write("Couldn't boot choice: ");
+                            cons.write(e.str());
+                            cons.write("\r\n");
+                            return e;
+                        }
+                        break;
+                    }
+                    '2' => {
+                        if let Err(e) = util::boot_image(BOOTPATH_2, image_handle) {
+                            cons.write("Couldn't boot choice: ");
+                            cons.write(e.str());
+                            cons.write("\r\n");
+                            return e;
+                        }
+                        break;
+                    }
+                    _ => {
+                        cons.write("Unrecognized option.");
+                    }
+                }
+            }
+            Ok(None) => {
+                cons.write("\r\nTaking default.\r\n");
+                if let Err(e) = util::boot_image(BOOTPATH_1, image_handle) {
+                    cons.write("Couldn't boot default: ");
+                    cons.write(e.str());
+                    cons.write("\r\n");
+                    return e;
+                }
+                break;
+            }
+            Err(_) => {
+                cons.write("\r\nCouldn't read key.");
+            }
+        }
+
+        cons.write("\r\n");
+    }
+
     Status::Success
 }

src/uefi_entry/mod.rs

@@ -14,10 +14,19 @@
 
 extern crate uefi;
 
+use uefi::{protocol, Status};
+
 #[no_mangle]
 pub extern "win64" fn efi_entry(image_handle: uefi::Handle,
                                 system_table: *const uefi::SystemTable)
                                 -> isize {
     uefi::set_system_table(system_table).console().reset();
-    ::efi_main() as isize
+
+    let loaded_image_proto: Result<&'static protocol::LoadedImageProtocol, Status> =
+        protocol::set_current_image(image_handle);
+    if let Err(status) = loaded_image_proto {
+        return status as isize;
+    }
+
+    ::efi_main(image_handle) as isize
 }

src/util/image.rs

@@ -0,0 +1,49 @@
+// Copyright 2017 CoreOS, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+extern crate uefi;
+
+use uefi::*;
+
+fn str_to_device_path(image: &str) -> Result<&protocol::DevicePathProtocol, Status> {
+    let bs = uefi::get_system_table().boot_services();
+    bs.locate_protocol::<protocol::DevicePathFromTextProtocol>(0 as *const CVoid)
+        .and_then(|from_text| from_text.text_to_device_path_node(image))
+}
+
+fn build_boot_path(file: &protocol::DevicePathProtocol)
+                   -> Result<*const protocol::DevicePathProtocol, Status> {
+    let bs = uefi::get_system_table().boot_services();
+
+    bs.handle_protocol::<protocol::DevicePathProtocol>(protocol::get_current_image().device_handle)
+        .and_then(|this_device_path| {
+            bs.locate_protocol::<protocol::DevicePathUtilitiesProtocol>(0 as *const CVoid)
+                .and_then(|utilities| {
+                              utilities
+                                  .append_device_node(this_device_path, file)
+                                  .map(|output| output as *const protocol::DevicePathProtocol)
+                          })
+        })
+}
+
+pub fn boot_image(image: &str, parent: Handle) -> Result<(), Status> {
+    let bs = uefi::get_system_table().boot_services();
+
+    str_to_device_path(image)
+        .and_then(build_boot_path)
+        .and_then(|full_path| {
+                      bs.load_image(true, parent, full_path)
+                          .and_then(|loaded_image| bs.start_image(loaded_image))
+                  })
+}

src/util/mod.rs

@@ -14,8 +14,10 @@
 
 extern crate uefi;
 
+mod image;
 mod input;
 
+pub use self::image::*;
 pub use self::input::*;
 
 use core::fmt::Arguments;