From 650cb8985f0b763ddc184e32c74abfbc1e09d9a2 Mon Sep 17 00:00:00 2001
From: Sebastiaan van Stijn <github@gone.nl>
Date: Thu, 7 Sep 2023 23:40:52 +0200
Subject: [PATCH] add "exclude" option to --install

This allows installing "all" architectures, but skipping one, for
example, to install all architectures except for riscv64;

    binfmt --install=all,-riscv64

Signed-off-by: Sebastiaan van Stijn <github@gone.nl>
---
 README.md               |  3 +++
 cmd/binfmt/list.go      | 48 +++++++++++++++++++++++++++++++++++++++++
 cmd/binfmt/list_test.go | 24 +++++++++++++++++++++
 cmd/binfmt/main.go      | 30 ++++++++++++--------------
 4 files changed, 89 insertions(+), 16 deletions(-)
 create mode 100644 cmd/binfmt/list.go
 create mode 100644 cmd/binfmt/list_test.go

diff --git a/README.md b/README.md
index 2f824fe9..35c72873 100644
--- a/README.md
+++ b/README.md
@@ -51,6 +51,9 @@ Prints similar to:
 ```bash
 docker run --privileged --rm tonistiigi/binfmt --install all
 docker run --privileged --rm tonistiigi/binfmt --install arm64,riscv64,arm
+
+# install all architectures, except for riscv64
+docker run --privileged --rm tonistiigi/binfmt --install all,-riscv64
 ```
 
 ## Uninstalling emulators
diff --git a/cmd/binfmt/list.go b/cmd/binfmt/list.go
new file mode 100644
index 00000000..88a317ae
--- /dev/null
+++ b/cmd/binfmt/list.go
@@ -0,0 +1,48 @@
+package main
+
+import "github.com/containerd/containerd/platforms"
+
+func newList() list {
+	return list{
+		included: map[string]bool{},
+	}
+}
+
+type list struct {
+	included map[string]bool
+	items    []string
+}
+
+func (a *list) add(platform ...string) {
+	for _, v := range platform {
+		p := getArch(v)
+		if a.included[p] {
+			continue
+		}
+		a.items = append(a.items, p)
+		a.included[p] = true
+	}
+}
+
+// exclude the given platform.
+func (a *list) exclude(platform string) {
+	if p := getArch(platform); a.included[p] {
+		b := a.items[:0]
+		for _, v := range a.items {
+			if v != p {
+				b = append(b, v)
+			}
+		}
+		a.items = b
+	} else {
+		a.included[p] = true
+	}
+}
+
+func getArch(platform string) string {
+	p, err := platforms.Parse(platform)
+	if err != nil {
+		return platform
+	}
+	return p.Architecture
+}
diff --git a/cmd/binfmt/list_test.go b/cmd/binfmt/list_test.go
new file mode 100644
index 00000000..d82af9f0
--- /dev/null
+++ b/cmd/binfmt/list_test.go
@@ -0,0 +1,24 @@
+package main
+
+import (
+	"testing"
+)
+
+func TestList(t *testing.T) {
+	all := []string{"amd64", "arm64", "s390x"}
+	allItems := len(all)
+
+	a := newList()
+	a.add(all...)
+	if len(a.items) != allItems {
+		t.Errorf("expected: %d, got: %d: %v", allItems, len(a.items), a.items)
+	}
+	a.add("amd64", "amd64")
+	if len(a.items) != allItems {
+		t.Errorf("expected: %d, got: %d: %v", allItems, len(a.items), a.items)
+	}
+	a.exclude("amd64")
+	if len(a.items) != allItems-1 {
+		t.Errorf("expected: %d, got: %d: %v", allItems-1, len(a.items), a.items)
+	}
+}
diff --git a/cmd/binfmt/main.go b/cmd/binfmt/main.go
index 4de4dd3b..e203951e 100644
--- a/cmd/binfmt/main.go
+++ b/cmd/binfmt/main.go
@@ -26,7 +26,7 @@ var (
 
 func init() {
 	flag.StringVar(&mount, "mount", "/proc/sys/fs/binfmt_misc", "binfmt_misc mount point")
-	flag.StringVar(&toInstall, "install", "", "architectures to install")
+	flag.StringVar(&toInstall, "install", "", `architectures to install or exclude (prefix with "-" to exclude)`)
 	flag.StringVar(&toUninstall, "uninstall", "", "architectures to uninstall")
 	flag.BoolVar(&flVersion, "version", false, "display version")
 }
@@ -145,19 +145,24 @@ func formatPlatforms(p []ocispecs.Platform) []string {
 	return str
 }
 
-func parseArch(in string) (out []string) {
+func parseInstall(in string) (out []string) {
 	if in == "" {
 		return
 	}
+	archList := newList()
 	for _, v := range strings.Split(in, ",") {
-		p, err := platforms.Parse(v)
-		if err != nil {
-			out = append(out, v)
-		} else {
-			out = append(out, p.Architecture)
+		if v == "" {
+			continue
+		}
+		if v[0] == '-' {
+			archList.exclude(v[1:])
+			continue
+		}
+		if v == "all" {
+			archList.add(allArch()...)
 		}
 	}
-	return
+	return archList.items
 }
 
 func parseUninstall(in string) (out []string) {
@@ -211,14 +216,7 @@ func run() error {
 		}
 	}
 
-	var installArchs []string
-	if toInstall == "all" {
-		installArchs = allArch()
-	} else {
-		installArchs = parseArch(toInstall)
-	}
-
-	for _, name := range installArchs {
+	for _, name := range parseInstall(toInstall) {
 		err := install(name)
 		if err == nil {
 			log.Printf("installing: %s OK", name)