Add fork identity wait plumbing

server/cmd/api/fork_identity.go

@@ -0,0 +1,99 @@
+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
+		}
+
+		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,107 @@
+package main
+
+import (
+	"encoding/json"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"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 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,9 @@ func main() {
 	})
 	oapi.HandlerFromMux(strictHandler, r)
 
+	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,78 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"os"
+	"path/filepath"
+	goruntime "runtime"
+
+	"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()
+		}
+		goruntime.Gosched()
+	}
+}
+
+func applyForkIdentityPayload(payload forkidentity.Payload) error {
+	for _, key := range forkidentity.ClearEnvKeys(payload) {
+		_ = os.Unsetenv(key)
+	}
+	for key, value := range forkidentity.Env(payload) {
+		_ = os.Setenv(key, value)
+	}
+	return nil
+}

server/cmd/wrapper/fork_identity_test.go

@@ -0,0 +1,50 @@
+package main
+
+import (
+	"os"
+	"testing"
+
+	"github.com/kernel/kernel-images/server/lib/forkidentity"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestApplyForkIdentityPayloadSetsAndClearsEnv(t *testing.T) {
+	t.Setenv("METRO_NAME", "old-metro")
+	t.Setenv("S2_STREAM", "old-stream")
+	t.Setenv("FUTURE_IDENTITY_FIELD_NAME", "old-future")
+
+	err := applyForkIdentityPayload(forkidentity.Payload{
+		"instance_name":               "browser-1",
+		"metro_name":                  "iad",
+		"xds_server":                  "xds.example.test",
+		"kernel_instance_jwt":         "jwt",
+		"metro_api_url":               "https://metro.example.test/browser/kernel",
+		"session_intel_url":           "https://intel.example.test",
+		"future_identity_field_name":  "future-value",
+		"empty_future_identity_field": "",
+	})
+	require.NoError(t, err)
+
+	assert.Equal(t, "browser-1", os.Getenv("INSTANCE_NAME"))
+	assert.Equal(t, "browser-1", os.Getenv("INST_NAME"))
+	assert.Equal(t, "iad", os.Getenv("METRO_NAME"))
+	assert.Equal(t, "xds.example.test", os.Getenv("XDS_SERVER"))
+	assert.Equal(t, "jwt", os.Getenv("KERNEL_INSTANCE_JWT"))
+	assert.Equal(t, "https://metro.example.test/browser/kernel", os.Getenv("KERNEL_METRO_API_BASE_URL"))
+	assert.Equal(t, "https://intel.example.test", os.Getenv("SESSION_INTEL_URL"))
+	assert.Equal(t, "future-value", os.Getenv("FUTURE_IDENTITY_FIELD_NAME"))
+	assert.Empty(t, os.Getenv("EMPTY_FUTURE_IDENTITY_FIELD"))
+	assert.Empty(t, os.Getenv("S2_STREAM"))
+}
+
+func TestForkIdentityURLPrecedence(t *testing.T) {
+	payload := forkidentity.Payload{
+		"metro_api_url":             "https://legacy.example.test/browser/kernel",
+		"kernel_metro_api_base_url": "https://metro.example.test/browser/kernel",
+		"session_intel_url":         "https://intel.example.test",
+	}
+
+	assert.Equal(t, "https://metro.example.test/browser/kernel", forkidentity.MetroAPIURL(payload))
+	assert.Equal(t, "https://intel.example.test", forkidentity.ExtensionAPIURL(payload))
+}