test(api,db): cover snapshot team scoping

Add tests for the ENG-3544 team scoping at every layer:
- db: GetLastSnapshot returns the snapshot for the owning team, ErrNoRows
  for another team, and the snapshot when no team is given.
- cache: SnapshotCache isolates teams by cache key and Invalidate clears
  the team-scoped key via prefix deletion.
- handlers: pauseHandleNotRunningSandbox returns 404 for a foreign team
  and 409 for the owner; get/connect/resume return 404 end to end when
  the snapshot belongs to another team (real Postgres+Redis containers,
  mocked orchestrator).

Also add the InsertTestSnapshot/EnvBuild test query helpers.

Author: adababys

packages/api/internal/cache/snapshots/snapshot_cache_test.go

@@ -0,0 +1,85 @@
+package snapshotcache
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/db/pkg/testutils"
+	redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis"
+)
+
+// TestSnapshotCache_TeamScopedIsolationAcrossTeams verifies that a snapshot fetched
+// through the cache is scoped to the owning team: the owner reads it, while another
+// team querying the identical sandbox ID gets a not-found result (no cross-team leak
+// and no cache-key collision between teams).
+func TestSnapshotCache_TeamScopedIsolationAcrossTeams(t *testing.T) {
+	t.Parallel()
+	db := testutils.SetupDatabase(t)
+	redisClient := redis_utils.SetupInstance(t)
+	ctx := t.Context()
+
+	sc := NewSnapshotCache(db.SqlcClient, redisClient)
+	defer sc.Close(ctx)
+
+	ownerTeamID := testutils.CreateTestTeam(t, db)
+	baseTemplateID := testutils.CreateTestTemplate(t, db, ownerTeamID)
+
+	sandboxID := "sandbox-" + uuid.New().String()
+	snapshotTemplateID := "snapshot-template-" + uuid.New().String()
+	result := testutils.UpsertTestSnapshot(t, ctx, db, snapshotTemplateID, sandboxID, ownerTeamID, baseTemplateID)
+
+	// Owner team reads its own snapshot.
+	info, err := sc.Get(ctx, sandboxID, ownerTeamID)
+	require.NoError(t, err)
+	assert.Equal(t, result.BuildID, info.EnvBuild.ID,
+		"owning team should read the snapshot build through the cache")
+
+	// A different team must not see it, even though the sandbox ID is identical.
+	otherTeamID := testutils.CreateTestTeam(t, db)
+	_, err = sc.Get(ctx, sandboxID, otherTeamID)
+	require.ErrorIs(t, err, ErrSnapshotNotFound,
+		"a non-owning team must not read another team's snapshot")
+}
+
+// TestSnapshotCache_InvalidateClearsTeamScopedKey verifies that Invalidate clears the
+// team-scoped cache key (sandboxID:teamID) and not only the plain sandboxID key.
+// It first caches a negative (not-found) result under the team-scoped key, then proves
+// the negative result is served from cache, and finally that Invalidate forces a
+// re-fetch that now finds the snapshot.
+func TestSnapshotCache_InvalidateClearsTeamScopedKey(t *testing.T) {
+	t.Parallel()
+	db := testutils.SetupDatabase(t)
+	redisClient := redis_utils.SetupInstance(t)
+	ctx := t.Context()
+
+	sc := NewSnapshotCache(db.SqlcClient, redisClient)
+	defer sc.Close(ctx)
+
+	teamID := testutils.CreateTestTeam(t, db)
+	baseTemplateID := testutils.CreateTestTemplate(t, db, teamID)
+
+	sandboxID := "sandbox-" + uuid.New().String()
+	snapshotTemplateID := "snapshot-template-" + uuid.New().String()
+
+	// No snapshot yet: not-found, cached under the team-scoped key.
+	_, err := sc.Get(ctx, sandboxID, teamID)
+	require.ErrorIs(t, err, ErrSnapshotNotFound)
+
+	// The snapshot now appears in the DB.
+	result := testutils.UpsertTestSnapshot(t, ctx, db, snapshotTemplateID, sandboxID, teamID, baseTemplateID)
+
+	// The cached negative result is still served, proving the team-scoped key is cached.
+	_, err = sc.Get(ctx, sandboxID, teamID)
+	require.ErrorIs(t, err, ErrSnapshotNotFound,
+		"the team-scoped negative result should be served from cache before invalidation")
+
+	// Invalidate must delete the team-scoped key (sandboxID:teamID) via prefix deletion.
+	sc.Invalidate(ctx, sandboxID)
+
+	info, err := sc.Get(ctx, sandboxID, teamID)
+	require.NoError(t, err, "after Invalidate the team-scoped key must be re-fetched from the DB")
+	assert.Equal(t, result.BuildID, info.EnvBuild.ID)
+}

packages/api/internal/handlers/sandbox_pause_test.go

@@ -0,0 +1,74 @@
+package handlers
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+
+	snapshotcache "github.com/e2b-dev/infra/packages/api/internal/cache/snapshots"
+	"github.com/e2b-dev/infra/packages/db/pkg/testutils"
+	redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis"
+)
+
+// newTestSnapshotCache builds a SnapshotCache backed by a real Postgres and Redis
+// test container, mirroring the wiring in NewSnapshotCache used by the API store.
+func newTestSnapshotCache(t *testing.T, db *testutils.Database) *snapshotcache.SnapshotCache {
+	t.Helper()
+
+	redisClient := redis_utils.SetupInstance(t)
+	sc := snapshotcache.NewSnapshotCache(db.SqlcClient, redisClient)
+	t.Cleanup(func() {
+		_ = sc.Close(t.Context())
+	})
+
+	return sc
+}
+
+// TestPauseHandleNotRunningSandbox_WrongTeamReturnsNotFound verifies the ENG-3544
+// behavior in the pause path: pausing a not-running sandbox whose snapshot belongs to
+// another team must return 404 (indistinguishable from a sandbox that does not exist),
+// rather than leaking that a paused snapshot exists.
+func TestPauseHandleNotRunningSandbox_WrongTeamReturnsNotFound(t *testing.T) {
+	t.Parallel()
+	db := testutils.SetupDatabase(t)
+	ctx := t.Context()
+
+	sc := newTestSnapshotCache(t, db)
+
+	ownerTeamID := testutils.CreateTestTeam(t, db)
+	baseTemplateID := testutils.CreateTestTemplate(t, db, ownerTeamID)
+
+	sandboxID := "sandbox-" + uuid.New().String()
+	snapshotTemplateID := "snapshot-template-" + uuid.New().String()
+	testutils.UpsertTestSnapshot(t, ctx, db, snapshotTemplateID, sandboxID, ownerTeamID, baseTemplateID)
+
+	otherTeamID := testutils.CreateTestTeam(t, db)
+
+	apiErr := pauseHandleNotRunningSandbox(ctx, sc, sandboxID, otherTeamID)
+	assert.Equal(t, http.StatusNotFound, apiErr.Code,
+		"another team's paused sandbox must be reported as not found")
+}
+
+// TestPauseHandleNotRunningSandbox_OwnerReturnsAlreadyPaused verifies that when the
+// owning team pauses a sandbox whose snapshot already exists, the handler reports the
+// already-paused conflict rather than a not-found.
+func TestPauseHandleNotRunningSandbox_OwnerReturnsAlreadyPaused(t *testing.T) {
+	t.Parallel()
+	db := testutils.SetupDatabase(t)
+	ctx := t.Context()
+
+	sc := newTestSnapshotCache(t, db)
+
+	teamID := testutils.CreateTestTeam(t, db)
+	baseTemplateID := testutils.CreateTestTemplate(t, db, teamID)
+
+	sandboxID := "sandbox-" + uuid.New().String()
+	snapshotTemplateID := "snapshot-template-" + uuid.New().String()
+	testutils.UpsertTestSnapshot(t, ctx, db, snapshotTemplateID, sandboxID, teamID, baseTemplateID)
+
+	apiErr := pauseHandleNotRunningSandbox(ctx, sc, sandboxID, teamID)
+	assert.Equal(t, http.StatusConflict, apiErr.Code,
+		"the owning team's already-paused sandbox should return a conflict")
+}

packages/api/internal/handlers/sandbox_snapshot_scope_test.go

@@ -0,0 +1,134 @@
+package handlers
+
+import (
+	"bytes"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/api/internal/api"
+	handlersmocks "github.com/e2b-dev/infra/packages/api/internal/handlers/mocks"
+	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
+	"github.com/e2b-dev/infra/packages/auth/pkg/auth"
+	authtypes "github.com/e2b-dev/infra/packages/auth/pkg/types"
+	authqueries "github.com/e2b-dev/infra/packages/db/pkg/auth/queries"
+	"github.com/e2b-dev/infra/packages/db/pkg/testutils"
+)
+
+// newCrossTeamSnapshotFixture creates a snapshot owned by one team and returns an
+// APIStore whose orchestrator is mocked (so the sandbox is reported not-running and
+// the handler falls through to the snapshot lookup) together with a *different*
+// requesting team. The returned sandboxID is a bare, dash-free lowercase ID so it
+// survives utils.ShortID parsing unchanged and matches the stored snapshot row.
+func newCrossTeamSnapshotFixture(t *testing.T) (store *APIStore, orch *handlersmocks.MockSandboxOrchestrator, requesterTeamID uuid.UUID, sandboxID string) {
+	t.Helper()
+
+	db := testutils.SetupDatabase(t)
+	ctx := t.Context()
+
+	sc := newTestSnapshotCache(t, db)
+	orch = handlersmocks.NewMockSandboxOrchestrator(t)
+
+	ownerTeamID := testutils.CreateTestTeam(t, db)
+	baseTemplateID := testutils.CreateTestTemplate(t, db, ownerTeamID)
+
+	sandboxID = strings.ReplaceAll(uuid.NewString(), "-", "")
+	snapshotTemplateID := "snapshot-template-" + uuid.New().String()
+	testutils.UpsertTestSnapshot(t, ctx, db, snapshotTemplateID, sandboxID, ownerTeamID, baseTemplateID)
+
+	requesterTeamID = testutils.CreateTestTeam(t, db)
+
+	store = &APIStore{
+		orchestrator:  orch,
+		snapshotCache: sc,
+	}
+
+	return store, orch, requesterTeamID, sandboxID
+}
+
+func newTeamGinContext(t *testing.T, teamID uuid.UUID, method, target string, body any) (*httptest.ResponseRecorder, *gin.Context) {
+	t.Helper()
+
+	rec := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(rec)
+
+	var raw []byte
+	if body != nil {
+		var err error
+		raw, err = json.Marshal(body)
+		require.NoError(t, err)
+	}
+
+	c.Request = httptest.NewRequestWithContext(t.Context(), method, target, bytes.NewReader(raw))
+	c.Request.Header.Set("Content-Type", "application/json")
+
+	auth.SetTeamInfoForTest(t, c, &authtypes.Team{
+		Team:   &authqueries.Team{ID: teamID, Slug: "team-" + teamID.String()[:8]},
+		Limits: &authtypes.TeamLimits{MaxLengthHours: 24},
+	})
+
+	return rec, c
+}
+
+// TestGetSandboxesSandboxID_CrossTeamSnapshotReturnsNotFound verifies ENG-3544 end to
+// end for GET: when the running sandbox is absent and the snapshot belongs to another
+// team, the requester receives 404 rather than the other team's snapshot.
+func TestGetSandboxesSandboxID_CrossTeamSnapshotReturnsNotFound(t *testing.T) {
+	t.Parallel()
+
+	store, orch, requesterTeamID, sandboxID := newCrossTeamSnapshotFixture(t)
+
+	orch.EXPECT().
+		GetSandbox(mock.Anything, requesterTeamID, sandboxID).
+		Return(sandbox.Sandbox{}, sandbox.ErrNotFound)
+
+	rec, c := newTeamGinContext(t, requesterTeamID, http.MethodGet, "/sandboxes/"+sandboxID, nil)
+	store.GetSandboxesSandboxID(c, sandboxID)
+
+	assert.Equal(t, http.StatusNotFound, rec.Code)
+}
+
+// TestPostSandboxesSandboxIDConnect_CrossTeamSnapshotReturnsNotFound verifies the same
+// for the connect path: KeepAliveFor reports the sandbox is not in the store, so the
+// handler falls through to the team-scoped snapshot lookup, which must not find another
+// team's snapshot.
+func TestPostSandboxesSandboxIDConnect_CrossTeamSnapshotReturnsNotFound(t *testing.T) {
+	t.Parallel()
+
+	store, orch, requesterTeamID, sandboxID := newCrossTeamSnapshotFixture(t)
+
+	orch.EXPECT().
+		KeepAliveFor(mock.Anything, requesterTeamID, sandboxID, mock.Anything, false).
+		Return(nil, &api.APIError{Code: http.StatusNotFound, ClientMsg: "not found", Err: sandbox.ErrNotFound})
+
+	rec, c := newTeamGinContext(t, requesterTeamID, http.MethodPost, "/sandboxes/"+sandboxID+"/connect", api.ConnectSandbox{Timeout: 30})
+	store.PostSandboxesSandboxIDConnect(c, sandboxID)
+
+	assert.Equal(t, http.StatusNotFound, rec.Code)
+}
+
+// TestPostSandboxesSandboxIDResume_CrossTeamSnapshotReturnsNotFound verifies the same
+// for the resume path: GetSandbox reports not-running, and the team-scoped snapshot
+// lookup must not resume another team's snapshot.
+func TestPostSandboxesSandboxIDResume_CrossTeamSnapshotReturnsNotFound(t *testing.T) {
+	t.Parallel()
+
+	store, orch, requesterTeamID, sandboxID := newCrossTeamSnapshotFixture(t)
+
+	orch.EXPECT().
+		GetSandbox(mock.Anything, requesterTeamID, sandboxID).
+		Return(sandbox.Sandbox{}, sandbox.ErrNotFound)
+
+	rec, c := newTeamGinContext(t, requesterTeamID, http.MethodPost, "/sandboxes/"+sandboxID+"/resume", api.ResumedSandbox{})
+	store.PostSandboxesSandboxIDResume(c, sandboxID)
+
+	assert.Equal(t, http.StatusNotFound, rec.Code)
+}

packages/db/pkg/tests/snapshots/snapshot_team_scope_test.go

@@ -0,0 +1,86 @@
+package snapshots
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/db/pkg/testutils"
+	"github.com/e2b-dev/infra/packages/db/queries"
+)
+
+// TestGetLastSnapshot_MatchingTeamReturnsSnapshot verifies that scoping the query
+// to the owning team returns the snapshot's latest build.
+func TestGetLastSnapshot_MatchingTeamReturnsSnapshot(t *testing.T) {
+	t.Parallel()
+	db := testutils.SetupDatabase(t)
+	ctx := t.Context()
+
+	teamID := testutils.CreateTestTeam(t, db)
+	baseTemplateID := testutils.CreateTestTemplate(t, db, teamID)
+
+	sandboxID := "sandbox-" + uuid.New().String()
+	snapshotTemplateID := "snapshot-template-" + uuid.New().String()
+	result := testutils.UpsertTestSnapshot(t, ctx, db, snapshotTemplateID, sandboxID, teamID, baseTemplateID)
+
+	snapshot, err := db.SqlcClient.GetLastSnapshot(ctx, queries.GetLastSnapshotParams{
+		SandboxID: sandboxID,
+		TeamID:    &teamID,
+	})
+	require.NoError(t, err)
+
+	assert.Equal(t, result.BuildID, snapshot.EnvBuild.ID,
+		"GetLastSnapshot scoped to the owning team should return the snapshot build")
+}
+
+// TestGetLastSnapshot_WrongTeamReturnsNotFound is the core security assertion for
+// ENG-3544: a team must not be able to read another team's snapshot. Scoping the
+// query to a non-owning team must yield no rows.
+func TestGetLastSnapshot_WrongTeamReturnsNotFound(t *testing.T) {
+	t.Parallel()
+	db := testutils.SetupDatabase(t)
+	ctx := t.Context()
+
+	ownerTeamID := testutils.CreateTestTeam(t, db)
+	baseTemplateID := testutils.CreateTestTemplate(t, db, ownerTeamID)
+
+	sandboxID := "sandbox-" + uuid.New().String()
+	snapshotTemplateID := "snapshot-template-" + uuid.New().String()
+	testutils.UpsertTestSnapshot(t, ctx, db, snapshotTemplateID, sandboxID, ownerTeamID, baseTemplateID)
+
+	otherTeamID := testutils.CreateTestTeam(t, db)
+
+	_, err := db.SqlcClient.GetLastSnapshot(ctx, queries.GetLastSnapshotParams{
+		SandboxID: sandboxID,
+		TeamID:    &otherTeamID,
+	})
+	require.ErrorIs(t, err, pgx.ErrNoRows,
+		"GetLastSnapshot scoped to a non-owning team must not return the snapshot")
+}
+
+// TestGetLastSnapshot_NilTeamReturnsSnapshot verifies the backwards-compatible path:
+// when no team is provided the query is not team-scoped and still returns the snapshot.
+func TestGetLastSnapshot_NilTeamReturnsSnapshot(t *testing.T) {
+	t.Parallel()
+	db := testutils.SetupDatabase(t)
+	ctx := t.Context()
+
+	teamID := testutils.CreateTestTeam(t, db)
+	baseTemplateID := testutils.CreateTestTemplate(t, db, teamID)
+
+	sandboxID := "sandbox-" + uuid.New().String()
+	snapshotTemplateID := "snapshot-template-" + uuid.New().String()
+	result := testutils.UpsertTestSnapshot(t, ctx, db, snapshotTemplateID, sandboxID, teamID, baseTemplateID)
+
+	snapshot, err := db.SqlcClient.GetLastSnapshot(ctx, queries.GetLastSnapshotParams{
+		SandboxID: sandboxID,
+		TeamID:    nil,
+	})
+	require.NoError(t, err)
+
+	assert.Equal(t, result.BuildID, snapshot.EnvBuild.ID,
+		"GetLastSnapshot without team scoping should still return the snapshot build")
+}