mirror of
https://github.com/fleetdm/fleet
synced 2026-05-20 15:38:39 +00:00
33a15831c0
#30877 We need to send `platform_like` during orbit enrollment for proper setup experience for Linux If some of the following don't apply, delete the relevant line. - [X] Changes file added for user-visible changes in `changes/`, `orbit/changes/` or `ee/fleetd-chrome/changes`. See [Changes files](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/guides/committing-changes.md#changes-files) for more information. ## Testing - [X] Added/updated automated tests - [X] QA'd all new/changed functionality manually ## fleetd/orbit/Fleet Desktop - [X] Verified compatibility with the latest released version of Fleet (see [Must rule](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/workflows/fleetd-development-and-release-strategy.md)) - [ ] Verified that fleetd runs on macOS, Linux and Windows - [x] Verified auto-update works from the released version of component to the new version (see [tools/tuf/test](../tools/tuf/test/README.md))
245 lines
7.6 KiB
Go
245 lines
7.6 KiB
Go
package mysqlredis
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/fleetdm/fleet/v4/server/datastore/redis/redistest"
|
|
"github.com/fleetdm/fleet/v4/server/fleet"
|
|
"github.com/fleetdm/fleet/v4/server/mock"
|
|
redigo "github.com/gomodule/redigo/redis"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestEnforceHostLimit(t *testing.T) {
|
|
const hostLimit = 3
|
|
|
|
oldBatchSize := redisSetMembersBatchSize
|
|
redisSetMembersBatchSize = 2
|
|
defer func() { redisSetMembersBatchSize = oldBatchSize }()
|
|
|
|
runTest := func(t *testing.T, pool fleet.RedisPool) {
|
|
var hostIDSeq uint
|
|
var expiredHostsIDs, incomingHostsIDs []uint
|
|
|
|
ctx := context.Background()
|
|
ds := new(mock.Store)
|
|
ds.EnrollOsqueryFunc = func(_ context.Context, opts ...fleet.DatastoreEnrollOsqueryOption) (*fleet.Host, error) {
|
|
config := &fleet.DatastoreEnrollOsqueryConfig{}
|
|
for _, opt := range opts {
|
|
opt(config)
|
|
}
|
|
hostIDSeq++
|
|
return &fleet.Host{
|
|
ID: hostIDSeq, OsqueryHostID: &config.OsqueryHostID, NodeKey: &config.NodeKey,
|
|
}, nil
|
|
}
|
|
ds.NewHostFunc = func(ctx context.Context, host *fleet.Host) (*fleet.Host, error) {
|
|
hostIDSeq++
|
|
host.ID = hostIDSeq
|
|
return host, nil
|
|
}
|
|
ds.DeleteHostFunc = func(ctx context.Context, hid uint) error {
|
|
return nil
|
|
}
|
|
ds.DeleteHostsFunc = func(ctx context.Context, ids []uint) error {
|
|
return nil
|
|
}
|
|
ds.CleanupExpiredHostsFunc = func(ctx context.Context) ([]uint, error) {
|
|
return expiredHostsIDs, nil
|
|
}
|
|
ds.CleanupIncomingHostsFunc = func(ctx context.Context, now time.Time) ([]uint, error) {
|
|
return incomingHostsIDs, nil
|
|
}
|
|
|
|
wrappedDS := New(ds, pool, WithEnforcedHostLimit(hostLimit))
|
|
|
|
requireInvokedAndReset := func(flag *bool) {
|
|
require.True(t, *flag)
|
|
*flag = false
|
|
}
|
|
requireCanEnroll := func(ok bool) {
|
|
canEnroll, err := wrappedDS.CanEnrollNewHost(ctx)
|
|
require.NoError(t, err)
|
|
require.Equal(t, ok, canEnroll)
|
|
}
|
|
|
|
// create a few hosts within the limit
|
|
h1, err := wrappedDS.NewHost(ctx, &fleet.Host{})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, h1)
|
|
requireInvokedAndReset(&ds.NewHostFuncInvoked)
|
|
requireCanEnroll(true)
|
|
h2, err := wrappedDS.EnrollOsquery(ctx,
|
|
fleet.WithEnrollOsqueryHostID("osquery-2"),
|
|
fleet.WithEnrollOsqueryNodeKey("node-2"),
|
|
fleet.WithEnrollOsqueryCooldown(time.Second),
|
|
)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, h2)
|
|
requireInvokedAndReset(&ds.EnrollOsqueryFuncInvoked)
|
|
requireCanEnroll(true)
|
|
h3, err := wrappedDS.EnrollOsquery(ctx,
|
|
fleet.WithEnrollOsqueryHostID("osquery-3"),
|
|
fleet.WithEnrollOsqueryNodeKey("node-3"),
|
|
fleet.WithEnrollOsqueryCooldown(time.Second),
|
|
)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, h3)
|
|
requireInvokedAndReset(&ds.EnrollOsqueryFuncInvoked)
|
|
requireCanEnroll(false)
|
|
|
|
// deleting h1 allows h4 to be created
|
|
err = wrappedDS.DeleteHost(ctx, h1.ID)
|
|
require.NoError(t, err)
|
|
requireCanEnroll(true)
|
|
h4, err := wrappedDS.EnrollOsquery(ctx,
|
|
fleet.WithEnrollOsqueryHostID("osquery-4"),
|
|
fleet.WithEnrollOsqueryNodeKey("node-4"),
|
|
fleet.WithEnrollOsqueryCooldown(time.Second),
|
|
)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, h4)
|
|
requireInvokedAndReset(&ds.EnrollOsqueryFuncInvoked)
|
|
requireCanEnroll(false)
|
|
|
|
// delete h1-h2-h3 (even if h1 is already deleted) should allow 2 more
|
|
err = wrappedDS.DeleteHosts(ctx, []uint{h1.ID, h2.ID, h3.ID})
|
|
require.NoError(t, err)
|
|
requireCanEnroll(true)
|
|
h5, err := wrappedDS.EnrollOsquery(ctx,
|
|
fleet.WithEnrollOsqueryHostID("osquery-5"),
|
|
fleet.WithEnrollOsqueryNodeKey("node-5"),
|
|
fleet.WithEnrollOsqueryCooldown(time.Second),
|
|
)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, h5)
|
|
requireInvokedAndReset(&ds.EnrollOsqueryFuncInvoked)
|
|
requireCanEnroll(true)
|
|
h6, err := wrappedDS.NewHost(ctx, &fleet.Host{})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, h6)
|
|
requireInvokedAndReset(&ds.NewHostFuncInvoked)
|
|
requireCanEnroll(false)
|
|
|
|
// cleanup expired removes h4
|
|
expiredHostsIDs = []uint{h4.ID}
|
|
_, err = wrappedDS.CleanupExpiredHosts(ctx)
|
|
require.NoError(t, err)
|
|
requireCanEnroll(true)
|
|
// cleanup incoming removes h4, h5
|
|
incomingHostsIDs = []uint{h4.ID, h5.ID}
|
|
_, err = wrappedDS.CleanupIncomingHosts(ctx, time.Now())
|
|
require.NoError(t, err)
|
|
requireCanEnroll(true)
|
|
|
|
// can now create 2 more
|
|
h7, err := wrappedDS.EnrollOsquery(ctx,
|
|
fleet.WithEnrollOsqueryHostID("osquery-7"),
|
|
fleet.WithEnrollOsqueryNodeKey("node-7"),
|
|
fleet.WithEnrollOsqueryCooldown(time.Second),
|
|
)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, h7)
|
|
requireInvokedAndReset(&ds.EnrollOsqueryFuncInvoked)
|
|
requireCanEnroll(true)
|
|
h8, err := wrappedDS.NewHost(ctx, &fleet.Host{})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, h8)
|
|
requireInvokedAndReset(&ds.NewHostFuncInvoked)
|
|
requireCanEnroll(false)
|
|
}
|
|
|
|
t.Run("standalone", func(t *testing.T) {
|
|
pool := redistest.SetupRedis(t, enrolledHostsSetKey, false, false, false)
|
|
runTest(t, pool)
|
|
})
|
|
|
|
t.Run("cluster", func(t *testing.T) {
|
|
pool := redistest.SetupRedis(t, enrolledHostsSetKey, true, true, false)
|
|
runTest(t, pool)
|
|
})
|
|
}
|
|
|
|
func TestSyncEnrolledHostIDs(t *testing.T) {
|
|
runTest := func(t *testing.T, pool fleet.RedisPool) {
|
|
var hostIDSeq uint
|
|
var enrolledHostCount int
|
|
var enrolledHostIDs []uint
|
|
ctx := context.Background()
|
|
|
|
ds := new(mock.Store)
|
|
ds.NewHostFunc = func(ctx context.Context, host *fleet.Host) (*fleet.Host, error) {
|
|
hostIDSeq++
|
|
host.ID = hostIDSeq
|
|
return host, nil
|
|
}
|
|
ds.CountEnrolledHostsFunc = func(ctx context.Context) (int, error) {
|
|
return enrolledHostCount, nil
|
|
}
|
|
ds.EnrolledHostIDsFunc = func(ctx context.Context) ([]uint, error) {
|
|
return enrolledHostIDs, nil
|
|
}
|
|
|
|
requireInvokedAndReset := func(flag *bool) {
|
|
require.True(t, *flag)
|
|
*flag = false
|
|
}
|
|
|
|
wrappedDS := New(ds, pool, WithEnforcedHostLimit(10)) // limit is irrelevant for this test
|
|
|
|
// create a few hosts kept in sync
|
|
h1, err := wrappedDS.NewHost(ctx, &fleet.Host{})
|
|
require.NoError(t, err)
|
|
h2, err := wrappedDS.NewHost(ctx, &fleet.Host{})
|
|
require.NoError(t, err)
|
|
h3, err := wrappedDS.NewHost(ctx, &fleet.Host{})
|
|
require.NoError(t, err)
|
|
|
|
conn := pool.Get()
|
|
defer conn.Close()
|
|
|
|
redisIDs, err := redigo.Strings(conn.Do("SMEMBERS", enrolledHostsSetKey))
|
|
require.NoError(t, err)
|
|
require.ElementsMatch(t, []string{fmt.Sprint(h1.ID), fmt.Sprint(h2.ID), fmt.Sprint(h3.ID)}, redisIDs)
|
|
|
|
// syncing with the correct count does not trigger a sync
|
|
enrolledHostCount = 3
|
|
err = wrappedDS.SyncEnrolledHostIDs(ctx)
|
|
require.NoError(t, err)
|
|
requireInvokedAndReset(&ds.CountEnrolledHostsFuncInvoked)
|
|
require.False(t, ds.EnrolledHostIDsFuncInvoked)
|
|
|
|
// syncing with a non-matching count triggers a sync
|
|
enrolledHostCount = 2
|
|
enrolledHostIDs = []uint{h1.ID, h3.ID} // will set the redis key to those values
|
|
err = wrappedDS.SyncEnrolledHostIDs(ctx)
|
|
require.NoError(t, err)
|
|
requireInvokedAndReset(&ds.CountEnrolledHostsFuncInvoked)
|
|
requireInvokedAndReset(&ds.EnrolledHostIDsFuncInvoked)
|
|
|
|
redisIDs, err = redigo.Strings(conn.Do("SMEMBERS", enrolledHostsSetKey))
|
|
require.NoError(t, err)
|
|
require.ElementsMatch(t, []string{fmt.Sprint(h1.ID), fmt.Sprint(h3.ID)}, redisIDs)
|
|
|
|
// syncing when enforcing the limit is disabled removes the set key
|
|
wrappedDS = New(ds, pool) // no limit enforced
|
|
err = wrappedDS.SyncEnrolledHostIDs(ctx)
|
|
require.NoError(t, err)
|
|
exists, err := redigo.Bool(conn.Do("EXISTS", enrolledHostsSetKey))
|
|
require.NoError(t, err)
|
|
require.False(t, exists)
|
|
}
|
|
|
|
t.Run("standalone", func(t *testing.T) {
|
|
pool := redistest.SetupRedis(t, enrolledHostsSetKey, false, false, false)
|
|
runTest(t, pool)
|
|
})
|
|
|
|
t.Run("cluster", func(t *testing.T) {
|
|
pool := redistest.SetupRedis(t, enrolledHostsSetKey, true, true, false)
|
|
runTest(t, pool)
|
|
})
|
|
}
|