fix(multitenancy): pin the live-query websocket stream to the tenant team

openframe/docs/mysql-multitenancy-feature.md

@@ -72,6 +72,11 @@ role-authz grouping upstream; here it is a hard boundary regardless of token rol
 - **enroll_secrets** (`app_configs.go`) — `GetEnrollSecrets`/`ApplyEnrollSecrets` force `teamID =
   pinned`; `VerifyEnrollSecret` only accepts a secret whose `team_id = pinned` (agent boundary).
 - **live-query targets** (`targets.go`) — `HostIDsInTargets`/`CountHostsInTargets` scoped.
+- **live-query campaigns** (`campaigns.go`) — `DistributedQueryCampaign`/
+  `DistributedQueryCampaignTargetIDs` fenced via the campaign's query team (`EXISTS` against
+  `queries.team_id`; pinned creation re-homes ad-hoc query rows, so every campaign's query carries
+  its tenant's team). Upstream's only guard is `campaign.UserID`, which is useless on the shared
+  Fleet where every tenant operates as the same Admin user.
 - **host-assignments** (`policy_hosts`/`query_hosts`) — parent verified in team + foreign host ids
   dropped (pre-existing `OPENFRAME(host-assignments)` feature, extended here).
 - **teams** — `TeamLite`/`ListTeams` read fence.
@@ -90,6 +95,12 @@ role-authz grouping upstream; here it is a hard boundary regardless of token rol
   the osquery header pre-auth paths (`osquery_header_auth.go`); fail-closed on a team-less host.
 - **Enrollment** — after `VerifyEnrollSecret`, `osquery.go`/`orbit.go` pin from `secret.TeamID`
   (reject if the secret has no team). The host row is then created carrying that `team_id`.
+- **Live-query results websocket** (`endpoint_campaigns.go`) — the sockjs handler rebuilds its
+  context from `context.Background()`, discarding the middleware-pinned upgrade-request context,
+  so it **re-pins from `session.Request().Context()`** (read once — polling transports mutate the
+  session request). Fail closed: in shared mode an unpinned session is rejected. Without the
+  re-pin the whole campaign stream ran unfenced and the `live_query` activity was stamped
+  `team_id NULL`.
 
 Agents send **no tenant header** — tenant identity flows in via the enroll secret and thereafter via
 the host record (node key → host → team). This is by design and stronger than a header.

server/datastore/mysql/campaigns.go

@@ -62,8 +62,18 @@ func (ds *Datastore) DistributedQueryCampaign(ctx context.Context, id uint) (*fl
 	sql := `
 		SELECT * FROM distributed_query_campaigns WHERE id = ?
 	`
+	args := []interface{}{id}
+	// >>> OPENFRAME(mysql-multitenancy): a campaign's tenant is its query's team (pinned creation
+	// re-homes ad-hoc query rows to the tenant team) — fence by-id reads so one tenant cannot load
+	// (and then stream) another tenant's campaign on the shared DB. No-op when unpinned.
+	// — openframe/docs/mysql-multitenancy-feature.md
+	if teamID, ok := fleet.OpenframeTeamID(ctx); ok {
+		sql += ` AND EXISTS (SELECT 1 FROM queries q WHERE q.id = distributed_query_campaigns.query_id AND q.team_id = ?)`
+		args = append(args, teamID)
+	}
+	// <<< OPENFRAME(mysql-multitenancy)
 	campaign := &fleet.DistributedQueryCampaign{}
-	if err := sqlx.GetContext(ctx, ds.reader(ctx), campaign, sql, id); err != nil {
+	if err := sqlx.GetContext(ctx, ds.reader(ctx), campaign, sql, args...); err != nil {
 		return nil, ctxerr.Wrap(ctx, err, "selecting distributed query campaign")
 	}
 
@@ -106,9 +116,21 @@ func (ds *Datastore) DistributedQueryCampaignTargetIDs(ctx context.Context, id u
 	sqlStatement := `
 		SELECT * FROM distributed_query_campaign_targets WHERE distributed_query_campaign_id = ?
 	`
+	args := []interface{}{id}
+	// >>> OPENFRAME(mysql-multitenancy): same fence as DistributedQueryCampaign — targets resolve
+	// via their campaign's query team. No-op when unpinned.
+	if teamID, ok := fleet.OpenframeTeamID(ctx); ok {
+		sqlStatement += ` AND EXISTS (
+			SELECT 1 FROM distributed_query_campaigns dqc
+			JOIN queries q ON q.id = dqc.query_id
+			WHERE dqc.id = distributed_query_campaign_targets.distributed_query_campaign_id AND q.team_id = ?
+		)`
+		args = append(args, teamID)
+	}
+	// <<< OPENFRAME(mysql-multitenancy)
 	targets := []fleet.DistributedQueryCampaignTarget{}
 
-	if err := sqlx.SelectContext(ctx, ds.reader(ctx), &targets, sqlStatement, id); err != nil {
+	if err := sqlx.SelectContext(ctx, ds.reader(ctx), &targets, sqlStatement, args...); err != nil {
 		return nil, ctxerr.Wrap(ctx, err, "select distributed campaign target")
 	}
 

server/datastore/mysql/campaigns_openframe_test.go

@@ -0,0 +1,73 @@
+package mysql
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/fleetdm/fleet/v4/server/fleet"
+	"github.com/fleetdm/fleet/v4/server/test"
+	"github.com/stretchr/testify/require"
+)
+
+// TestOpenframeCampaignTeamFence verifies the campaign read fence on a shared DB: a campaign's
+// tenant is its query's team, so a pinned context must not load (and then stream) another
+// tenant's campaign by id. Unpinned contexts keep upstream behavior. Runs only under MYSQL_TEST=1.
+func TestOpenframeCampaignTeamFence(t *testing.T) {
+	ds := CreateMySQLDS(t)
+	ctx := context.Background()
+
+	user := test.NewUser(t, ds, "Fence", "fence@openframe.local", true)
+	teamA, err := ds.NewTeam(ctx, &fleet.Team{Name: "openframe-campaign-a"})
+	require.NoError(t, err)
+	teamB, err := ds.NewTeam(ctx, &fleet.Team{Name: "openframe-campaign-b"})
+	require.NoError(t, err)
+
+	queryA := test.NewQuery(t, ds, &teamA.ID, "campaign-fence-a", "SELECT 1", user.ID, true)
+	queryB := test.NewQuery(t, ds, &teamB.ID, "campaign-fence-b", "SELECT 1", user.ID, true)
+	campA := test.NewCampaign(t, ds, queryA.ID, fleet.QueryRunning, time.Now())
+	campB := test.NewCampaign(t, ds, queryB.ID, fleet.QueryRunning, time.Now())
+
+	_, err = ds.NewDistributedQueryCampaignTarget(ctx, &fleet.DistributedQueryCampaignTarget{
+		Type: fleet.TargetTeam, DistributedQueryCampaignID: campA.ID, TargetID: teamA.ID,
+	})
+	require.NoError(t, err)
+	_, err = ds.NewDistributedQueryCampaignTarget(ctx, &fleet.DistributedQueryCampaignTarget{
+		Type: fleet.TargetTeam, DistributedQueryCampaignID: campB.ID, TargetID: teamB.ID,
+	})
+	require.NoError(t, err)
+
+	ctxA := fleet.NewOpenframeTeamContext(ctx, teamA.ID)
+
+	// Pinned to A: own campaign loads, B's campaign is invisible.
+	got, err := ds.DistributedQueryCampaign(ctxA, campA.ID)
+	require.NoError(t, err)
+	require.Equal(t, campA.ID, got.ID)
+
+	_, err = ds.DistributedQueryCampaign(ctxA, campB.ID)
+	require.Error(t, err, "a foreign campaign must not load under another tenant's pin")
+	require.True(t, errors.Is(err, sql.ErrNoRows))
+
+	// Same fence for target ids: own targets load, foreign campaign's targets come back empty.
+	targetsA, err := ds.DistributedQueryCampaignTargetIDs(ctxA, campA.ID)
+	require.NoError(t, err)
+	require.Equal(t, []uint{teamA.ID}, targetsA.TeamIDs)
+
+	targetsB, err := ds.DistributedQueryCampaignTargetIDs(ctxA, campB.ID)
+	require.NoError(t, err)
+	require.Empty(t, targetsB.TeamIDs)
+	require.Empty(t, targetsB.HostIDs)
+	require.Empty(t, targetsB.LabelIDs)
+
+	// Unpinned: upstream behavior — both campaigns load.
+	for _, id := range []uint{campA.ID, campB.ID} {
+		got, err := ds.DistributedQueryCampaign(ctx, id)
+		require.NoError(t, err)
+		require.Equal(t, id, got.ID)
+	}
+	targetsB, err = ds.DistributedQueryCampaignTargetIDs(ctx, campB.ID)
+	require.NoError(t, err)
+	require.Equal(t, []uint{teamB.ID}, targetsB.TeamIDs)
+}

server/service/endpoint_campaigns.go

@@ -84,6 +84,27 @@ func makeStreamDistributedQueryCampaignResultsHandler(config config.ServerConfig
 
 			ctx := viewer.NewContext(context.Background(), *vc)
 
+			// >>> OPENFRAME(mysql-multitenancy): sockjs handlers rebuild their context from
+			// Background, discarding the upgrade request's context that WithOpenframeTenant pinned
+			// with the tenant team — leaving the whole campaign stream unfenced and the live_query
+			// activity stamped with a NULL team_id. Re-apply the pin from the upgrade request (read
+			// once: polling transports mutate the session request under a lock). The base ctx stays
+			// Background so the stream's cancellation semantics are unchanged. Fail closed in shared
+			// mode: the middleware 401s a headerless upgrade before the handler runs, so an unpinned
+			// session here means the handler was mounted outside the middleware.
+			// — openframe/docs/mysql-multitenancy-feature.md
+			if req := session.Request(); req != nil {
+				if teamID, ok := fleet.OpenframeTeamID(req.Context()); ok {
+					ctx = fleet.NewOpenframeTeamContext(ctx, teamID)
+				}
+			}
+			if _, ok := fleet.OpenframeTeamID(ctx); !ok && fleet.IsOpenframeSharedMode() {
+				logger.ErrorContext(ctx, "openframe shared mode: rejecting campaign stream without tenant pin")
+				conn.WriteJSONError("missing tenant") //nolint:errcheck
+				return
+			}
+			// <<< OPENFRAME(mysql-multitenancy)
+
 			msg, err := conn.ReadJSONMessage()
 			if err != nil {
 				logger.ErrorContext(ctx, "reading select_campaign JSON", "err", err)

server/service/endpoint_campaigns_openframe_test.go

@@ -0,0 +1,119 @@
+// OPENFRAME(mysql-multitenancy): tests for the live-query websocket tenant re-pin — the sockjs
+// handler rebuilds its context from Background, so it must copy the team pin from the upgrade
+// request's context (set by the tenant middleware) onto the stream context. The fail-closed
+// shared-mode guard reads the cached mode env and is not exercisable here (same limitation as
+// openframe_middleware_test.go); these tests cover the re-pin and the unpinned pass-through.
+package service
+
+import (
+	"context"
+	"errors"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/fleetdm/fleet/v4/server/config"
+	"github.com/fleetdm/fleet/v4/server/fleet"
+	"github.com/fleetdm/fleet/v4/server/mock"
+	"github.com/fleetdm/fleet/v4/server/ptr"
+	ws "github.com/fleetdm/fleet/v4/server/websocket"
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/require"
+)
+
+type observedPin struct {
+	teamID uint
+	ok     bool
+}
+
+// streamCampaignAndObservePin drives the websocket handler (optionally wrapped in the tenant
+// middleware body, as WithOpenframeTenant wraps it in shared mode) through auth + select_campaign
+// and returns the team pin observed inside the campaign datastore lookup.
+func streamCampaignAndObservePin(t *testing.T, withTenantMiddleware bool) observedPin {
+	t.Helper()
+
+	ds := new(mock.Store)
+	svc, _ := newTestService(t, ds, nil, nil)
+
+	// Auth plumbing for AuthViewer (token sent as the first websocket message).
+	ds.SessionByKeyFunc = func(ctx context.Context, key string) (*fleet.Session, error) {
+		return &fleet.Session{
+			CreateTimestamp: fleet.CreateTimestamp{CreatedAt: time.Now()},
+			ID:              1,
+			AccessedAt:      time.Now(),
+			UserID:          1,
+			Key:             key,
+		}, nil
+	}
+	ds.MarkSessionAccessedFunc = func(context.Context, *fleet.Session) error { return nil }
+	ds.UserByIDFunc = func(ctx context.Context, id uint) (*fleet.User, error) {
+		return &fleet.User{ID: id, GlobalRole: ptr.String(fleet.RoleAdmin)}, nil
+	}
+
+	// Capture the pin at the first datastore touch of the stream, then short-circuit the
+	// handler with a non-ErrNoRows error (ErrNoRows would make it poll for 5s).
+	pinCh := make(chan observedPin, 1)
+	ds.DistributedQueryCampaignFunc = func(ctx context.Context, id uint) (*fleet.DistributedQueryCampaign, error) {
+		teamID, ok := fleet.OpenframeTeamID(ctx)
+		select {
+		case pinCh <- observedPin{teamID: teamID, ok: ok}:
+		default:
+		}
+		return nil, errors.New("pin captured, stop the stream")
+	}
+
+	pathHandler := makeStreamDistributedQueryCampaignResultsHandler(
+		config.TestConfig().Server, svc, slog.New(slog.DiscardHandler))
+	handler := pathHandler("/api/{fleetversion:(?:latest)}/fleet/results/")
+	if withTenantMiddleware {
+		handler = openframeTenantHandler(&fakeTeamEnsurer{teamID: 42}, slog.New(slog.DiscardHandler), handler)
+	}
+	s := httptest.NewServer(handler)
+	defer s.Close()
+
+	u := "ws" + strings.TrimPrefix(s.URL, "http") + "/api/latest/fleet/results/websocket"
+	dialer := &websocket.Dialer{HandshakeTimeout: 10 * time.Second}
+	var header http.Header
+	if withTenantMiddleware {
+		header = http.Header{"X-Tenant-Id": []string{"3f1a9b2c-0000-4d5e-8f00-000000000001"}}
+	}
+	conn, _, err := dialer.Dial(u, header)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	require.NoError(t, conn.WriteJSON(ws.JSONMessage{
+		Type: "auth",
+		Data: map[string]interface{}{"token": "test-token"},
+	}))
+	require.NoError(t, conn.WriteJSON(ws.JSONMessage{
+		Type: "select_campaign",
+		Data: map[string]interface{}{"campaign_id": 1},
+	}))
+
+	select {
+	case pin := <-pinCh:
+		return pin
+	case <-time.After(10 * time.Second):
+		t.Fatal("the campaign datastore lookup was never reached")
+		return observedPin{}
+	}
+}
+
+// TestOpenframeCampaignStreamRePinsFromUpgradeRequest verifies that the tenant pin placed on the
+// websocket upgrade request by the tenant middleware survives into the campaign stream's context
+// (and from there into the datastore fences and the live_query activity team stamping).
+func TestOpenframeCampaignStreamRePinsFromUpgradeRequest(t *testing.T) {
+	pin := streamCampaignAndObservePin(t, true)
+	require.True(t, pin.ok, "stream ctx must carry the tenant pin from the upgrade request")
+	require.Equal(t, uint(42), pin.teamID)
+}
+
+// TestOpenframeCampaignStreamUnpinnedPassThrough verifies upstream behavior is preserved when no
+// pin is present on the upgrade request (flag off / single-tenant): no pin is invented.
+func TestOpenframeCampaignStreamUnpinnedPassThrough(t *testing.T) {
+	pin := streamCampaignAndObservePin(t, false)
+	require.False(t, pin.ok, "an unpinned upgrade request must leave the stream ctx unpinned")
+}