diff --git a/cmd/mcpproxy/listen_flag_test.go b/cmd/mcpproxy/listen_flag_test.go index e38f14c66..03e1c791e 100644 --- a/cmd/mcpproxy/listen_flag_test.go +++ b/cmd/mcpproxy/listen_flag_test.go @@ -55,7 +55,7 @@ func TestLoadConfig_ListenFlag(t *testing.T) { t.Fatal(err) } - cfg, err := loadConfig(cmd) + cfg, _, err := loadConfig(cmd) if err != nil { t.Fatalf("loadConfig: %v", err) } diff --git a/internal/upstream/core/instance.go b/internal/upstream/core/instance.go index db6ec93e0..f08736182 100644 --- a/internal/upstream/core/instance.go +++ b/internal/upstream/core/instance.go @@ -10,47 +10,143 @@ import ( "github.com/google/uuid" ) +const instanceIDFileName = "instance-id" + var ( instanceID string instanceIDOnce sync.Once + + dataDirMu sync.Mutex + dataDir string + + // legacyInstanceIDPath is a var (not a const) so tests can point it at a + // scratch path instead of the real host-wide file. + legacyInstanceIDPath = func() string { + return filepath.Join(os.TempDir(), "mcpproxy-instance-id") + } ) -// getInstanceID returns a unique identifier for this mcpproxy instance -// The ID is persisted across restarts and used to label Docker containers +// SetInstanceDataDir records the data directory this process's instance ID +// should be persisted under (normally cfg.DataDir). Call it once at startup, +// before the first GetInstanceID() call — e.g. from upstream.NewManager, +// which every entry point constructs early with the loaded config. +// +// Concurrent mcpproxy processes on the same host must already use distinct +// data directories (BBolt takes an exclusive lock on config.db), so keying +// the instance ID off the data dir rather than a single shared file under +// os.TempDir() gives each process a genuinely distinct ID. A call after the +// ID has already been resolved, or no call at all, has no effect beyond that +// first read: GetInstanceID() then falls back to a fresh id for the +// process's lifetime instead of reusing a host-wide file. +func SetInstanceDataDir(dir string) { + dataDirMu.Lock() + dataDir = dir + dataDirMu.Unlock() +} + +// getInstanceID returns a unique identifier for this mcpproxy instance, +// resolved once per process and cached for the process's lifetime. func getInstanceID() string { instanceIDOnce.Do(func() { - // Try to load from file first - if id, err := loadInstanceID(); err == nil && id != "" { - instanceID = id - return - } - - // Generate new instance ID - instanceID = uuid.New().String() - _ = saveInstanceID(instanceID) // Best effort save + dataDirMu.Lock() + dir := dataDir + dataDirMu.Unlock() + instanceID = resolveInstanceID(dir) }) return instanceID } +// resolveInstanceID contains the actual id-resolution logic, kept free of +// package-level state so it can be exercised directly (and repeatedly, with +// different dirs) in tests without the sync.Once in getInstanceID hiding +// everything but the first call. +func resolveInstanceID(dir string) string { + if dir == "" { + // No data directory known at labeling time: use a fresh id for this + // process's lifetime rather than the old host-wide shared temp file, + // which made every mcpproxy process on a machine collide on one ID. + return uuid.New().String() + } + + if id, err := loadInstanceID(dir); err == nil && id != "" { + return id + } + + // First run under this data dir: adopt the legacy host-wide id if one is + // still there, so containers created before this fix stay manageable by + // whichever instance starts first after the upgrade. Then retire the + // legacy file so no other data dir can adopt the same id afterwards -- + // otherwise every future data dir would keep re-adopting it forever, + // recreating the exact bug this is fixing. + if id := adoptLegacyInstanceID(dir); id != "" { + return id + } + + id := uuid.New().String() + _ = saveInstanceID(dir, id) // Best effort save + return id +} + +// adoptLegacyInstanceID migrates the pre-fix, host-wide shared instance id +// (if present) into dataDir. Returns "" if there is no legacy file to adopt. +// +// Claiming the legacy file happens via os.Rename to a globally-unique path +// (a fresh UUID, not e.g. os.Getpid()) rather than a plain read-then-remove: +// rename atomically fails if the source is already gone, so when two +// processes race to adopt the same legacy file at upgrade time, exactly one +// wins and the other correctly falls through to generating its own fresh id. +// A read-then-remove would let both processes read the same id before +// either removed the file, recreating the original host-wide-shared-id bug +// for that pair. The claim destination must itself never collide between +// racers -- a PID-based suffix alone doesn't guarantee that (PIDs repeat +// across a crashed-and-restarted process, and are identical across +// goroutines within one process), so a random UUID is used instead. +func adoptLegacyInstanceID(dataDir string) string { + claimPath := fmt.Sprintf("%s.claimed-%s", legacyInstanceIDPath(), uuid.New().String()) + if err := os.Rename(legacyInstanceIDPath(), claimPath); err != nil { + // No legacy file, or another process already claimed it. + return "" + } + + data, err := os.ReadFile(claimPath) + if err != nil { + return "" + } + id := strings.TrimSpace(string(data)) + if id == "" { + _ = os.Remove(claimPath) + return "" + } + + if err := saveInstanceID(dataDir, id); err != nil { + // Persistence under the new per-data-dir location failed: leave the + // claimed content at claimPath instead of deleting it, so it isn't + // silently lost and can still be recovered manually. This process + // still uses id for its own lifetime -- the same best-effort + // tolerance the fresh-uuid path below already has for a failed save. + return id + } + _ = os.Remove(claimPath) + return id +} + // GetInstanceID returns the unique identifier for this mcpproxy instance (exported for use by manager) func GetInstanceID() string { return getInstanceID() } -// loadInstanceID attempts to load the instance ID from disk -func loadInstanceID() (string, error) { - instanceFile := filepath.Join(os.TempDir(), "mcpproxy-instance-id") - data, err := os.ReadFile(instanceFile) +// loadInstanceID attempts to load the instance ID from disk under dataDir +func loadInstanceID(dataDir string) (string, error) { + data, err := os.ReadFile(filepath.Join(dataDir, instanceIDFileName)) if err != nil { return "", err } return strings.TrimSpace(string(data)), nil } -// saveInstanceID persists the instance ID to disk -func saveInstanceID(id string) error { - instanceFile := filepath.Join(os.TempDir(), "mcpproxy-instance-id") - return os.WriteFile(instanceFile, []byte(id), 0644) +// saveInstanceID persists the instance ID to disk under dataDir +func saveInstanceID(dataDir, id string) error { + return os.WriteFile(filepath.Join(dataDir, instanceIDFileName), []byte(id), 0o600) } // formatContainerLabels returns Docker labels for container ownership tracking diff --git a/internal/upstream/core/instance_test.go b/internal/upstream/core/instance_test.go new file mode 100644 index 000000000..a8c60fb06 --- /dev/null +++ b/internal/upstream/core/instance_test.go @@ -0,0 +1,339 @@ +package core + +import ( + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "testing" + + "github.com/google/uuid" +) + +// withLegacyInstanceIDPath points the legacy (pre-fix, host-wide) instance id +// file at a path under the test's own temp dir instead of the real +// os.TempDir(), so tests don't read or clobber a real machine's legacy file +// and don't race other tests/processes that touch it. +func withLegacyInstanceIDPath(t *testing.T, path string) { + t.Helper() + original := legacyInstanceIDPath + legacyInstanceIDPath = func() string { return path } + t.Cleanup(func() { legacyInstanceIDPath = original }) +} + +func TestResolveInstanceIDUniquePerDataDir(t *testing.T) { + withLegacyInstanceIDPath(t, filepath.Join(t.TempDir(), "no-legacy-file")) + + dir1 := t.TempDir() + dir2 := t.TempDir() + + id1 := resolveInstanceID(dir1) + id2 := resolveInstanceID(dir2) + + if id1 == id2 { + t.Fatalf("expected distinct instance ids for distinct data dirs, got %q for both", id1) + } + if _, err := uuid.Parse(id1); err != nil { + t.Errorf("id1 %q is not a valid UUID: %v", id1, err) + } + if _, err := uuid.Parse(id2); err != nil { + t.Errorf("id2 %q is not a valid UUID: %v", id2, err) + } +} + +func TestResolveInstanceIDPersistsAcrossCalls(t *testing.T) { + withLegacyInstanceIDPath(t, filepath.Join(t.TempDir(), "no-legacy-file")) + + dir := t.TempDir() + + first := resolveInstanceID(dir) + second := resolveInstanceID(dir) + + if first != second { + t.Fatalf("expected the same data dir to resolve the same instance id across calls (simulating a restart), got %q then %q", first, second) + } +} + +func TestResolveInstanceIDNoDataDirReturnsFreshID(t *testing.T) { + withLegacyInstanceIDPath(t, filepath.Join(t.TempDir(), "no-legacy-file")) + + id1 := resolveInstanceID("") + id2 := resolveInstanceID("") + + if id1 == id2 { + t.Fatalf("expected fresh ids each time no data dir is available, got %q for both", id1) + } +} + +func TestResolveInstanceIDAdoptsLegacySharedFileOnce(t *testing.T) { + legacyPath := filepath.Join(t.TempDir(), "mcpproxy-instance-id") + withLegacyInstanceIDPath(t, legacyPath) + + legacyID := uuid.New().String() + if err := os.WriteFile(legacyPath, []byte(legacyID), 0o600); err != nil { + t.Fatalf("failed to seed legacy instance id file: %v", err) + } + + dataDir := t.TempDir() + adopted := resolveInstanceID(dataDir) + + if adopted != legacyID { + t.Fatalf("expected the first data dir to adopt the legacy shared id %q, got %q", legacyID, adopted) + } + if _, err := os.Stat(legacyPath); !os.IsNotExist(err) { + t.Fatalf("expected the legacy shared file to be removed after adoption, stat err = %v", err) + } + + // A second, distinct data dir must NOT also adopt the (now-consumed) + // legacy id -- that would recreate the original host-wide-shared-id bug + // for every data dir created after the upgrade. + otherDataDir := t.TempDir() + otherID := resolveInstanceID(otherDataDir) + if otherID == legacyID { + t.Fatalf("a second data dir must not also adopt the already-consumed legacy id %q", legacyID) + } +} + +func TestAdoptLegacyInstanceIDConcurrentClaimIsExclusive(t *testing.T) { + legacyPath := filepath.Join(t.TempDir(), "mcpproxy-instance-id") + withLegacyInstanceIDPath(t, legacyPath) + + legacyID := uuid.New().String() + if err := os.WriteFile(legacyPath, []byte(legacyID), 0o600); err != nil { + t.Fatalf("failed to seed legacy instance id file: %v", err) + } + + const racers = 8 + dataDirs := make([]string, racers) + for i := range dataDirs { + dataDirs[i] = t.TempDir() + } + + start := make(chan struct{}) + results := make([]string, racers) + var wg sync.WaitGroup + for i := 0; i < racers; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + results[i] = adoptLegacyInstanceID(dataDirs[i]) + }(i) + } + close(start) + wg.Wait() + + adopters := 0 + for _, id := range results { + if id == legacyID { + adopters++ + } else if id != "" { + t.Errorf("adoptLegacyInstanceID returned an unexpected non-empty, non-legacy id %q", id) + } + } + if adopters != 1 { + t.Fatalf("expected exactly one of %d concurrent racers to adopt the legacy id, got %d", racers, adopters) + } + if _, err := os.Stat(legacyPath); !os.IsNotExist(err) { + t.Fatalf("expected the legacy shared file to be gone after the race, stat err = %v", err) + } +} + +func TestAdoptLegacyInstanceIDPreservesClaimWhenSaveFails(t *testing.T) { + legacyPath := filepath.Join(t.TempDir(), "mcpproxy-instance-id") + withLegacyInstanceIDPath(t, legacyPath) + + legacyID := uuid.New().String() + if err := os.WriteFile(legacyPath, []byte(legacyID), 0o600); err != nil { + t.Fatalf("failed to seed legacy instance id file: %v", err) + } + + // A plain file (not a directory) as the "data dir" makes saveInstanceID's + // os.WriteFile(filepath.Join(dataDir, instanceIDFileName), ...) fail on + // every OS, without relying on permission bits that behave differently + // on Windows. + unwritableDataDir := filepath.Join(t.TempDir(), "not-a-directory") + if err := os.WriteFile(unwritableDataDir, []byte("x"), 0o600); err != nil { + t.Fatalf("failed to seed the not-a-directory stand-in: %v", err) + } + + id := adoptLegacyInstanceID(unwritableDataDir) + if id != legacyID { + t.Fatalf("expected adoptLegacyInstanceID to still return the legacy id %q despite the failed save, got %q", legacyID, id) + } + + matches, err := filepath.Glob(legacyPath + ".claimed-*") + if err != nil { + t.Fatalf("glob failed: %v", err) + } + if len(matches) != 1 { + t.Fatalf("expected the claimed legacy content to survive a failed save (found %d claim files, want 1): %v", len(matches), matches) + } + claimed, err := os.ReadFile(matches[0]) + if err != nil { + t.Fatalf("failed to read the surviving claim file: %v", err) + } + if strings.TrimSpace(string(claimed)) != legacyID { + t.Fatalf("surviving claim file content = %q, want %q", claimed, legacyID) + } +} + +// helperProcessEnvVar and friends drive a re-exec of this test binary as a +// standalone helper process, so GetInstanceID's process-wide sync.Once is +// exercised fresh -- other tests in this package (e.g. isolation_*_test.go, +// via BuildDockerArgs) already call GetInstanceID indirectly, which would +// otherwise cache a result before this test runs and make in-process testing +// of the singleton order-dependent. +// +// helperProcessLegacyPathEnvVar is required, not optional: without it the +// helper process would fall through to the real legacyInstanceIDPath() +// default (the actual host-wide os.TempDir() file), and could adopt-and-delete +// a genuine machine's pre-upgrade migration file as a side effect of running +// this test suite. +const ( + helperProcessEnvVar = "MCPPROXY_INSTANCE_ID_TEST_HELPER" + helperProcessDataDirEnvVar = "MCPPROXY_INSTANCE_ID_TEST_DATA_DIR" + helperProcessLegacyPathEnvVar = "MCPPROXY_INSTANCE_ID_TEST_LEGACY_PATH" +) + +// TestHelperProcess is not a real test; it's invoked as a subprocess by the +// tests below. See https://pkg.go.dev/os/exec#Command for this pattern. +func TestHelperProcess(t *testing.T) { + if os.Getenv(helperProcessEnvVar) != "1" { + return + } + legacyPath := os.Getenv(helperProcessLegacyPathEnvVar) + if legacyPath == "" { + t.Fatal("helper process requires " + helperProcessLegacyPathEnvVar + " to avoid touching the real host-wide legacy file") + } + legacyInstanceIDPath = func() string { return legacyPath } + SetInstanceDataDir(os.Getenv(helperProcessDataDirEnvVar)) + fmt.Print(GetInstanceID()) + os.Exit(0) +} + +// runInstanceIDHelperProcess runs GetInstanceID() (via SetInstanceDataDir) in +// a fresh subprocess scoped to dataDir, with the legacy shared-id path +// scoped to a scratch location under dataDir so the helper never touches a +// real machine's host-wide legacy file. Returns what it printed. +func runInstanceIDHelperProcess(t *testing.T, dataDir string) string { + t.Helper() + out, err := runInstanceIDHelperProcessWithLegacyPath(dataDir, filepath.Join(dataDir, "no-legacy-file")) + if err != nil { + t.Fatalf("%v", err) + } + return out +} + +// runInstanceIDHelperProcessWithLegacyPath is like runInstanceIDHelperProcess +// but lets the caller point multiple helper processes at the SAME legacy +// path, to exercise the real cross-process legacy-claim race. It returns an +// error instead of calling t.Fatalf directly so it's safe to invoke from a +// spawned goroutine (testing.T.Fatal/FailNow must only be called from the +// test's own goroutine). +func runInstanceIDHelperProcessWithLegacyPath(dataDir, legacyPath string) (string, error) { + cmd := exec.Command(os.Args[0], "-test.run=TestHelperProcess") + cmd.Env = append(os.Environ(), + helperProcessEnvVar+"=1", + helperProcessDataDirEnvVar+"="+dataDir, + helperProcessLegacyPathEnvVar+"="+legacyPath, + ) + out, err := cmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("helper process failed: %w\noutput: %s", err, out) + } + return string(out), nil +} + +func TestGetInstanceIDReturnsValidUUIDPersistedUnderDataDir(t *testing.T) { + dir := t.TempDir() + + id := runInstanceIDHelperProcess(t, dir) + if _, err := uuid.Parse(id); err != nil { + t.Fatalf("GetInstanceID() = %q is not a valid UUID: %v", id, err) + } + + persisted, err := loadInstanceID(dir) + if err != nil { + t.Fatalf("expected instance id to be persisted under the data dir: %v", err) + } + if persisted != id { + t.Fatalf("persisted instance id %q does not match GetInstanceID() %q", persisted, id) + } +} + +func TestGetInstanceIDDistinctAcrossConcurrentProcesses(t *testing.T) { + dir1 := t.TempDir() + dir2 := t.TempDir() + + var id1, id2 string + var err1, err2 error + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + id1, err1 = runInstanceIDHelperProcessWithLegacyPath(dir1, filepath.Join(dir1, "no-legacy-file")) + }() + go func() { + defer wg.Done() + id2, err2 = runInstanceIDHelperProcessWithLegacyPath(dir2, filepath.Join(dir2, "no-legacy-file")) + }() + wg.Wait() + + if err1 != nil { + t.Fatalf("%v", err1) + } + if err2 != nil { + t.Fatalf("%v", err2) + } + if id1 == id2 { + t.Fatalf("expected two mcpproxy processes with distinct data dirs to get distinct instance ids, got %q for both", id1) + } +} + +// TestGetInstanceIDCrossProcessLegacyClaimIsExclusive is the real-world +// counterpart to TestAdoptLegacyInstanceIDConcurrentClaimIsExclusive: that +// test proves the claim is exclusive between goroutines in one process +// (which, unrealistically, all share one PID); this one starts two genuinely +// separate OS processes racing over the SAME seeded legacy file and checks +// exactly one of them adopts it. +func TestGetInstanceIDCrossProcessLegacyClaimIsExclusive(t *testing.T) { + legacyPath := filepath.Join(t.TempDir(), "mcpproxy-instance-id") + legacyID := uuid.New().String() + if err := os.WriteFile(legacyPath, []byte(legacyID), 0o600); err != nil { + t.Fatalf("failed to seed legacy instance id file: %v", err) + } + + dir1 := t.TempDir() + dir2 := t.TempDir() + + var id1, id2 string + var err1, err2 error + var wg sync.WaitGroup + wg.Add(2) + go func() { defer wg.Done(); id1, err1 = runInstanceIDHelperProcessWithLegacyPath(dir1, legacyPath) }() + go func() { defer wg.Done(); id2, err2 = runInstanceIDHelperProcessWithLegacyPath(dir2, legacyPath) }() + wg.Wait() + + if err1 != nil { + t.Fatalf("%v", err1) + } + if err2 != nil { + t.Fatalf("%v", err2) + } + + adopters := 0 + for _, id := range []string{id1, id2} { + if id == legacyID { + adopters++ + } + } + if adopters != 1 { + t.Fatalf("expected exactly one of 2 concurrent mcpproxy processes to adopt the legacy id %q, got %d (id1=%q id2=%q)", legacyID, adopters, id1, id2) + } + if _, err := os.Stat(legacyPath); !os.IsNotExist(err) { + t.Fatalf("expected the legacy shared file to be gone after the race, stat err = %v", err) + } +} diff --git a/internal/upstream/manager.go b/internal/upstream/manager.go index 0f6d3e6f5..42bf50ea3 100644 --- a/internal/upstream/manager.go +++ b/internal/upstream/manager.go @@ -193,6 +193,15 @@ func cloneServerConfig(cfg *config.ServerConfig) *config.ServerConfig { // NewManager creates a new upstream manager func NewManager(logger *zap.Logger, globalConfig *config.Config, boltStorage *storage.BoltDB, secretResolver *secret.Resolver, storageMgr *storage.Manager) *Manager { + // Scope this process's Docker container-ownership instance ID to its data + // dir instead of a host-wide shared file, so concurrent mcpproxy + // processes on one host (which already require distinct data dirs, since + // BBolt locks config.db) get distinct IDs. Must happen before any code + // path calls core.GetInstanceID(), so it's done here, at the earliest + // point every entry point (serve, tray, CLI subcommands) has the loaded + // config available. + core.SetInstanceDataDir(globalConfig.DataDir) + shutdownCtx, shutdownCancel := context.WithCancel(context.Background()) manager := &Manager{ clients: make(map[string]*managed.Client),