Skip to content

Commit 6b8bf29

Browse files
adababysadababys
committed
fix(api): scope snapshot lookups by team (ENG-3544)
GetLastSnapshot previously filtered only by sandbox_id, so the API handlers fetched the snapshot and then compared Snapshot.TeamID to the requesting team after the fact. Push the team check into the query with an optional team_id filter and drop the post-fetch ownership checks in the connect/get/pause/resume handlers, so a sandbox_id belonging to another team is indistinguishable from a missing one. The snapshot cache Get now takes an optional teamID, uses a team-scoped cache key, and Invalidate clears both the plain and team-scoped keys.
1 parent 17038b8 commit 6b8bf29

8 files changed

Lines changed: 48 additions & 55 deletions

File tree

‎packages/api/internal/cache/snapshots/snapshot_cache.go‎

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"fmt"
77
"time"
88

9+
"github.com/google/uuid"
910
"github.com/redis/go-redis/v9"
1011
"go.opentelemetry.io/otel"
1112

@@ -57,11 +58,24 @@ func NewSnapshotCache(db *sqlcdb.Client, redisClient redis.UniversalClient) *Sna
5758
}
5859

5960
// Get returns the last snapshot for a sandbox, using cache with DB fallback.
60-
func (c *SnapshotCache) Get(ctx context.Context, sandboxID string) (*SnapshotInfo, error) {
61+
// When teamID is provided, the query filters by team at the DB level and uses
62+
// a team-scoped cache key, preventing cross-team data access without needing
63+
// a separate method or post-fetch ownership check.
64+
func (c *SnapshotCache) Get(ctx context.Context, sandboxID string, teamID ...uuid.UUID) (*SnapshotInfo, error) {
6165
ctx, span := tracer.Start(ctx, "get last snapshot")
6266
defer span.End()
6367

64-
info, err := c.cache.GetOrSet(ctx, sandboxID, c.fetchFromDB)
68+
var tid *uuid.UUID
69+
cacheKey := sandboxID
70+
71+
if len(teamID) > 0 {
72+
tid = &teamID[0]
73+
cacheKey = sandboxID + ":" + tid.String()
74+
}
75+
76+
info, err := c.cache.GetOrSet(ctx, cacheKey, func(ctx context.Context, _ string) (*SnapshotInfo, error) {
77+
return c.fetchFromDB(ctx, sandboxID, tid)
78+
})
6579
if err != nil {
6680
return nil, err
6781
}
@@ -73,11 +87,14 @@ func (c *SnapshotCache) Get(ctx context.Context, sandboxID string) (*SnapshotInf
7387
return info, nil
7488
}
7589

76-
func (c *SnapshotCache) fetchFromDB(ctx context.Context, sandboxID string) (*SnapshotInfo, error) {
90+
func (c *SnapshotCache) fetchFromDB(ctx context.Context, sandboxID string, teamID *uuid.UUID) (*SnapshotInfo, error) {
7791
ctx, span := tracer.Start(ctx, "fetch last snapshot from DB")
7892
defer span.End()
7993

80-
row, err := c.db.GetLastSnapshot(ctx, sandboxID)
94+
row, err := c.db.GetLastSnapshot(ctx, queries.GetLastSnapshotParams{
95+
SandboxID: sandboxID,
96+
TeamID: teamID,
97+
})
8198
if err != nil {
8299
if dberrors.IsNotFoundError(err) {
83100
return errNotFoundSentinel, nil
@@ -95,8 +112,10 @@ func (c *SnapshotCache) fetchFromDB(ctx context.Context, sandboxID string) (*Sna
95112
}
96113

97114
// Invalidate removes the cached snapshot for a sandbox.
115+
// It deletes both the plain sandboxID key and any team-scoped keys.
98116
func (c *SnapshotCache) Invalidate(ctx context.Context, sandboxID string) {
99117
c.cache.Delete(ctx, sandboxID)
118+
c.cache.DeleteByPrefix(ctx, sandboxID+":")
100119
}
101120

102121
func (c *SnapshotCache) Close(ctx context.Context) error {

‎packages/api/internal/handlers/sandbox_connect.go‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -115,8 +115,7 @@ func (a *APIStore) PostSandboxesSandboxIDConnect(c *gin.Context, sandboxID api.S
115115
continue
116116
}
117117

118-
// TODO: ENG-3544 scope GetLastSnapshot query by teamID to avoid post-fetch ownership check.
119-
lastSnapshot, err := a.snapshotCache.Get(ctx, sandboxID)
118+
lastSnapshot, err := a.snapshotCache.Get(ctx, sandboxID, teamID)
120119
if err != nil {
121120
if errors.Is(err, snapshotcache.ErrSnapshotNotFound) {
122121
logger.L().Debug(ctx, "Snapshot not found", logger.WithSandboxID(sandboxID))
@@ -131,13 +130,6 @@ func (a *APIStore) PostSandboxesSandboxIDConnect(c *gin.Context, sandboxID api.S
131130
return
132131
}
133132

134-
if lastSnapshot.Snapshot.TeamID != teamID {
135-
telemetry.ReportError(ctx, fmt.Sprintf("snapshot for sandbox '%s' doesn't belong to team '%s'", sandboxID, teamID.String()), nil)
136-
a.sendAPIStoreError(c, http.StatusNotFound, utils.SandboxNotFoundMsg(sandboxID))
137-
138-
return
139-
}
140-
141133
// A paused filesystem-only snapshot resumes by cold-booting (reboot) from its
142134
// rootfs; the orchestrator selects reboot-vs-memory-resume from the snapshot
143135
// metadata, so the generic resume path below handles it. In-memory state was
@@ -156,7 +148,7 @@ func (a *APIStore) PostSandboxesSandboxIDConnect(c *gin.Context, sandboxID api.S
156148
sandboxID,
157149
timeout,
158150
teamInfo,
159-
a.buildResumeSandboxData(sandboxID, nil),
151+
a.buildResumeSandboxData(sandboxID, nil, teamID),
160152
&c.Request.Header,
161153
true,
162154
nil, // mcp

‎packages/api/internal/handlers/sandbox_get.go‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -175,8 +175,7 @@ func (a *APIStore) GetSandboxesSandboxID(c *gin.Context, id string) {
175175
}
176176

177177
// If sandbox not found try to get the latest snapshot
178-
// TODO: ENG-3544 scope GetLastSnapshot query by teamID to avoid post-fetch ownership check.
179-
lastSnapshot, err := a.snapshotCache.Get(ctx, sandboxId)
178+
lastSnapshot, err := a.snapshotCache.Get(ctx, sandboxId, team.ID)
180179
if err != nil {
181180
if errors.Is(err, snapshotcache.ErrSnapshotNotFound) {
182181
telemetry.ReportError(ctx, "snapshot not found", err, telemetry.WithSandboxID(sandboxId))
@@ -191,13 +190,6 @@ func (a *APIStore) GetSandboxesSandboxID(c *gin.Context, id string) {
191190
return
192191
}
193192

194-
if lastSnapshot.Snapshot.TeamID != team.ID {
195-
telemetry.ReportError(ctx, fmt.Sprintf("snapshot for sandbox '%s' doesn't belong to team '%s'", sandboxId, team.ID.String()), nil)
196-
a.sendAPIStoreError(c, http.StatusNotFound, utils.SandboxNotFoundMsg(id))
197-
198-
return
199-
}
200-
201193
memoryMB := int32(lastSnapshot.EnvBuild.RamMb)
202194
cpuCount := int32(lastSnapshot.EnvBuild.Vcpu)
203195

‎packages/api/internal/handlers/sandbox_pause.go‎

Lines changed: 1 addition & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -94,18 +94,8 @@ func (a *APIStore) PostSandboxesSandboxIDPause(c *gin.Context, sandboxID api.San
9494
}
9595

9696
func pauseHandleNotRunningSandbox(ctx context.Context, cache *snapshotcache.SnapshotCache, sandboxID string, teamID uuid.UUID) api.APIError {
97-
// TODO: ENG-3544 scope GetLastSnapshot query by teamID to avoid post-fetch ownership check.
98-
snap, err := cache.Get(ctx, sandboxID)
97+
_, err := cache.Get(ctx, sandboxID, teamID)
9998
if err == nil {
100-
if snap.Snapshot.TeamID != teamID {
101-
logger.L().Debug(ctx, "Snapshot team mismatch on pause", logger.WithSandboxID(sandboxID), logger.WithTeamID(teamID.String()))
102-
103-
return api.APIError{
104-
Code: http.StatusNotFound,
105-
ClientMsg: utils.SandboxNotFoundMsg(sandboxID),
106-
}
107-
}
108-
10999
logger.L().Warn(ctx, "Sandbox is already paused", logger.WithSandboxID(sandboxID))
110100

111101
return api.APIError{

‎packages/api/internal/handlers/sandbox_resume.go‎

Lines changed: 5 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"time"
99

1010
"github.com/gin-gonic/gin"
11+
"github.com/google/uuid"
1112
"go.opentelemetry.io/otel/trace"
1213
"go.uber.org/zap"
1314

@@ -120,8 +121,7 @@ func (a *APIStore) PostSandboxesSandboxIDResume(c *gin.Context, sandboxID api.Sa
120121
}
121122
}
122123

123-
// TODO: ENG-3544 scope GetLastSnapshot query by teamID to avoid post-fetch ownership check.
124-
lastSnapshot, err := a.snapshotCache.Get(ctx, sandboxID)
124+
lastSnapshot, err := a.snapshotCache.Get(ctx, sandboxID, teamID)
125125
if err != nil {
126126
if errors.Is(err, snapshotcache.ErrSnapshotNotFound) {
127127
logger.L().Debug(ctx, "Snapshot not found", logger.WithSandboxID(sandboxID))
@@ -139,13 +139,6 @@ func (a *APIStore) PostSandboxesSandboxIDResume(c *gin.Context, sandboxID api.Sa
139139
return
140140
}
141141

142-
if lastSnapshot.Snapshot.TeamID != teamID {
143-
telemetry.ReportError(ctx, fmt.Sprintf("snapshot for sandbox '%s' doesn't belong to team '%s'", sandboxID, teamID.String()), nil)
144-
a.sendAPIStoreError(c, http.StatusNotFound, utils.SandboxNotFoundMsg(sandboxID))
145-
146-
return
147-
}
148-
149142
sbxlogger.E(&sbxlogger.SandboxMetadata{
150143
SandboxID: sandboxID,
151144
TemplateID: lastSnapshot.Snapshot.EnvID,
@@ -157,7 +150,7 @@ func (a *APIStore) PostSandboxesSandboxIDResume(c *gin.Context, sandboxID api.Sa
157150
sandboxID,
158151
timeout,
159152
teamInfo,
160-
a.buildResumeSandboxData(sandboxID, body.AutoPause),
153+
a.buildResumeSandboxData(sandboxID, body.AutoPause, teamID),
161154
&c.Request.Header,
162155
true,
163156
nil, // mcp
@@ -189,9 +182,9 @@ func convertDatabaseMountsToOrchestratorMounts(volumes []*types.SandboxVolumeMou
189182
// buildResumeSandboxData returns a SandboxDataFetcher that fetches snapshot data
190183
// from the cache and builds SandboxMetadata for resume operations.
191184
// The returned callback is called inside the sandbox lock to prevent race conditions.
192-
func (a *APIStore) buildResumeSandboxData(sandboxID string, autoPauseOverride *bool) orchestrator.SandboxDataFetcher {
185+
func (a *APIStore) buildResumeSandboxData(sandboxID string, autoPauseOverride *bool, teamID ...uuid.UUID) orchestrator.SandboxDataFetcher {
193186
return func(ctx context.Context) (orchestrator.SandboxMetadata, *api.APIError) {
194-
lastSnapshot, err := a.snapshotCache.Get(ctx, sandboxID)
187+
lastSnapshot, err := a.snapshotCache.Get(ctx, sandboxID, teamID...)
195188
if err != nil {
196189
return orchestrator.SandboxMetadata{}, &api.APIError{
197190
Code: http.StatusInternalServerError,

‎packages/db/pkg/tests/snapshots/snapshot_latest_assignment_test.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ func TestGetLastSnapshot_ReturnsLatestAssignment(t *testing.T) {
4343
require.NotEqual(t, build1ID, build2ID, "Each upsert should create a new build")
4444

4545
// Execute GetLastSnapshot - should return the latest build (build2)
46-
snapshot, err := db.SqlcClient.GetLastSnapshot(ctx, sandboxID)
46+
snapshot, err := db.SqlcClient.GetLastSnapshot(ctx, queries.GetLastSnapshotParams{SandboxID: sandboxID})
4747
require.NoError(t, err)
4848

4949
assert.Equal(t, build2ID, snapshot.EnvBuild.ID,
@@ -74,7 +74,7 @@ func TestGetLastSnapshot_OnlyReturnsSuccessBuilds(t *testing.T) {
7474
testutils.UpsertTestSnapshotWithStatus(t, ctx, db, snapshotTemplateID, sandboxID, teamID, baseTemplateID, types.BuildStatusSnapshotting)
7575

7676
// GetLastSnapshot should return the success build, not the snapshotting one
77-
snapshot, err := db.SqlcClient.GetLastSnapshot(ctx, sandboxID)
77+
snapshot, err := db.SqlcClient.GetLastSnapshot(ctx, queries.GetLastSnapshotParams{SandboxID: sandboxID})
7878
require.NoError(t, err)
7979

8080
assert.Equal(t, successBuildID, snapshot.EnvBuild.ID,
@@ -149,7 +149,7 @@ func TestGetLastSnapshot_BuildSharedWithOtherTemplate(t *testing.T) {
149149

150150
// GetLastSnapshot should still return build2 (latest for THIS template),
151151
// not build1 even though build1 has a newer assignment to another template
152-
snapshot, err := db.SqlcClient.GetLastSnapshot(ctx, sandboxID)
152+
snapshot, err := db.SqlcClient.GetLastSnapshot(ctx, queries.GetLastSnapshotParams{SandboxID: sandboxID})
153153
require.NoError(t, err)
154154

155155
assert.Equal(t, build2ID, snapshot.EnvBuild.ID,
@@ -183,7 +183,7 @@ func TestGetLastSnapshot_IgnoresNonDefaultTags(t *testing.T) {
183183
testutils.CreateTestBuildAssignment(t, ctx, db, result1.TemplateID, otherBuildID, "v1")
184184

185185
// GetLastSnapshot should return the default-tagged build, not the v1-tagged one
186-
snapshot, err := db.SqlcClient.GetLastSnapshot(ctx, sandboxID)
186+
snapshot, err := db.SqlcClient.GetLastSnapshot(ctx, queries.GetLastSnapshotParams{SandboxID: sandboxID})
187187
require.NoError(t, err)
188188

189189
assert.Equal(t, defaultBuildID, snapshot.EnvBuild.ID,
@@ -219,7 +219,7 @@ func TestGetLastSnapshot_AssignmentOrderDifferentFromBuildOrder(t *testing.T) {
219219
testutils.CreateSnapshotRecord(t, ctx, db, snapshotTemplateID, sandboxID, teamID, baseTemplateID)
220220

221221
// GetLastSnapshot should return build1 (latest assignment), not build2 (latest build)
222-
snapshot, err := db.SqlcClient.GetLastSnapshot(ctx, sandboxID)
222+
snapshot, err := db.SqlcClient.GetLastSnapshot(ctx, queries.GetLastSnapshotParams{SandboxID: sandboxID})
223223
require.NoError(t, err)
224224

225225
assert.Equal(t, build1ID, snapshot.EnvBuild.ID,

‎packages/db/queries/get_last_snapshot.sql.go‎

Lines changed: 10 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/db/queries/snapshots/get_last_snapshot.sql‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,4 +18,4 @@ LEFT JOIN LATERAL (
1818
FROM "public"."env_aliases"
1919
WHERE env_id = s.base_env_id
2020
) ea ON TRUE
21-
WHERE s.sandbox_id = $1;
21+
WHERE s.sandbox_id = $1 AND (sqlc.narg(team_id)::uuid IS NULL OR s.team_id = sqlc.narg(team_id)::uuid);

0 commit comments

Comments
 (0)