From c96d7de9d755621ecd94aa64964009efdb7e1909 Mon Sep 17 00:00:00 2001
From: Mark Dietzer <git@doridian.net>
Date: Tue, 12 Dec 2023 23:37:17 -0800
Subject: [PATCH] Update gokr-update-firmware to also pull *.dtbo from overlays
 (#309)

---
 .github/workflows/pull.yml           |  2 +-
 cmd/gokr-update-firmware/firmware.go | 72 +++++++++++++++++++++++-----
 2 files changed, 61 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml
index e534f68..68d641c 100644
--- a/.github/workflows/pull.yml
+++ b/.github/workflows/pull.yml
@@ -54,7 +54,7 @@ jobs:
         BOOTERY_URL: ${{ secrets.BOOTERY_URL }}
       if: ${{ env.GH_USER != 0 }}
       run: |
-        if ! gokr-has-label please-merge && ! gokr-has-label please-boot; then go install ./cmd/gokr-update-firmware && gokr-update-firmware && GOPROXY=direct go install github.com/gokrazy/autoupdate/cmd/gokr-amend@latest && gokr-amend -set_label=please-boot *.bin *.elf *.dat; fi
+        if ! gokr-has-label please-merge && ! gokr-has-label please-boot; then go install ./cmd/gokr-update-firmware && gokr-update-firmware && GOPROXY=direct go install github.com/gokrazy/autoupdate/cmd/gokr-amend@latest && gokr-amend -set_label=please-boot *.bin *.elf *.dat overlays; fi
 
     - name: Merge if boot successful
       env:
diff --git a/cmd/gokr-update-firmware/firmware.go b/cmd/gokr-update-firmware/firmware.go
index ab96e83..0845ef7 100644
--- a/cmd/gokr-update-firmware/firmware.go
+++ b/cmd/gokr-update-firmware/firmware.go
@@ -12,6 +12,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"sync"
 	"time"
 
 	"context"
@@ -78,7 +79,7 @@ func main() {
 	}
 
 	var firmwareFiles []string
-	for _, pattern := range []string{"*.elf", "*.bin", "*.dat"} {
+	for _, pattern := range []string{"*.elf", "*.bin", "*.dat", "overlays/*.dtbo"} {
 		files, err := filepath.Glob(pattern)
 		if err != nil {
 			log.Fatal(err)
@@ -87,10 +88,11 @@ func main() {
 	}
 
 	// Calculate the git blob hash of each file
-	firmwareHashes := make([]string, len(firmwareFiles))
+	firmwareHashes := make(map[string]string)
+	var firmwareHashesLock sync.Mutex
 	var eg errgroup.Group
-	for idx, path := range firmwareFiles {
-		idx, path := idx, path // copy
+	for _, path := range firmwareFiles {
+		path := path // copy
 		eg.Go(func() error {
 			hash := sha1.New()
 			f, err := os.Open(path)
@@ -108,7 +110,10 @@ func main() {
 			if _, err := io.Copy(hash, f); err != nil {
 				return err
 			}
-			firmwareHashes[idx] = fmt.Sprintf("%x", hash.Sum(nil))
+
+			firmwareHashesLock.Lock()
+			defer firmwareHashesLock.Unlock()
+			firmwareHashes[path] = fmt.Sprintf("%x", hash.Sum(nil))
 			return nil
 		})
 	}
@@ -117,23 +122,66 @@ func main() {
 		log.Fatal(err)
 	}
 
+	// Build a map of all files we want to make sure are up-to-date from GitHub
+	// map["start.elf"] = &contentEntry{...} | nil
+	filesToCheck := make(map[string]*contentEntry, 0)
+
 	contents, err := githubContents("https://api.github.com/repos/raspberrypi/firmware/contents/boot?ref=" + firmwareRef)
 	if err != nil {
 		log.Fatal(err)
 	}
 
+	// Here we only handle files directly in /boot where we only want to mirror existing files
+	for _, path := range firmwareFiles {
+		// We handle overlays below
+		if filepath.Base(filepath.Dir(path)) == "overlays" {
+			continue
+		}
+		entry, ok := contents[path]
+		if ok {
+			filesToCheck[path] = &entry
+		} else {
+			filesToCheck[path] = nil
+		}
+	}
+
+	contents, err = githubContents("https://api.github.com/repos/raspberrypi/firmware/contents/boot/overlays?ref=" + firmwareRef)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// Here we handle /boot/overlays where we want to mirror all files (especially new ones)
+	for path := range contents {
+		if filepath.Ext(path) != ".dtbo" {
+			continue
+		}
+		downloadPath := filepath.Join("overlays", filepath.Base(path))
+		entry, ok := contents[path]
+		if ok {
+			filesToCheck[downloadPath] = &entry
+		} else {
+			filesToCheck[downloadPath] = nil
+		}
+	}
+
 	ctx, canc := context.WithDeadline(context.Background(), time.Now().Add(1*time.Minute))
 	defer canc()
 	deg, ctx := errgroup.WithContext(ctx)
-	for idx, path := range firmwareFiles {
-		fn := filepath.Base(path)
-		githubContent, ok := contents[fn]
-		if !ok {
-			log.Printf("file %q not found on GitHub, obsolete?", fn)
+	for path, githubContent := range filesToCheck {
+		githubContent := githubContent
+
+		if githubContent == nil {
+			log.Printf("file %q not found on GitHub, obsolete?", path)
 			continue
 		}
-		if got, want := firmwareHashes[idx], githubContent.Sha; got != want {
-			log.Printf("getting %s (local %s, GitHub %s)", fn, got, want)
+
+		dirName := filepath.Dir(path)
+		if dirName != "" {
+			os.MkdirAll(dirName, 0755)
+		}
+
+		if got, want := firmwareHashes[path], githubContent.Sha; got != want {
+			log.Printf("getting %s (local %s, GitHub %s)", path, got, want)
 			path := path // copy
 			deg.Go(func() error {
 				req, err := http.NewRequest(http.MethodGet, githubContent.GitURL, nil)