Add fork identity wait plumbing

server/cmd/api/fork_identity.go

@@ -0,0 +1,109 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"log/slog"
+	"net/http"
+	"os"
+	"time"
+
+	"github.com/kernel/kernel-images/server/lib/forkidentity"
+)
+
+func forkIdentityHandler(log *slog.Logger) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		enabled, err := forkidentity.WaitEnabled()
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+		if !enabled {
+			http.Error(w, "fork identity wait is disabled", http.StatusConflict)
+			return
+		}
+		appliedInstance, err := forkidentity.ReadAppliedMarker()
+		if err != nil {
+			log.Error("fork identity applied marker read failed", "err", err)
+			http.Error(w, "failed to read fork identity", http.StatusInternalServerError)
+			return
+		}
+		if appliedInstance != "" {
+			http.Error(w, "fork identity already applied", http.StatusConflict)
+			return
+		}
+
+		var payload forkidentity.Payload
+		dec := json.NewDecoder(http.MaxBytesReader(w, r.Body, forkidentity.MaxPayloadBytes))
+		if err := dec.Decode(&payload); err != nil {
+			http.Error(w, fmt.Sprintf("decode payload: %v", err), http.StatusBadRequest)
+			return
+		}
+		if err := payload.Validate(); err != nil {
+			http.Error(w, err.Error(), http.StatusBadRequest)
+			return
+		}
+		if err := os.Remove(forkidentity.AppliedFile); err != nil && !os.IsNotExist(err) {
+			log.Error("fork identity applied marker reset failed", "err", err)
+			http.Error(w, "failed to reset fork identity", http.StatusInternalServerError)
+			return
+		}
+		if err := forkidentity.WritePayload(payload); err != nil {
+			log.Error("fork identity payload write failed", "err", err)
+			http.Error(w, "failed to write fork identity", http.StatusInternalServerError)
+			return
+		}
+		if err := forkidentity.WaitAppliedMarker(payload.InstanceName(), 30*time.Second); err != nil {
+			log.Error("fork identity apply wait failed", "err", err)
+			http.Error(w, "fork identity not applied", http.StatusInternalServerError)
+			return
+		}
+		w.WriteHeader(http.StatusNoContent)
+	}
+}
+
+func forkIdentityConfigHandler(log *slog.Logger) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		payload, err := forkidentity.ReadPayload()
+		if err != nil {
+			if os.IsNotExist(err) {
+				enabled, parseErr := forkidentity.WaitEnabled()
+				if parseErr != nil {
+					http.Error(w, parseErr.Error(), http.StatusInternalServerError)
+					return
+				}
+				if enabled {
+					w.WriteHeader(http.StatusAccepted)
+					return
+				}
+				http.NotFound(w, r)
+				return
+			}
+			log.Error("fork identity config read failed", "err", err)
+			http.Error(w, "failed to read fork identity", http.StatusInternalServerError)
+			return
+		}
+		enabled, err := forkidentity.WaitEnabled()
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+		if enabled {
+			applied, err := forkidentity.AppliedMarkerMatches(payload.InstanceName())
+			if err != nil {
+				log.Error("fork identity applied marker read failed", "err", err)
+				http.Error(w, "failed to read fork identity", http.StatusInternalServerError)
+				return
+			}
+			if !applied {
+				w.WriteHeader(http.StatusAccepted)
+				return
+			}
+		}
+		resp := forkidentity.ExtensionConfigFromPayload(payload)
+		w.Header().Set("Content-Type", "application/json")
+		if err := json.NewEncoder(w).Encode(resp); err != nil {
+			log.Error("fork identity config encode failed", "err", err)
+		}
+	}
+}

server/cmd/api/fork_identity_test.go

@@ -0,0 +1,180 @@
+package main
+
+import (
+	"encoding/json"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/kernel/kernel-images/server/lib/forkidentity"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestForkIdentityConfigHandlerReturnsNotFoundWithoutPayload(t *testing.T) {
+	useTempForkIdentityFiles(t)
+
+	req := httptest.NewRequest(http.MethodGet, "/internal/fork-identity/config", nil)
+	rec := httptest.NewRecorder()
+	forkIdentityConfigHandler(slog.Default()).ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusNotFound, rec.Code)
+}
+
+func TestForkIdentityConfigHandlerReturnsAcceptedWhileWaiting(t *testing.T) {
+	useTempForkIdentityFiles(t)
+	t.Setenv(forkidentity.WaitEnv, "true")
+
+	req := httptest.NewRequest(http.MethodGet, "/internal/fork-identity/config", nil)
+	rec := httptest.NewRecorder()
+	forkIdentityConfigHandler(slog.Default()).ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusAccepted, rec.Code)
+}
+
+func TestForkIdentityConfigHandlerReturnsExtensionConfig(t *testing.T) {
+	useTempForkIdentityFiles(t)
+	payload := forkidentity.Payload{
+		"instance_name":              "browser-1",
+		"metro_api_url":              "https://metro.example.test/browser/kernel",
+		"kernel_metro_api_base_url":  "https://kernel-metro.example.test/browser/kernel",
+		"session_intel_url":          "https://intel.example.test",
+		"future_identity_field_name": "future-value",
+	}
+	writeForkIdentityPayloadForTest(t, payload)
+
+	req := httptest.NewRequest(http.MethodGet, "/internal/fork-identity/config", nil)
+	rec := httptest.NewRecorder()
+	forkIdentityConfigHandler(slog.Default()).ServeHTTP(rec, req)
+
+	require.Equal(t, http.StatusOK, rec.Code)
+	var got forkidentity.ExtensionConfig
+	require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
+	assert.Equal(t, forkidentity.ExtensionConfig{
+		InstanceName: "browser-1",
+		MetroAPIURL:  "https://intel.example.test",
+	}, got)
+}
+
+func TestForkIdentityConfigHandlerReturnsAcceptedUntilPayloadApplied(t *testing.T) {
+	useTempForkIdentityFiles(t)
+	t.Setenv(forkidentity.WaitEnv, "true")
+	payload := forkidentity.Payload{
+		"instance_name":     "browser-1",
+		"session_intel_url": "https://intel.example.test",
+	}
+	writeForkIdentityPayloadForTest(t, payload)
+
+	req := httptest.NewRequest(http.MethodGet, "/internal/fork-identity/config", nil)
+	rec := httptest.NewRecorder()
+	forkIdentityConfigHandler(slog.Default()).ServeHTTP(rec, req)
+	require.Equal(t, http.StatusAccepted, rec.Code)
+
+	require.NoError(t, forkidentity.WriteAppliedMarker("browser-1"))
+	rec = httptest.NewRecorder()
+	forkIdentityConfigHandler(slog.Default()).ServeHTTP(rec, req)
+
+	require.Equal(t, http.StatusOK, rec.Code)
+	var got forkidentity.ExtensionConfig
+	require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
+	assert.Equal(t, forkidentity.ExtensionConfig{
+		InstanceName: "browser-1",
+		MetroAPIURL:  "https://intel.example.test",
+	}, got)
+}
+
+func TestForkIdentityHandlerReturnsConflictWhenDisabled(t *testing.T) {
+	useTempForkIdentityFiles(t)
+
+	req := httptest.NewRequest(http.MethodPost, "/internal/fork-identity", strings.NewReader(`{
+		"instance_name": "browser-1",
+		"session_intel_url": "https://intel.example.test"
+	}`))
+	rec := httptest.NewRecorder()
+	forkIdentityHandler(slog.Default()).ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusConflict, rec.Code)
+}
+
+func TestForkIdentityHandlerRejectsBadPayload(t *testing.T) {
+	useTempForkIdentityFiles(t)
+	t.Setenv(forkidentity.WaitEnv, "true")
+
+	req := httptest.NewRequest(http.MethodPost, "/internal/fork-identity", strings.NewReader(`{}`))
+	rec := httptest.NewRecorder()
+	forkIdentityHandler(slog.Default()).ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusBadRequest, rec.Code)
+}
+
+func TestForkIdentityHandlerRejectsAfterApplied(t *testing.T) {
+	useTempForkIdentityFiles(t)
+	t.Setenv(forkidentity.WaitEnv, "true")
+	require.NoError(t, forkidentity.WriteAppliedMarker("browser-1"))
+
+	req := httptest.NewRequest(http.MethodPost, "/internal/fork-identity", strings.NewReader(`{
+		"instance_name": "browser-2",
+		"session_intel_url": "https://intel.example.test"
+	}`))
+	rec := httptest.NewRecorder()
+	forkIdentityHandler(slog.Default()).ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusConflict, rec.Code)
+}
+
+func TestForkIdentityHandlerWritesPayloadAndWaitsForAppliedMarker(t *testing.T) {
+	useTempForkIdentityFiles(t)
+	t.Setenv(forkidentity.WaitEnv, "true")
+
+	req := httptest.NewRequest(http.MethodPost, "/internal/fork-identity", strings.NewReader(`{
+		"instance_name": "browser-1",
+		"session_intel_url": "https://intel.example.test"
+	}`))
+	rec := httptest.NewRecorder()
+	done := make(chan struct{})
+	go func() {
+		forkIdentityHandler(slog.Default()).ServeHTTP(rec, req)
+		close(done)
+	}()
+
+	require.Eventually(t, func() bool {
+		_, err := os.Stat(forkidentity.PayloadFile)
+		return err == nil
+	}, time.Second, 10*time.Millisecond)
+	payload, err := forkidentity.ReadPayload()
+	require.NoError(t, err)
+	assert.Equal(t, "browser-1", payload.InstanceName())
+
+	require.NoError(t, forkidentity.WriteAppliedMarker("browser-1"))
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("handler did not return after applied marker")
+	}
+	assert.Equal(t, http.StatusNoContent, rec.Code)
+}
+
+func useTempForkIdentityFiles(t *testing.T) {
+	t.Helper()
+	dir := t.TempDir()
+	oldPayloadFile := forkidentity.PayloadFile
+	oldAppliedFile := forkidentity.AppliedFile
+	forkidentity.PayloadFile = filepath.Join(dir, "fork-identity.json")
+	forkidentity.AppliedFile = filepath.Join(dir, "fork-identity-applied")
+	t.Cleanup(func() {
+		forkidentity.PayloadFile = oldPayloadFile
+		forkidentity.AppliedFile = oldAppliedFile
+	})
+}
+
+func writeForkIdentityPayloadForTest(t *testing.T, payload forkidentity.Payload) {
+	t.Helper()
+	data, err := json.Marshal(payload)
+	require.NoError(t, err)
+	require.NoError(t, os.WriteFile(forkidentity.PayloadFile, data, 0o600))
+}

server/cmd/api/main.go

@@ -147,6 +147,10 @@ func main() {
 	})
 	oapi.HandlerFromMux(strictHandler, r)
 
+	// Fork identity endpoints - not part of OpenAPI spec.
+	r.Post("/internal/fork-identity", forkIdentityHandler(slogger))
+	r.Get("/internal/fork-identity/config", forkIdentityConfigHandler(slogger))
+
 	// endpoints to expose the spec
 	r.Get("/spec.yaml", func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "application/vnd.oai.openapi")

server/cmd/wrapper/fork_identity.go

@@ -0,0 +1,86 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"os"
+	"path/filepath"
+	"time"
+
+	"github.com/kernel/kernel-images/server/lib/forkidentity"
+)
+
+func forkIdentityWaitEnabled() (bool, error) {
+	return forkidentity.WaitEnabled()
+}
+
+func waitForForkIdentityIfEnabled(ctx context.Context, enabled bool) bool {
+	if !enabled {
+		return true
+	}
+	stopAll("envoy")
+
+	for _, path := range []string{forkidentity.AppliedFile, forkidentity.PayloadFile} {
+		if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
+			fatalf("fork identity reset %s: %v", path, err)
+		}
+	}
+	if err := os.MkdirAll(filepath.Dir(forkidentity.ReadyFile), 0o755); err != nil {
+		fatalf("fork identity ready dir: %v", err)
+	}
+	if err := os.WriteFile(forkidentity.ReadyFile, []byte("waiting\n"), 0o644); err != nil {
+		fatalf("fork identity ready file: %v", err)
+	}
+
+	logf("fork identity waiting payload=%s", forkidentity.PayloadFile)
+	payload, err := waitForForkIdentityPayload(ctx)
+	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			logf("fork identity wait canceled")
+			return false
+		}
+		fatalf("fork identity payload wait: %v", err)
+	}
+	if err := applyForkIdentityPayload(payload); err != nil {
+		fatalf("fork identity apply: %v", err)
+	}
+	if err := forkidentity.WriteAppliedMarker(payload.InstanceName()); err != nil {
+		fatalf("fork identity applied file: %v", err)
+	}
+	logf("fork identity applied instance=%s", payload.InstanceName())
+	return true
+}
+
+func waitForForkIdentityPayload(ctx context.Context) (forkidentity.Payload, error) {
+	for {
+		payload, err := forkidentity.ReadPayload()
+		if err == nil {
+			return payload, nil
+		}
+		if !os.IsNotExist(err) {
+			return nil, err
+		}
+		if err := ctx.Err(); err != nil {
+			return nil, ctx.Err()
+		}
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-time.After(20 * time.Millisecond):
+		}
+	}
+}
+
+func applyForkIdentityPayload(payload forkidentity.Payload) error {
+	for _, key := range forkidentity.ClearEnvKeys(payload) {
+		if err := os.Unsetenv(key); err != nil {
+			return err
+		}
+	}
+	for key, value := range forkidentity.Env(payload) {
+		if err := os.Setenv(key, value); err != nil {
+			return err
+		}
+	}
+	return nil
+}