Fetch configure extensions by reference from the metro API

server/cmd/api/api/api.go

@@ -14,6 +14,7 @@ import (
 	"github.com/kernel/kernel-images/server/lib/devtoolsproxy"
 	"github.com/kernel/kernel-images/server/lib/egresspolicy"
 	"github.com/kernel/kernel-images/server/lib/events"
+	"github.com/kernel/kernel-images/server/lib/extensionfetch"
 	"github.com/kernel/kernel-images/server/lib/logger"
 	"github.com/kernel/kernel-images/server/lib/nekoclient"
 	oapi "github.com/kernel/kernel-images/server/lib/oapi"
@@ -177,6 +178,10 @@ type ApiService struct {
 	// connection; PUT /webauthn/hook turns it on and off.
 	webauthnHook *webauthnhook.Hook
 
+	// extensionFetcher downloads the extensions a configure request names
+	// instead of carrying; nil fails such requests.
+	extensionFetcher *extensionfetch.Fetcher
+
 	lifecycleCtx    context.Context
 	lifecycleCancel context.CancelFunc
 }
@@ -196,6 +201,7 @@ func New(
 	s2Storage S2Storage,
 	egressPolicy *egresspolicy.State,
 	webauthnHook *webauthnhook.Hook,
+	extensionFetcher *extensionfetch.Fetcher,
 	processOutputConfig ProcessOutputConfig,
 ) (*ApiService, error) {
 	switch {
@@ -244,6 +250,7 @@ func New(
 		s2Storage:         s2Storage,
 		egressPolicy:      egressPolicy,
 		webauthnHook:      webauthnHook,
+		extensionFetcher:  extensionFetcher,
 		webmcp:            webmcpclient.NewManager(upstreamMgr),
 		browserRepl:       newBrowserReplManager(),
 		lifecycleCtx:      ctx,

server/cmd/api/api/api_test.go

@@ -433,7 +433,7 @@ func testWebAuthnWithMetro(t *testing.T, metroURL string) *webauthnhook.Hook {
 func newSvc(t *testing.T, mgr recorder.RecordManager) (*ApiService, error) {
 	t.Helper()
 	ts, es := newTelemetrySession(t)
-	return New(mgr, newMockFactory(), newTestUpstreamManager(), scaletozero.NewNoopController(), newMockNekoClient(t), ts, es, 0, nil, nil, testEgressPolicy(t), testWebAuthn(t), ProcessOutputConfig{})
+	return New(mgr, newMockFactory(), newTestUpstreamManager(), scaletozero.NewNoopController(), newMockNekoClient(t), ts, es, 0, nil, nil, testEgressPolicy(t), testWebAuthn(t), nil, ProcessOutputConfig{})
 }
 
 func TestApiService_PatchChromiumFlags(t *testing.T) {

server/cmd/api/api/chromium_configure.go

@@ -13,9 +13,11 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/kernel/kernel-images/server/lib/cdpclient"
+	"github.com/kernel/kernel-images/server/lib/extensionfetch"
 	"github.com/kernel/kernel-images/server/lib/logger"
 	oapi "github.com/kernel/kernel-images/server/lib/oapi"
 	"github.com/kernel/kernel-images/server/lib/policy"
@@ -25,6 +27,7 @@ import (
 const userDataProfileDir = "/home/kernel/user-data"
 const maxStartURLLen = 2048
 const startURLDispatchTimeout = 3 * time.Second
+const maxExtensionRefBytes = 8 << 10
 
 type chromiumConfigureState struct {
 	displayJSON        *string
@@ -39,6 +42,7 @@ type chromiumConfigureState struct {
 	startURLRaw *string
 
 	extItems []extensionZipItem // zipTemp paths; merged with chromiumCfgParseExtensions
+	extRefs  []extensionfetch.Ref
 
 	allTemps []string
 }
@@ -73,6 +77,10 @@ func (s *ApiService) ChromiumConfigure(ctx context.Context, request oapi.Chromiu
 	}
 	defer st.cleanup()
 
+	if err := s.fetchExtensionRefs(ctx, st); err != nil {
+		return cfg500ConfigureStep(chromiumConfigureStepExtensionFetch, err.Error()), nil
+	}
+
 	spec, msgs := chromiumStartURLSpec(st.startURLRaw)
 	if msgs != "" {
 		return cfg400(msgs), nil
@@ -355,6 +363,9 @@ const (
 	chromiumConfigureStepDisplay    chromiumConfigureStep = "display"
 	chromiumConfigureStepFlags      chromiumConfigureStep = "chromium_flags"
 	chromiumConfigureStepProfile    chromiumConfigureStep = "profile"
+	// chromiumConfigureStepExtensionFetch fails before anything is applied,
+	// so the caller can resend the request with the archives uploaded.
+	chromiumConfigureStepExtensionFetch chromiumConfigureStep = "extension_fetch"
 )
 
 func chromiumStartURLSpec(raw *string) (startURLParsed, string) {
@@ -646,6 +657,19 @@ func chromiumCfgParseMultipart(body interface{}, st *chromiumConfigureState) err
 			}
 			cur.zipTmp = tmp.Name()
 			cur.gotZip = true
+		case "extensions.ref":
+			b, err := io.ReadAll(io.LimitReader(part, maxExtensionRefBytes+1))
+			if err != nil {
+				return cfgParseInternal("read extensions.ref")
+			}
+			var ref extensionfetch.Ref
+			if len(b) > maxExtensionRefBytes || json.Unmarshal(b, &ref) != nil {
+				return cfgParseBadRequest("invalid extensions.ref")
+			}
+			if err := ref.Validate(); err != nil {
+				return cfgParseBadRequest(fmt.Sprintf("invalid extensions.ref: %v", err))
+			}
+			st.extRefs = append(st.extRefs, ref)
 		case "extensions.name":
 			if cur == nil {
 				cur = &pend{}
@@ -673,6 +697,48 @@ func chromiumCfgParseMultipart(body interface{}, st *chromiumConfigureState) err
 	if cur != nil && (!cur.gotZip || cur.name == "") {
 		return cfgParseBadRequest("each extension pair needs extensions.zip_file plus extensions.name")
 	}
+	if len(st.extItems)+len(st.extRefs) > maxExtensionUploadCount {
+		return cfgParseBadRequest(fmt.Sprintf("too many extensions; maximum is %d", maxExtensionUploadCount))
+	}
+	return nil
+}
+
+// fetchExtensionRefs downloads the extensions the request names from the
+// metro API, in parallel, and adds them to the request's extensions as if
+// they had been uploaded. It runs before anything is applied.
+func (s *ApiService) fetchExtensionRefs(ctx context.Context, st *chromiumConfigureState) error {
+	if len(st.extRefs) == 0 {
+		return nil
+	}
+	if s.extensionFetcher == nil {
+		return errors.New("extension fetch is not configured")
+	}
+	items := make([]extensionZipItem, len(st.extRefs))
+	for i, ref := range st.extRefs {
+		tmp, err := os.CreateTemp("", "bcc-ext-*.zip")
+		if err != nil {
+			return fmt.Errorf("temp file for %s: %w", ref.Name, err)
+		}
+		_ = tmp.Close()
+		st.allTemps = append(st.allTemps, tmp.Name())
+		items[i] = extensionZipItem{zipTemp: tmp.Name(), name: ref.Name}
+	}
+	errs := make([]error, len(st.extRefs))
+	var wg sync.WaitGroup
+	for i, ref := range st.extRefs {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if err := s.extensionFetcher.Fetch(ctx, ref, items[i].zipTemp); err != nil {
+				errs[i] = fmt.Errorf("fetch %s: %w", ref.Name, err)
+			}
+		}()
+	}
+	wg.Wait()
+	if err := errors.Join(errs...); err != nil {
+		return err
+	}
+	st.extItems = append(st.extItems, items...)
 	return nil
 }
 

server/cmd/api/api/chromium_configure_refs_test.go

@@ -0,0 +1,219 @@
+package api
+
+import (
+	"bytes"
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"mime/multipart"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/kernel/kernel-images/server/lib/extensionfetch"
+	oapi "github.com/kernel/kernel-images/server/lib/oapi"
+	"github.com/stretchr/testify/require"
+)
+
+func refPart(t *testing.T, w *multipart.Writer, ref extensionfetch.Ref) {
+	t.Helper()
+	b, err := json.Marshal(ref)
+	require.NoError(t, err)
+	require.NoError(t, w.WriteField("extensions.ref", string(b)))
+}
+
+func sumOf(b []byte) string {
+	s := sha256.Sum256(b)
+	return hex.EncodeToString(s[:])
+}
+
+func metroFetcher(t *testing.T, archives map[string][]byte) *extensionfetch.Fetcher {
+	t.Helper()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, ok := archives[strings.TrimPrefix(r.URL.Path, "/extension-archive/")]
+		if !ok {
+			http.Error(w, "not found", http.StatusNotFound)
+			return
+		}
+		_, _ = w.Write(body)
+	}))
+	t.Cleanup(srv.Close)
+	return &extensionfetch.Fetcher{
+		BaseURL: func() string { return srv.URL },
+		Token:   func() string { return "instance-jwt" },
+		Client:  srv.Client(),
+	}
+}
+
+func TestChromiumCfgParseMultipartExtensionRefs(t *testing.T) {
+	buf := bytes.NewBuffer(nil)
+	w := multipart.NewWriter(buf)
+	refPart(t, w, extensionfetch.Ref{Name: "ext1", SHA256: sumOf([]byte("a")), Grant: "g1"})
+	refPart(t, w, extensionfetch.Ref{Name: "ext2", SHA256: sumOf([]byte("b")), Grant: "g2"})
+	require.NoError(t, w.Close())
+
+	st := &chromiumConfigureState{}
+	err := chromiumCfgParseMultipart(multipart.NewReader(buf, w.Boundary()), st)
+	defer st.cleanup()
+	require.NoError(t, err)
+	require.Len(t, st.extRefs, 2)
+	require.Equal(t, "ext1", st.extRefs[0].Name)
+	require.Equal(t, "g2", st.extRefs[1].Grant)
+	require.Empty(t, st.extItems, "references are fetched after parsing")
+}
+
+func TestChromiumCfgParseMultipartExtensionRefValidation(t *testing.T) {
+	for name, value := range map[string]string{
+		"not json":       "{",
+		"bad checksum":   `{"name":"ext1","sha256":"abc","grant":"g"}`,
+		"path traversal": `{"name":"../ext1","sha256":"` + sumOf(nil) + `","grant":"g"}`,
+		"too large":      `{"name":"ext1","sha256":"` + sumOf(nil) + `","grant":"` + strings.Repeat("g", maxExtensionRefBytes) + `"}`,
+	} {
+		t.Run(name, func(t *testing.T) {
+			buf := bytes.NewBuffer(nil)
+			w := multipart.NewWriter(buf)
+			require.NoError(t, w.WriteField("extensions.ref", value))
+			require.NoError(t, w.Close())
+
+			st := &chromiumConfigureState{}
+			err := chromiumCfgParseMultipart(multipart.NewReader(buf, w.Boundary()), st)
+			defer st.cleanup()
+			require.ErrorContains(t, err, "invalid extensions.ref")
+		})
+	}
+}
+
+func TestChromiumCfgParseMultipartExtensionLimit(t *testing.T) {
+	zipPair := func(t *testing.T, w *multipart.Writer, name string) {
+		t.Helper()
+		part, err := w.CreateFormFile("extensions.zip_file", name+".zip")
+		require.NoError(t, err)
+		_, err = part.Write([]byte("zip"))
+		require.NoError(t, err)
+		require.NoError(t, w.WriteField("extensions.name", name))
+	}
+	ref := func(t *testing.T, w *multipart.Writer, name string) {
+		t.Helper()
+		refPart(t, w, extensionfetch.Ref{Name: name, SHA256: sumOf(nil), Grant: "g"})
+	}
+	parse := func(t *testing.T, build func(*multipart.Writer)) error {
+		t.Helper()
+		buf := bytes.NewBuffer(nil)
+		w := multipart.NewWriter(buf)
+		build(w)
+		require.NoError(t, w.Close())
+		st := &chromiumConfigureState{}
+		defer st.cleanup()
+		return chromiumCfgParseMultipart(multipart.NewReader(buf, w.Boundary()), st)
+	}
+
+	// The cap counts uploads and references together, whatever their order.
+	for name, build := range map[string]func(*multipart.Writer){
+		"references": func(w *multipart.Writer) {
+			for i := 0; i <= maxExtensionUploadCount; i++ {
+				ref(t, w, fmt.Sprintf("r%d", i))
+			}
+		},
+		"uploads": func(w *multipart.Writer) {
+			for i := 0; i <= maxExtensionUploadCount; i++ {
+				zipPair(t, w, fmt.Sprintf("u%d", i))
+			}
+		},
+		"references first": func(w *multipart.Writer) {
+			for i := 0; i < maxExtensionUploadCount; i++ {
+				ref(t, w, fmt.Sprintf("r%d", i))
+			}
+			zipPair(t, w, "u")
+		},
+		"uploads first": func(w *multipart.Writer) {
+			for i := 0; i < maxExtensionUploadCount; i++ {
+				zipPair(t, w, fmt.Sprintf("u%d", i))
+			}
+			ref(t, w, "r")
+		},
+	} {
+		t.Run(name, func(t *testing.T) {
+			require.ErrorContains(t, parse(t, build), "too many extensions")
+		})
+	}
+
+	t.Run("at the cap", func(t *testing.T) {
+		require.NoError(t, parse(t, func(w *multipart.Writer) {
+			for i := 0; i < maxExtensionUploadCount/2; i++ {
+				zipPair(t, w, fmt.Sprintf("u%d", i))
+				ref(t, w, fmt.Sprintf("r%d", i))
+			}
+		}))
+	})
+}
+
+func TestFetchExtensionRefs(t *testing.T) {
+	archives := map[string][]byte{"ext1": []byte("zip-one"), "ext2": []byte("zip-two")}
+	s := &ApiService{extensionFetcher: metroFetcher(t, archives)}
+	st := &chromiumConfigureState{
+		extItems: []extensionZipItem{{zipTemp: "uploaded.zip", name: "uploaded"}},
+		extRefs: []extensionfetch.Ref{
+			{Name: "ext1", SHA256: sumOf(archives["ext1"]), Grant: "g"},
+			{Name: "ext2", SHA256: sumOf(archives["ext2"]), Grant: "g"},
+		},
+	}
+	defer st.cleanup()
+
+	require.NoError(t, s.fetchExtensionRefs(context.Background(), st))
+	require.Len(t, st.extItems, 3)
+	require.Equal(t, "uploaded", st.extItems[0].name)
+	for i, name := range []string{"ext1", "ext2"} {
+		item := st.extItems[i+1]
+		require.Equal(t, name, item.name)
+		got, err := os.ReadFile(item.zipTemp)
+		require.NoError(t, err)
+		require.Equal(t, archives[name], got)
+	}
+}
+
+func TestFetchExtensionRefsFailure(t *testing.T) {
+	archives := map[string][]byte{"ext1": []byte("zip-one")}
+	refs := []extensionfetch.Ref{
+		{Name: "ext1", SHA256: sumOf(archives["ext1"]), Grant: "g"},
+		{Name: "missing", SHA256: sumOf([]byte("x")), Grant: "g"},
+	}
+
+	t.Run("one fetch fails", func(t *testing.T) {
+		s := &ApiService{extensionFetcher: metroFetcher(t, archives)}
+		st := &chromiumConfigureState{extRefs: refs}
+		defer st.cleanup()
+		err := s.fetchExtensionRefs(context.Background(), st)
+		require.ErrorContains(t, err, "fetch missing: status 404")
+		require.Empty(t, st.extItems)
+	})
+
+	t.Run("no fetcher", func(t *testing.T) {
+		s := &ApiService{}
+		st := &chromiumConfigureState{extRefs: refs}
+		defer st.cleanup()
+		require.ErrorContains(t, s.fetchExtensionRefs(context.Background(), st), "not configured")
+	})
+}
+
+func TestChromiumConfigureExtensionFetchFailureAppliesNothing(t *testing.T) {
+	s := &ApiService{extensionFetcher: metroFetcher(t, nil)}
+
+	buf := bytes.NewBuffer(nil)
+	w := multipart.NewWriter(buf)
+	refPart(t, w, extensionfetch.Ref{Name: "ext1", SHA256: sumOf([]byte("a")), Grant: "g"})
+	require.NoError(t, w.WriteField("chromium_flags", `{"flags":["--kiosk"]}`))
+	require.NoError(t, w.Close())
+
+	resp, err := s.ChromiumConfigure(context.Background(), oapi.ChromiumConfigureRequestObject{
+		Body: multipart.NewReader(buf, w.Boundary()),
+	})
+	require.NoError(t, err)
+	failure, ok := resp.(oapi.ChromiumConfigure500JSONResponse)
+	require.True(t, ok, "got %T", resp)
+	require.NotNil(t, failure.Step)
+	require.Equal(t, oapi.ExtensionFetch, *failure.Step)
+}

server/cmd/api/api/display_test.go

@@ -39,7 +39,7 @@ func testFFmpegFactory(t *testing.T, tempDir string) recorder.FFmpegRecorderFact
 func newTestServiceWithFactory(t *testing.T, mgr recorder.RecordManager, factory recorder.FFmpegRecorderFactory) *ApiService {
 	t.Helper()
 	ts, es := newTelemetrySession(t)
-	svc, err := New(mgr, factory, newTestUpstreamManager(), scaletozero.NewNoopController(), newMockNekoClient(t), ts, es, 0, nil, nil, testEgressPolicy(t), testWebAuthn(t), ProcessOutputConfig{})
+	svc, err := New(mgr, factory, newTestUpstreamManager(), scaletozero.NewNoopController(), newMockNekoClient(t), ts, es, 0, nil, nil, testEgressPolicy(t), testWebAuthn(t), nil, ProcessOutputConfig{})
 	require.NoError(t, err)
 	return svc
 }

server/cmd/api/api/telemetry_test.go

@@ -409,7 +409,7 @@ func (m *mockRecordManager) StopAll(_ context.Context) error
 func newTestService(t *testing.T, mgr recorder.RecordManager) *ApiService {
 	t.Helper()
 	ts, es := newTelemetrySession(t)
-	svc, err := New(mgr, newMockFactory(), newTestUpstreamManager(), scaletozero.NewNoopController(), newMockNekoClient(t), ts, es, 0, nil, nil, testEgressPolicy(t), testWebAuthn(t), ProcessOutputConfig{})
+	svc, err := New(mgr, newMockFactory(), newTestUpstreamManager(), scaletozero.NewNoopController(), newMockNekoClient(t), ts, es, 0, nil, nil, testEgressPolicy(t), testWebAuthn(t), nil, ProcessOutputConfig{})
 	require.NoError(t, err)
 	svc.cdpMonitor = &stubCdpMonitor{}
 	return svc