core/cmd/mkimage: clean up
This cleans up a few minor issues with the mkimage tool:
- error logging now goes to stderr; it going to stdout made my life
painful when debugging a build failure as the output got mangled by
Bazel
- repeated actions (copies, mkdirs) are now expressed less verbosely
- flags are now declared in a more canonical way
- flags are now named in a more canonical way (underscores vs. hyphens;
these haven't been used anywhere in the build, so we're not breaking
anything).
Test Plan: Refactor, no change behaviour (apart from stderr logging).
X-Origin-Diff: phab/D558
GitOrigin-RevId: 2531562acc46ec7aa0a6ac89b883dfd45cddcf82
diff --git a/core/cmd/mkimage/main.go b/core/cmd/mkimage/main.go
index 574842f..9f49f0a 100644
--- a/core/cmd/mkimage/main.go
+++ b/core/cmd/mkimage/main.go
@@ -16,10 +16,14 @@
package main
+// mkimage is a tool to generate a Smalltown disk image containing the given EFI payload, and optionally, a given external
+// initramfs image and enrolment credentials.
+
import (
"flag"
"fmt"
"io/ioutil"
+ "log"
"os"
diskfs "github.com/diskfs/go-diskfs"
@@ -31,12 +35,12 @@
var SmalltownDataPartition gpt.Type = gpt.Type("9eeec464-6885-414a-b278-4305c51f7966")
var (
- efiPayloadPath = flag.String("efi", "", "UEFI payload")
- outputPath = flag.String("out", "", "Output disk image")
- initramfsPath = flag.String("initramfs", "", "External initramfs [optional]")
- enrolmentCredentialsPath = flag.String("enrolment-credentials", "", "Enrolment credentials [optional]")
- dataPartitionSizeMiB = flag.Uint64("data-partition-size", 2048, "Override the data partition size (default 2048 MiB)")
- espPartitionSizeMiB = flag.Uint64("esp-partition-size", 512, "Override the ESP partition size (default: 512MiB)")
+ flagEFI string
+ flagOut string
+ flagInitramfs string
+ flagEnrolmentCredentials string
+ flagDataPartitionSize uint64
+ flagESPPartitionSize uint64
)
func mibToSectors(size uint64) uint64 {
@@ -44,17 +48,22 @@
}
func main() {
+ flag.StringVar(&flagEFI, "efi", "", "UEFI payload")
+ flag.StringVar(&flagOut, "out", "", "Output disk image")
+ flag.StringVar(&flagInitramfs, "initramfs", "", "External initramfs [optional]")
+ flag.StringVar(&flagEnrolmentCredentials, "enrolment_credentials", "", "Enrolment credentials [optional]")
+ flag.Uint64Var(&flagDataPartitionSize, "data_partition_size", 2048, "Override the data partition size (default 2048 MiB)")
+ flag.Uint64Var(&flagESPPartitionSize, "esp_partition_size", 512, "Override the ESP partition size (default: 512MiB)")
flag.Parse()
- if *efiPayloadPath == "" || *outputPath == "" {
- flag.PrintDefaults()
- os.Exit(2)
+
+ if flagEFI == "" || flagOut == "" {
+ log.Fatalf("efi and initramfs must be set")
}
- _ = os.Remove(*outputPath)
- diskImg, err := diskfs.Create(*outputPath, 3*1024*1024*1024, diskfs.Raw)
+ _ = os.Remove(flagOut)
+ diskImg, err := diskfs.Create(flagOut, 3*1024*1024*1024, diskfs.Raw)
if err != nil {
- fmt.Printf("Failed to create disk: %v", err)
- os.Exit(1)
+ log.Fatalf("diskfs.Create(%q): %v", flagOut, err)
}
table := &gpt.Table{
@@ -67,97 +76,66 @@
Type: gpt.EFISystemPartition,
Name: "ESP",
Start: mibToSectors(1),
- End: mibToSectors(*espPartitionSizeMiB) - 1,
+ End: mibToSectors(flagESPPartitionSize) - 1,
},
{
Type: SmalltownDataPartition,
Name: "SIGNOS-DATA",
- Start: mibToSectors(*espPartitionSizeMiB),
- End: mibToSectors(*espPartitionSizeMiB+*dataPartitionSizeMiB) - 1,
+ Start: mibToSectors(flagESPPartitionSize),
+ End: mibToSectors(flagESPPartitionSize+flagDataPartitionSize) - 1,
},
},
}
if err := diskImg.Partition(table); err != nil {
- fmt.Printf("Failed to apply partition table: %v", err)
- os.Exit(1)
+ log.Fatalf("Failed to apply partition table: %v", err)
}
fs, err := diskImg.CreateFilesystem(disk.FilesystemSpec{Partition: 1, FSType: filesystem.TypeFat32, VolumeLabel: "ESP"})
if err != nil {
- fmt.Printf("Failed to create filesystem: %v", err)
- os.Exit(1)
+ log.Fatalf("Failed to create filesystem: %v", err)
}
- if err := fs.Mkdir("/EFI"); err != nil {
- panic(err)
- }
- if err := fs.Mkdir("/EFI/BOOT"); err != nil {
- panic(err)
- }
- if err := fs.Mkdir("/EFI/smalltown"); err != nil {
- panic(err)
- }
- efiPayload, err := fs.OpenFile("/EFI/BOOT/BOOTX64.EFI", os.O_CREATE|os.O_RDWR)
- if err != nil {
- fmt.Printf("Failed to open EFI payload for writing: %v", err)
- os.Exit(1)
- }
- efiPayloadSrc, err := os.Open(*efiPayloadPath)
- if err != nil {
- fmt.Printf("Failed to open EFI payload for reading: %v", err)
- os.Exit(1)
- }
- defer efiPayloadSrc.Close()
- // If this is streamed (e.g. using io.Copy) it exposes a bug in diskfs, so do it in one go.
- efiPayloadFull, err := ioutil.ReadAll(efiPayloadSrc)
- if err != nil {
- panic(err)
- }
- if _, err := efiPayload.Write(efiPayloadFull); err != nil {
- fmt.Printf("Failed to write EFI payload: %v", err)
- os.Exit(1)
- }
- initramfs, err := fs.OpenFile("/EFI/smalltown/initramfs.cpio.lz4", os.O_CREATE|os.O_RDWR)
- if err != nil {
- fmt.Printf("Failed to open initramfs for writing: %v", err)
- os.Exit(1)
- }
- // If we have more than two arguments, the second one is the initramfs
- if *initramfsPath != "" {
- initramfsSrc, err := os.Open(*initramfsPath)
- if err != nil {
- fmt.Printf("Failed to open initramfs for reading: %v", err)
- os.Exit(1)
- }
- initramfsFull, err := ioutil.ReadAll(initramfsSrc)
- if err != nil {
- fmt.Printf("Failed to read initramfs: %v", err)
- os.Exit(1)
- }
- if _, err := initramfs.Write(initramfsFull); err != nil {
- fmt.Printf("Failed to write initramfs: %v", err)
- os.Exit(1)
+
+ // Create EFI partition structure.
+ for _, dir := range []string{"/EFI", "/EFI/BOOT", "/EFI/smalltown"} {
+ if err := fs.Mkdir(dir); err != nil {
+ log.Fatalf("Mkdir(%q): %v", dir, err)
}
}
- if *enrolmentCredentialsPath != "" {
- enrolmentCredentials, err := fs.OpenFile("/EFI/smalltown/enrolment.pb", os.O_CREATE|os.O_RDWR)
- enrolmentCredentialsSrc, err := os.Open(*enrolmentCredentialsPath)
- if err != nil {
- fmt.Printf("Failed to open enrolment credentials for reading: %v", err)
- os.Exit(1)
- }
- enrolmentCredentialsFull, err := ioutil.ReadAll(enrolmentCredentialsSrc)
- if err != nil {
- fmt.Printf("Failed to read enrolment credentials: %v", err)
- os.Exit(1)
- }
- if _, err := enrolmentCredentials.Write(enrolmentCredentialsFull); err != nil {
- fmt.Printf("Failed to write enrolment credentials")
- os.Exit(1)
- }
+
+ put(fs, flagEFI, "/EFI/BOOT/BOOTX64.EFI")
+
+ if flagInitramfs != "" {
+ put(fs, flagInitramfs, "/EFI/smalltown/initramfs.cpio.lz4")
}
+
+ if flagEnrolmentCredentials != "" {
+ put(fs, flagEnrolmentCredentials, "/EFI/smalltown/enrolment.pb")
+ }
+
if err := diskImg.File.Close(); err != nil {
- fmt.Printf("Failed to write image: %v", err)
+ log.Fatalf("Failed to finalize image: %v", err)
+ }
+ log.Printf("Success! You can now boot %v", flagOut)
+}
+
+// put copies a file from the host filesystem into the target image.
+func put(fs filesystem.FileSystem, src, dst string) {
+ target, err := fs.OpenFile(dst, os.O_CREATE|os.O_RDWR)
+ if err != nil {
+ log.Fatalf("fs.OpenFile(%q): %v", dst, err)
+ }
+ source, err := os.Open(src)
+ if err != nil {
+ log.Fatalf("os.Open(%q): %v", src, err)
+ }
+ defer source.Close()
+ // If this is streamed (e.g. using io.Copy) it exposes a bug in diskfs, so do it in one go.
+ data, err := ioutil.ReadAll(source)
+ if err != nil {
+ log.Fatalf("Reading %q: %v", src, err)
+ }
+ if _, err := target.Write(data); err != nil {
+ fmt.Printf("writing file %q: %v", dst, err)
os.Exit(1)
}
- fmt.Println("Success! You can now boot smalltown.img")
}