From 4f0e4270256792c368f7f5e4eddb69f080a74df8 Mon Sep 17 00:00:00 2001
From: Michael Stapelberg <michael@stapelberg.de>
Date: Sun, 17 Jun 2018 08:54:44 +0200
Subject: [PATCH] reboot: use kexec where possible
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This doesn’t work on the Raspberry Pi yet (due to missing kexec_file_load), but
I tested it in qemu.
---
 reboot.go       |  9 +++++++
 reboot_amd64.go | 65 +++++++++++++++++++++++++++++++++++++++++++++++++
 update.go       |  5 ++--
 3 files changed, 77 insertions(+), 2 deletions(-)
 create mode 100644 reboot.go
 create mode 100644 reboot_amd64.go

diff --git a/reboot.go b/reboot.go
new file mode 100644
index 0000000..081e9d2
--- /dev/null
+++ b/reboot.go
@@ -0,0 +1,9 @@
+// +build !amd64
+
+package gokrazy
+
+import "golang.org/x/sys/unix"
+
+func reboot() error {
+	return unix.Reboot(unix.LINUX_REBOOT_CMD_RESTART)
+}
diff --git a/reboot_amd64.go b/reboot_amd64.go
new file mode 100644
index 0000000..f6cc644
--- /dev/null
+++ b/reboot_amd64.go
@@ -0,0 +1,65 @@
+package gokrazy
+
+import (
+	"io/ioutil"
+	"log"
+	"os"
+	"path/filepath"
+	"syscall"
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+)
+
+// TODO: get these constants into sys/unix
+const (
+	KEXEC_ARCH_DEFAULT      = (0 << 16)
+	KEXEC_FILE_NO_INITRAMFS = 0x00000004
+)
+
+func kexecReboot() error {
+	tmpdir, err := ioutil.TempDir("", "kexec")
+	if err != nil {
+		return err
+	}
+
+	if err := syscall.Mount(mustFindRootDevice()+"1", tmpdir, "vfat", 0, ""); err != nil {
+		return err
+	}
+
+	kernel, err := os.Open(filepath.Join(tmpdir, "vmlinuz"))
+	if err != nil {
+		return err
+	}
+	defer kernel.Close()
+	cmdline, err := ioutil.ReadFile("/proc/cmdline")
+	if err != nil {
+		return err
+	}
+	rep := rootRe.ReplaceAllLiteral(cmdline, []byte("root="+mustFindRootDevice()+inactiveRootPartition))
+	// NUL-terminate cmdline
+	cmdlinebuf := make([]byte, len(rep)+1)
+	copy(cmdlinebuf, rep)
+	_, _, errno := unix.Syscall6(
+		unix.SYS_KEXEC_FILE_LOAD,
+		uintptr(kernel.Fd()), // kernel_fd
+		0,                    // initrd_fd
+		uintptr(len(cmdline)+1),                    // cmdline_len
+		uintptr(unsafe.Pointer(&cmdlinebuf[0])),    // cmdline
+		KEXEC_ARCH_DEFAULT|KEXEC_FILE_NO_INITRAMFS, // flags
+		0)
+	if errno != 0 {
+		// errno is syscall.ENOSYS on kernels without CONFIG_KEXEC_FILE_LOAD=y
+		return errno
+	}
+
+	return unix.Reboot(unix.LINUX_REBOOT_CMD_KEXEC)
+}
+
+func reboot() error {
+	if err := kexecReboot(); err != nil {
+		log.Printf("kexec reboot failed: %v", err)
+		return unix.Reboot(unix.LINUX_REBOOT_CMD_RESTART)
+	}
+	return nil
+}
diff --git a/update.go b/update.go
index 545a61b..cda7c26 100644
--- a/update.go
+++ b/update.go
@@ -21,6 +21,8 @@ import (
 var (
 	rootRe       = regexp.MustCompile(`root=/dev/(?:mmcblk0p|sda)([2-3])`)
 	rootDeviceRe = regexp.MustCompile(`root=(/dev/(?:mmcblk0p|sda))`)
+
+	inactiveRootPartition string
 )
 
 // mustFindRootDevice returns the device from which gokrazy was booted. It is
@@ -136,7 +138,6 @@ func initUpdate() error {
 	}
 
 	rootPartition := matches[1]
-	var inactiveRootPartition string
 	switch rootPartition {
 	case "2":
 		inactiveRootPartition = "3"
@@ -176,7 +177,7 @@ func initUpdate() error {
 				log.Printf("unmounting /perm failed: %v", err)
 			}
 
-			if err := unix.Reboot(unix.LINUX_REBOOT_CMD_RESTART); err != nil {
+			if err := reboot(); err != nil {
 				http.Error(w, err.Error(), http.StatusInternalServerError)
 			}
 		}()