feat: add arm64/aarch64 multi-arch support (#44)

* feat: add arm64/aarch64 multi-arch support

Add support for arm64 architecture:

- Map arm64 to aarch64 in GetArch() for kernel/binary paths
- Map arm64 to aarch64 in cloud-hypervisor binary extraction
- Add platform-specific OCI image pulling with remote.WithPlatform()
  to correctly pull the right architecture for multi-arch images
- Add golang/protobuf dependency for grpc compatibility

* Refactor inspectManifest to use remote.Image for platform-specific digests (#46)

Co-authored-by: Cursor Agent <[email protected]>

---------

Co-authored-by: Cursor Agent <[email protected]>

Author: rgarcia

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/ghodss/yaml v1.0.0
 	github.com/go-chi/chi/v5 v5.2.3
 	github.com/golang-jwt/jwt/v5 v5.3.0
+	github.com/golang/protobuf v1.5.4
 	github.com/google/go-containerregistry v0.20.6
 	github.com/google/wire v0.7.0
 	github.com/gorilla/websocket v1.5.3
@@ -42,7 +43,6 @@ require (
 	golang.org/x/sync v0.17.0
 	golang.org/x/sys v0.38.0
 	google.golang.org/grpc v1.77.0
-	google.golang.org/protobuf v1.36.10
 	gvisor.dev/gvisor v0.0.0-20251125014920-fc40e232ff54
 )
 
@@ -112,6 +112,7 @@ require (
 	golang.org/x/tools v0.37.0 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20251022142026-3a174f9686a8 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20251022142026-3a174f9686a8 // indirect
+	google.golang.org/protobuf v1.36.10 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	gotest.tools/v3 v3.5.2 // indirect

lib/images/oci.go

@@ -4,10 +4,12 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"runtime"
 	"strings"
 
 	"github.com/google/go-containerregistry/pkg/authn"
 	"github.com/google/go-containerregistry/pkg/name"
+	gcr "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/empty"
 	"github.com/google/go-containerregistry/pkg/v1/layout"
 	"github.com/google/go-containerregistry/pkg/v1/remote"
@@ -60,24 +62,42 @@ func newOCIClient(cacheDir string) (*ociClient, error) {
 	return &ociClient{cacheDir: cacheDir}, nil
 }
 
+// currentPlatform returns the platform for the current host
+func currentPlatform() gcr.Platform {
+	return gcr.Platform{
+		Architecture: runtime.GOARCH,
+		OS:           runtime.GOOS,
+	}
+}
+
 // inspectManifest synchronously inspects a remote image to get its digest
 // without pulling the image. This is used for upfront digest discovery.
+// For multi-arch images, it returns the platform-specific manifest digest
+// (matching the current host platform) rather than the manifest index digest.
 func (c *ociClient) inspectManifest(ctx context.Context, imageRef string) (string, error) {
 	ref, err := name.ParseReference(imageRef)
 	if err != nil {
 		return "", fmt.Errorf("parse image reference: %w", err)
 	}
 
-	// Use system authentication (reads from ~/.docker/config.json, etc.)
-	// Default retry: only on network errors, max ~1.3s total
-	descriptor, err := remote.Head(ref,
+	// Use remote.Image with platform filtering to get the platform-specific digest.
+	// For multi-arch images, this resolves the manifest index to the correct platform.
+	// This matches what pullToOCILayout does to ensure cache key consistency.
+	// Note: remote.Image is lazy - it only fetches the manifest, not layer blobs.
+	img, err := remote.Image(ref,
 		remote.WithContext(ctx),
-		remote.WithAuthFromKeychain(authn.DefaultKeychain))
+		remote.WithAuthFromKeychain(authn.DefaultKeychain),
+		remote.WithPlatform(currentPlatform()))
 	if err != nil {
 		return "", fmt.Errorf("fetch manifest: %w", wrapRegistryError(err))
 	}
 
-	return descriptor.Digest.String(), nil
+	digest, err := img.Digest()
+	if err != nil {
+		return "", fmt.Errorf("get image digest: %w", err)
+	}
+
+	return digest.String(), nil
 }
 
 // pullResult contains the metadata and digest from pulling an image
@@ -126,9 +146,11 @@ func (c *ociClient) pullToOCILayout(ctx context.Context, imageRef, layoutTag str
 
 	// Use system authentication (reads from ~/.docker/config.json, etc.)
 	// Default retry: only on network errors, max ~1.3s total
+	// WithPlatform ensures we pull the correct architecture for multi-arch images
 	img, err := remote.Image(ref,
 		remote.WithContext(ctx),
-		remote.WithAuthFromKeychain(authn.DefaultKeychain))
+		remote.WithAuthFromKeychain(authn.DefaultKeychain),
+		remote.WithPlatform(currentPlatform()))
 	if err != nil {
 		// Rate limits fail here immediately (429 is not retried by default)
 		return fmt.Errorf("fetch image manifest: %w", wrapRegistryError(err))

lib/system/versions.go

@@ -74,5 +74,8 @@ func GetArch() string {
 	if arch == "amd64" {
 		return "x86_64"
 	}
+	if arch == "arm64" {
+		return "aarch64"
+	}
 	return arch
 }

lib/vmm/binaries.go

@@ -30,6 +30,8 @@ func ExtractBinary(p *paths.Paths, version CHVersion) (string, error) {
 	arch := runtime.GOARCH
 	if arch == "amd64" {
 		arch = "x86_64"
+	} else if arch == "arm64" {
+		arch = "aarch64"
 	}
 
 	embeddedPath := fmt.Sprintf("binaries/cloud-hypervisor/%s/%s/cloud-hypervisor", version, arch)