fleet/server/service/service_campaign_test.go

package service

import (
	"context"
	"crypto/tls"
	"database/sql"
	"log/slog"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/WatchBeam/clock"
	"github.com/fleetdm/fleet/v4/server/config"
	"github.com/fleetdm/fleet/v4/server/contexts/ctxdb"
	"github.com/fleetdm/fleet/v4/server/contexts/viewer"
	"github.com/fleetdm/fleet/v4/server/datastore/mysql"
	"github.com/fleetdm/fleet/v4/server/datastore/mysql/mysqltest"
	"github.com/fleetdm/fleet/v4/server/datastore/redis/redistest"
	"github.com/fleetdm/fleet/v4/server/fleet"
	"github.com/fleetdm/fleet/v4/server/live_query/live_query_mock"
	"github.com/fleetdm/fleet/v4/server/mock"
	"github.com/fleetdm/fleet/v4/server/ptr"
	"github.com/fleetdm/fleet/v4/server/pubsub"
	ws "github.com/fleetdm/fleet/v4/server/websocket"
	"github.com/gorilla/websocket"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestStreamCampaignResultsClosesReditOnWSClose(t *testing.T) {
	t.Skip("Seems to be a bit problematic in CI")

	store := pubsub.NewRedisQueryResults(
		redistest.SetupRedis(t, "zz", false, false, false),
		false,
		slog.New(slog.DiscardHandler),
	)

	mockClock := clock.NewMockClock()
	ds := new(mock.Store)
	lq := live_query_mock.New(t)
	svc, ctx := newTestServiceWithClock(t, ds, store, lq, mockClock)

	campaign := &fleet.DistributedQueryCampaign{ID: 42}

	ds.LabelQueriesForHostFunc = func(ctx context.Context, host *fleet.Host) (map[string]string, error) {
		return map[string]string{}, nil
	}
	ds.AppConfigFunc = func(ctx context.Context) (*fleet.AppConfig, error) {
		return &fleet.AppConfig{}, nil
	}
	ds.NewQueryFunc = func(ctx context.Context, query *fleet.Query, opts ...fleet.OptionalArg) (*fleet.Query, error) {
		return query, nil
	}
	ds.NewDistributedQueryCampaignFunc = func(ctx context.Context, camp *fleet.DistributedQueryCampaign) (*fleet.DistributedQueryCampaign, error) {
		return camp, nil
	}
	ds.NewDistributedQueryCampaignTargetFunc = func(ctx context.Context, target *fleet.DistributedQueryCampaignTarget) (*fleet.DistributedQueryCampaignTarget, error) {
		return target, nil
	}
	ds.HostIDsInTargetsFunc = func(ctx context.Context, filter fleet.TeamFilter, targets fleet.HostTargets) ([]uint, error) {
		return []uint{1}, nil
	}
	ds.CountHostsInTargetsFunc = func(ctx context.Context, filter fleet.TeamFilter, targets fleet.HostTargets, now time.Time) (fleet.TargetMetrics, error) {
		return fleet.TargetMetrics{TotalHosts: 1}, nil
	}
	ds.SessionByKeyFunc = func(ctx context.Context, key string) (*fleet.Session, error) {
		return &fleet.Session{
			CreateTimestamp: fleet.CreateTimestamp{CreatedAt: time.Now()},
			ID:              42,
			AccessedAt:      time.Now(),
			UserID:          999,
			Key:             "asd",
		}, nil
	}

	host := &fleet.Host{ID: 1, Platform: "windows"}

	lq.On("QueriesForHost", uint(1)).Return(
		map[string]string{
			strconv.Itoa(int(campaign.ID)): "select * from time",
		},
		nil,
	)
	lq.On("QueryCompletedByHost", strconv.Itoa(int(campaign.ID)), host.ID).Return(nil)
	lq.On("RunQuery", "0", "select year, month, day, hour, minutes, seconds from time", []uint{1}).Return(nil)
	viewerCtx := viewer.NewContext(ctx, viewer.Viewer{
		User: &fleet.User{
			ID:         0,
			GlobalRole: ptr.String(fleet.RoleAdmin),
		},
	})
	q := "select year, month, day, hour, minutes, seconds from time"
	_, err := svc.NewDistributedQueryCampaign(viewerCtx, q, nil, fleet.HostTargets{HostIDs: []uint{2}, LabelIDs: []uint{1}})
	require.NoError(t, err)

	pathHandler := makeStreamDistributedQueryCampaignResultsHandler(config.TestConfig().Server, svc, slog.New(slog.DiscardHandler))
	s := httptest.NewServer(pathHandler("/api/latest/fleet/results/"))
	defer s.Close()
	// Convert http://127.0.0.1 to ws://127.0.0.1
	u := "ws" + strings.TrimPrefix(s.URL, "http") + "/api/latest/fleet/results/websocket"

	// Connect to the server
	dialer := &websocket.Dialer{
		Proxy:            http.ProxyFromEnvironment,
		HandshakeTimeout: 45 * time.Second,
		TLSClientConfig:  &tls.Config{InsecureSkipVerify: true},
	}

	conn, _, err := dialer.Dial(u, nil)
	if err != nil {
		t.Fatalf("%v", err)
	}
	defer conn.Close()

	err = conn.WriteJSON(ws.JSONMessage{
		Type: "auth",
		Data: map[string]interface{}{"token": "asd"},
	})
	require.NoError(t, err)

	err = conn.WriteJSON(ws.JSONMessage{
		Type: "select_campaign",
		Data: map[string]interface{}{"campaign_id": campaign.ID},
	})
	require.NoError(t, err)

	ds.MarkSessionAccessedFunc = func(context.Context, *fleet.Session) error {
		return nil
	}
	ds.UserByIDFunc = func(ctx context.Context, id uint) (*fleet.User, error) {
		return &fleet.User{GlobalRole: ptr.String(fleet.RoleAdmin)}, nil
	}
	ds.DistributedQueryCampaignFunc = func(ctx context.Context, id uint) (*fleet.DistributedQueryCampaign, error) {
		return campaign, nil
	}
	ds.SaveDistributedQueryCampaignFunc = func(ctx context.Context, camp *fleet.DistributedQueryCampaign) error {
		return nil
	}
	ds.DistributedQueryCampaignTargetIDsFunc = func(ctx context.Context, id uint) (targets *fleet.HostTargets, err error) {
		return &fleet.HostTargets{HostIDs: []uint{1}}, nil
	}
	ds.QueryFunc = func(ctx context.Context, id uint) (*fleet.Query, error) {
		return &fleet.Query{}, nil
	}

	/*****************************************************************************************/
	/* THE ACTUAL TEST BEGINS HERE                                                           */
	/*****************************************************************************************/
	prevActiveConn := 0
	for prevActiveConn < 3 {
		time.Sleep(2 * time.Second)

		for _, s := range store.Pool().Stats() {
			prevActiveConn = s.ActiveCount
		}
	}

	conn.Close()
	time.Sleep(10 * time.Second)

	newActiveConn := prevActiveConn
	for _, s := range store.Pool().Stats() {
		newActiveConn = s.ActiveCount
	}
	require.Equal(t, prevActiveConn-1, newActiveConn)
}

// TestStreamCampaignResultsDoesNotLeakCampaignExistence ensures that
// selecting a nonexistent campaign and selecting a campaign owned by another
// user produce the exact same websocket error message. Distinct messages
// here would let an authenticated user enumerate other users' live query
// campaigns by observing which response they get.
func TestStreamCampaignResultsDoesNotLeakCampaignExistence(t *testing.T) {
	ds := new(mock.Store)
	svc, _ := newTestService(t, ds, nil, nil)

	ds.AppConfigFunc = func(ctx context.Context) (*fleet.AppConfig, error) {
		return &fleet.AppConfig{}, nil
	}
	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:             "observer-token",
		}, nil
	}
	ds.UserByIDFunc = func(ctx context.Context, id uint) (*fleet.User, error) {
		return &fleet.User{ID: 1, GlobalRole: new(fleet.RoleObserver)}, nil
	}
	ds.MarkSessionAccessedFunc = func(context.Context, *fleet.Session) error {
		return nil
	}

	othersCampaignID := uint(99)
	ds.DistributedQueryCampaignFunc = func(ctx context.Context, id uint) (*fleet.DistributedQueryCampaign, error) {
		if id == othersCampaignID {
			// Owned by a different user than the one authenticating below.
			return &fleet.DistributedQueryCampaign{ID: othersCampaignID, UserID: 2}, nil
		}
		return nil, sql.ErrNoRows
	}

	pathHandler := makeStreamDistributedQueryCampaignResultsHandler(config.TestConfig().Server, svc, slog.New(slog.DiscardHandler))
	s := httptest.NewServer(pathHandler("/api/{fleetversion:(?:v1|2022-04)}/fleet/results/"))
	defer s.Close()
	u := "ws" + strings.TrimPrefix(s.URL, "http") + "/api/2022-04/fleet/results/websocket"

	dialer := &websocket.Dialer{
		Proxy:            http.ProxyFromEnvironment,
		HandshakeTimeout: 45 * time.Second,
		TLSClientConfig:  &tls.Config{InsecureSkipVerify: true},
	}

	selectCampaign := func(t *testing.T, campaignID uint) string {
		conn, _, err := dialer.Dial(u, nil)
		require.NoError(t, err)
		defer conn.Close()

		require.NoError(t, conn.WriteJSON(ws.JSONMessage{
			Type: "auth",
			Data: map[string]any{"token": "observer-token"},
		}))
		require.NoError(t, conn.WriteJSON(ws.JSONMessage{
			Type: "select_campaign",
			Data: map[string]any{"campaign_id": campaignID},
		}))

		_, msg, err := conn.ReadMessage()
		require.NoError(t, err)
		return string(msg)
	}

	nonExistentMsg := selectCampaign(t, 999999) // waits out the replica-lag retry/timeout
	otherUsersMsg := selectCampaign(t, othersCampaignID)

	assert.Equal(t, nonExistentMsg, otherUsersMsg,
		"a nonexistent campaign and one owned by another user must be indistinguishable")
	assert.Contains(t, nonExistentMsg, "forbidden")
}

func testUpdateStats(t *testing.T, ds *mysql.Datastore, usingReplica bool) {
	t.Cleanup(
		func() {
			overwriteLastExecuted = false
		},
	)
	s, ctx := newTestService(t, ds, nil, nil)
	svc := s.(validationMiddleware).Service.(*Service)

	tracker := statsTracker{}
	// NOOP cases
	svc.updateStats(ctx, 0, svc.logger, nil, false)
	svc.updateStats(ctx, 0, svc.logger, &tracker, false)

	// More NOOP cases
	tracker.saveStats = true
	svc.updateStats(ctx, 0, svc.logger, nil, false)
	assert.True(t, tracker.saveStats)
	svc.updateStats(ctx, 0, svc.logger, nil, true)
	assert.True(t, tracker.saveStats)

	// Populate a batch of data
	hostIDs := []uint{}
	queryID := uint(1)
	myHostID := uint(10000)
	myWallTime := uint64(5)
	myUserTime := uint64(6)
	mySystemTime := uint64(7)
	myMemory := uint64(8)
	myOutputSize := uint64(9)
	tracker.stats = append(
		tracker.stats, statsToSave{
			hostID: myHostID,
			Stats: &fleet.Stats{
				WallTimeMs: myWallTime,
				UserTime:   myUserTime,
				SystemTime: mySystemTime,
				Memory:     myMemory,
			},
			outputSize: myOutputSize,
		},
	)
	hostIDs = append(hostIDs, myHostID)

	for i := uint(1); i < statsBatchSize; i++ {
		tracker.stats = append(
			tracker.stats, statsToSave{
				hostID: i,
				Stats: &fleet.Stats{
					WallTimeMs: rand.Uint64(),
					UserTime:   rand.Uint64(),
					SystemTime: rand.Uint64(),
					Memory:     rand.Uint64(),
				},
				outputSize: rand.Uint64(),
			},
		)
		hostIDs = append(hostIDs, i)
	}
	tracker.saveStats = true
	// We overwrite the last executed time to ensure that these stats have a different timestamp than later stats
	overwriteLastExecuted = true
	overwriteLastExecutedTime = time.Now().Add(-2 * time.Second).Round(time.Second)
	svc.updateStats(ctx, queryID, svc.logger, &tracker, false)
	assert.True(t, tracker.saveStats)
	assert.Equal(t, 0, len(tracker.stats))
	assert.True(t, tracker.aggregationNeeded)
	assert.Equal(t, overwriteLastExecutedTime, tracker.lastStatsEntry.LastExecuted)

	// Aggregate stats
	svc.updateStats(ctx, queryID, svc.logger, &tracker, true)
	overwriteLastExecuted = false

	// Check that aggregated stats were created. Since we read aggregated stats from the replica, we may need to wait for it to catch up.
	var err error
	var aggStats fleet.AggregatedStats
	done := make(chan struct{}, 1)
	go func() {
		for {
			aggStats, err = mysqltest.GetAggregatedStats(ctx, svc.ds.(*mysql.Datastore), fleet.AggregatedStatsTypeScheduledQuery, queryID)
			if usingReplica && err != nil {
				time.Sleep(30 * time.Millisecond)
			} else {
				done <- struct{}{}
				return
			}
		}
	}()
	select {
	case <-time.After(2 * time.Second):
		// We fail the test. Gather information for debug.
		// Grab stats from primary DB (master)
		lastErr := ""
		if err != nil {
			lastErr = err.Error()
		}
		aggStats, err = mysqltest.GetAggregatedStats(
			ctxdb.RequirePrimary(ctx, true), svc.ds.(*mysql.Datastore), fleet.AggregatedStatsTypeScheduledQuery, queryID,
		)
		if err != nil {
			t.Logf("Error getting aggregated stats from primary DB: %s", err.Error())
		} else {
			t.Logf("Aggregated stats from primary DB: totalExecutions=%f %#v", *aggStats.TotalExecutions, aggStats)
		}
		replicaStatus, err := ds.ReplicaStatus(ctx)
		assert.NoError(t, err)
		t.Logf("Replica status: %v", replicaStatus)
		t.Fatalf("Timeout waiting for aggregated stats. Last error: %s", lastErr)
	case <-done:
		// Continue
	}
	require.NoError(t, err)
	assert.Equal(t, statsBatchSize, int(*aggStats.TotalExecutions))
	// Sanity checks. Complete testing done in aggregated_stats_test.go
	assert.True(t, *aggStats.SystemTimeP50 > 0)
	assert.True(t, *aggStats.SystemTimeP95 > 0)
	assert.True(t, *aggStats.UserTimeP50 > 0)
	assert.True(t, *aggStats.UserTimeP95 > 0)

	// Get the stats from DB and make sure they match
	currentStats, err := svc.ds.GetLiveQueryStats(ctx, queryID, hostIDs)
	assert.NoError(t, err)
	assert.Equal(t, statsBatchSize, len(currentStats))
	currentStats, err = svc.ds.GetLiveQueryStats(ctx, queryID, []uint{myHostID})
	assert.NoError(t, err)
	require.Equal(t, 1, len(currentStats))
	myStat := currentStats[0]
	assert.Equal(t, myHostID, myStat.HostID)
	assert.Equal(t, uint64(1), myStat.Executions)
	assert.Equal(t, myWallTime, myStat.WallTime)
	assert.Equal(t, myUserTime, myStat.UserTime)
	assert.Equal(t, mySystemTime, myStat.SystemTime)
	assert.Equal(t, myMemory, myStat.AverageMemory)
	assert.Equal(t, myOutputSize, myStat.OutputSize)

	// Write new stats (update) for the same query/hosts
	myNewWallTime := uint64(15)
	myNewUserTime := uint64(16)
	myNewSystemTime := uint64(17)
	myNewMemory := uint64(18)
	myNewOutputSize := uint64(19)
	tracker.stats = append(
		tracker.stats, statsToSave{
			hostID: myHostID,
			Stats: &fleet.Stats{
				WallTimeMs: myNewWallTime,
				UserTime:   myNewUserTime,
				SystemTime: myNewSystemTime,
				Memory:     myNewMemory,
			},
			outputSize: myNewOutputSize,
		},
	)

	for i := uint(1); i < statsBatchSize; i++ {
		tracker.stats = append(
			tracker.stats, statsToSave{
				hostID: i,
				Stats: &fleet.Stats{
					WallTimeMs: rand.Uint64(),
					UserTime:   rand.Uint64() % 100, // Keep these values small to ensure the update will be noticeable
					SystemTime: rand.Uint64() % 100, // Keep these values small to ensure the update will be noticeable
					Memory:     rand.Uint64(),
				},
				outputSize: rand.Uint64(),
			},
		)
	}
	tracker.saveStats = true
	svc.updateStats(ctx, queryID, svc.logger, &tracker, true)
	assert.True(t, tracker.saveStats)
	assert.Equal(t, 0, len(tracker.stats))
	assert.False(t, tracker.aggregationNeeded)

	// Check that aggregated stats were updated. Since we read aggregated stats from the replica, we may need to wait for it to catch up.
	var newAggStats fleet.AggregatedStats
	done = make(chan struct{}, 1)
	go func() {
		for {
			newAggStats, err = mysqltest.GetAggregatedStats(ctx, svc.ds.(*mysql.Datastore), fleet.AggregatedStatsTypeScheduledQuery, queryID)
			if usingReplica && (*aggStats.SystemTimeP50 == *newAggStats.SystemTimeP50 ||
				*aggStats.SystemTimeP95 == *newAggStats.SystemTimeP95 ||
				*aggStats.UserTimeP50 == *newAggStats.UserTimeP50 ||
				*aggStats.UserTimeP95 == *newAggStats.UserTimeP95) {
				time.Sleep(30 * time.Millisecond)
			} else {
				done <- struct{}{}
				return
			}
		}
	}()
	select {
	case <-time.After(2 * time.Second):
		t.Fatal("Timeout waiting for aggregated stats")
	case <-done:
		// Continue
	}

	require.NoError(t, err)
	assert.Equal(t, statsBatchSize*2, int(*newAggStats.TotalExecutions))
	// Sanity checks. Complete testing done in aggregated_stats_test.go
	assert.True(t, *newAggStats.SystemTimeP50 > 0)
	assert.True(t, *newAggStats.SystemTimeP95 > 0)
	assert.True(t, *newAggStats.UserTimeP50 > 0)
	assert.True(t, *newAggStats.UserTimeP95 > 0)
	assert.NotEqual(t, *aggStats.SystemTimeP50, *newAggStats.SystemTimeP50)
	assert.NotEqual(t, *aggStats.SystemTimeP95, *newAggStats.SystemTimeP95)
	assert.NotEqual(t, *aggStats.UserTimeP50, *newAggStats.UserTimeP50)
	assert.NotEqual(t, *aggStats.UserTimeP95, *newAggStats.UserTimeP95)

	// Check that stats were updated
	currentStats, err = svc.ds.GetLiveQueryStats(ctx, queryID, []uint{myHostID})
	assert.NoError(t, err)
	require.Equal(t, 1, len(currentStats))
	myStat = currentStats[0]
	assert.Equal(t, myHostID, myStat.HostID)
	assert.Equal(t, uint64(2), myStat.Executions)
	assert.Equal(t, myWallTime+myNewWallTime, myStat.WallTime)
	assert.Equal(t, myUserTime+myNewUserTime, myStat.UserTime)
	assert.Equal(t, mySystemTime+myNewSystemTime, myStat.SystemTime)
	assert.Equal(t, (myMemory+myNewMemory)/2, myStat.AverageMemory)
	assert.Equal(t, myOutputSize+myNewOutputSize, myStat.OutputSize)
}

func TestUpdateStats(t *testing.T) {
	ds := mysqltest.CreateMySQLDS(t)
	defer mysqltest.TruncateTables(t, ds)
	testUpdateStats(t, ds, false)
}

func TestIntegrationsUpdateStatsOnReplica(t *testing.T) {
	ds := mysqltest.CreateMySQLDSWithReplica(t, nil)
	defer mysqltest.TruncateTables(t, ds)
	testUpdateStats(t, ds, true)
}

func TestCalculateOutputSize(t *testing.T) {
	createResult := func() *fleet.DistributedQueryResult {
		result := fleet.DistributedQueryResult{}
		result.Rows = append(result.Rows, nil)
		result.Rows = append(result.Rows, map[string]string{})
		result.Rows = append(result.Rows, map[string]string{"a": "b", "a1": "b1"})
		result.Rows = append(result.Rows, map[string]string{"c": "d"})
		result.Stats = &fleet.Stats{}
		return &result
	}
	t.Run(
		"output size save disabled", func(t *testing.T) {
			tracker := statsTracker{saveStats: false}
			size := calculateOutputSize(&tracker, createResult())
			require.Equal(t, uint64(0), size)
		},
	)
	t.Run(
		"output size empty", func(t *testing.T) {
			tracker := statsTracker{saveStats: true}
			size := calculateOutputSize(&tracker, &fleet.DistributedQueryResult{})
			require.Equal(t, uint64(0), size)
		},
	)
	t.Run(
		"output size calculate", func(t *testing.T) {
			tracker := statsTracker{saveStats: true}
			size := calculateOutputSize(&tracker, createResult())
			expected := uint64(8) // manually calculated
			require.Equal(t, expected, size)
		},
	)
}