From c38d498d4d24fc5de013966c6df6511209de332e Mon Sep 17 00:00:00 2001 From: Oleg Tkachuk Date: Thu, 13 Aug 2026 11:32:56 +0300 Subject: [PATCH] add unit tests for the openframe fork surfaces with no coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Guard the fork-only code against regressions during upstream syncs: - server/service/openframe: the whole agent token-auth pipeline had zero tests — AES-GCM decrypt (all key sizes, short/tampered/wrong-key payloads, error-count reset), token file extraction, the RW-mutex authorization manager (race-checked), and the refresh skip logic (empty/unchanged token, extract failures). - redis key-prefix: prefixedConn Do/Send/DoWithTimeout/Bind/ReadOnly forwarding, the PUBSUB introspection rule, WATCH/store-variant/LMOVE/ ZRANGESTORE/blocking-pop/FCALL table cases, input-slice immutability, and toInt/toString/prefixOne edge cases. - host-assignments service layer: the 8 add/remove/replace/list svc methods previously had only a manual curl script — now pin the openframe-mode gate, authz ordering (observer write rejection), host existence validation with de-duplication, and datastore delegation, plus direct tests for verifyHostsToAssociate. - mysql: openframeForeignTeam pin/foreign/unpinned decision (pure ctx logic, runs without MYSQL_TEST). --- .../mysql/openframe_foreign_team_test.go | 31 ++ server/datastore/redis/keyprefix_conn_test.go | 255 +++++++++++++++ .../host_assignments_openframe_test.go | 296 ++++++++++++++++++ server/service/openframe/openframe_test.go | 260 +++++++++++++++ 4 files changed, 842 insertions(+) create mode 100644 server/datastore/mysql/openframe_foreign_team_test.go create mode 100644 server/datastore/redis/keyprefix_conn_test.go create mode 100644 server/service/host_assignments_openframe_test.go create mode 100644 server/service/openframe/openframe_test.go diff --git a/server/datastore/mysql/openframe_foreign_team_test.go b/server/datastore/mysql/openframe_foreign_team_test.go new file mode 100644 index 00000000000..a2a18d8da5b --- /dev/null +++ b/server/datastore/mysql/openframe_foreign_team_test.go @@ -0,0 +1,31 @@ +// OPENFRAME(mysql-multitenancy): unit test for the explicit-team rejection +// helper — openframe/docs/mysql-multitenancy-feature.md +// +// openframeForeignTeam backs every "caller passed an explicit fleet_id for +// another tenant" fence; it is pure context logic, so it runs without +// MYSQL_TEST and catches a sync that breaks the pin plumbing. +package mysql + +import ( + "testing" + + "github.com/fleetdm/fleet/v4/server/fleet" + "github.com/stretchr/testify/require" +) + +func TestOpenframeForeignTeam(t *testing.T) { + ctx := t.Context() + + // Unpinned process: no tenant scope, so no team is foreign (upstream + // behavior must stay unchanged). This assumes no earlier test in this + // package pinned the process via fleet.SetOpenframeTeamID or the + // FLEET_OPENFRAME_* env vars — the pin and the env decision are cached + // process-globals. + require.False(t, openframeForeignTeam(ctx, 1)) + require.False(t, openframeForeignTeam(ctx, 0)) + + pinned := fleet.NewOpenframeTeamContext(ctx, 7) + require.False(t, openframeForeignTeam(pinned, 7), "own team is never foreign") + require.True(t, openframeForeignTeam(pinned, 8), "another tenant's team must be foreign") + require.True(t, openframeForeignTeam(pinned, 0)) +} diff --git a/server/datastore/redis/keyprefix_conn_test.go b/server/datastore/redis/keyprefix_conn_test.go new file mode 100644 index 00000000000..52618d9f7e5 --- /dev/null +++ b/server/datastore/redis/keyprefix_conn_test.go @@ -0,0 +1,255 @@ +// OPENFRAME(redis-key-prefix): unit tests for the prefixedConn wrapper plumbing +// and the prefix rules not covered by keyprefix_test.go — +// openframe/docs/redis-key-prefix.md +// +// These guard the paths where a regression is silent: a Do/Send that stops +// routing through prefixArgs, a PUBSUB introspection that leaks other tenants' +// channels, or a Bind that registers unprefixed keys with redisc. +package redis + +import ( + "reflect" + "testing" + "time" + + redigo "github.com/gomodule/redigo/redis" +) + +// recordingConn captures the command and args the wrapper forwards to the +// inner conn. It also implements ConnWithTimeout, Bind, and ReadOnly so the +// wrapper's optional-interface forwarding is testable. +type recordingConn struct { + cmds []string + args [][]interface{} + bound []string + readOnly bool +} + +func (c *recordingConn) Close() error { return nil } +func (c *recordingConn) Err() error { return nil } +func (c *recordingConn) Do(cmd string, args ...interface{}) (interface{}, error) { + c.cmds = append(c.cmds, cmd) + c.args = append(c.args, args) + return nil, nil +} +func (c *recordingConn) Send(cmd string, args ...interface{}) error { + c.cmds = append(c.cmds, cmd) + c.args = append(c.args, args) + return nil +} +func (c *recordingConn) Flush() error { return nil } +func (c *recordingConn) Receive() (interface{}, error) { return nil, nil } +func (c *recordingConn) DoWithTimeout(_ time.Duration, cmd string, args ...interface{}) (interface{}, error) { + return c.Do(cmd, args...) +} +func (c *recordingConn) ReceiveWithTimeout(time.Duration) (interface{}, error) { return nil, nil } +func (c *recordingConn) Bind(keys ...string) error { + c.bound = append(c.bound, keys...) + return nil +} +func (c *recordingConn) ReadOnly() error { + c.readOnly = true + return nil +} + +func (c *recordingConn) lastArgs(t *testing.T) []interface{} { + t.Helper() + if len(c.args) == 0 { + t.Fatal("inner conn received no command") + } + return c.args[len(c.args)-1] +} + +func TestPrefixedConnDoPrefixesKeys(t *testing.T) { + rc := &recordingConn{} + pc := newPrefixedConn(rc, "t:") + + if _, err := pc.Do("SET", "k", "v"); err != nil { + t.Fatalf("Do: %v", err) + } + if want := []interface{}{"t:k", "v"}; !reflect.DeepEqual(rc.lastArgs(t), want) { + t.Errorf("Do forwarded %v, want %v", rc.lastArgs(t), want) + } +} + +func TestPrefixedConnSendPrefixesKeys(t *testing.T) { + rc := &recordingConn{} + pc := newPrefixedConn(rc, "t:") + + if err := pc.Send("DEL", "a", "b"); err != nil { + t.Fatalf("Send: %v", err) + } + if want := []interface{}{"t:a", "t:b"}; !reflect.DeepEqual(rc.lastArgs(t), want) { + t.Errorf("Send forwarded %v, want %v", rc.lastArgs(t), want) + } +} + +func TestPrefixedConnDoWithTimeoutPrefixesKeys(t *testing.T) { + rc := &recordingConn{} + pc := newPrefixedConn(rc, "t:").(*prefixedConn) + + if _, err := pc.DoWithTimeout(time.Second, "GET", "k"); err != nil { + t.Fatalf("DoWithTimeout: %v", err) + } + if want := []interface{}{"t:k"}; !reflect.DeepEqual(rc.lastArgs(t), want) { + t.Errorf("DoWithTimeout forwarded %v, want %v", rc.lastArgs(t), want) + } +} + +func TestPrefixedConnBind(t *testing.T) { + rc := &recordingConn{} + pc := newPrefixedConn(rc, "t:").(*prefixedConn) + + if err := pc.Bind("k1", "k2"); err != nil { + t.Fatalf("Bind: %v", err) + } + if want := []string{"t:k1", "t:k2"}; !reflect.DeepEqual(rc.bound, want) { + t.Errorf("Bind forwarded %v, want %v", rc.bound, want) + } +} + +func TestPrefixedConnBindInnerWithoutBind(t *testing.T) { + var fc redigo.Conn = fakeConn{} // fakeConn has no Bind + pc := newPrefixedConn(fc, "t:").(*prefixedConn) + if err := pc.Bind("k"); err == nil { + t.Error("Bind on an inner conn without Bind must error, not panic or no-op") + } +} + +func TestPrefixedConnReadOnly(t *testing.T) { + rc := &recordingConn{} + pc := newPrefixedConn(rc, "t:").(*prefixedConn) + + if err := pc.ReadOnly(); err != nil { + t.Fatalf("ReadOnly: %v", err) + } + if !rc.readOnly { + t.Error("ReadOnly was not forwarded to the inner conn") + } + + var fc redigo.Conn = fakeConn{} + pcPlain := newPrefixedConn(fc, "t:").(*prefixedConn) + if err := pcPlain.ReadOnly(); err == nil { + t.Error("ReadOnly on an inner conn without ReadOnly must error") + } +} + +// TestPrefixArgs_MoreRules covers the rules and commands the main table in +// keyprefix_test.go does not: PUBSUB introspection, WATCH, store-variant +// multi-key commands, SORT without STORE, FCALL, and numKeys-as-string EVAL. +func TestPrefixArgs_MoreRules(t *testing.T) { + const p = "t:" + + cases := []struct { + name string + cmd string + args []interface{} + want []interface{} + }{ + // pubsubArgs: only the channel-taking subcommands get prefixed + {"PUBSUB CHANNELS", "PUBSUB", []interface{}{"CHANNELS", "ch*"}, []interface{}{"CHANNELS", "t:ch*"}}, + {"PUBSUB NUMSUB two channels", "PUBSUB", []interface{}{"NUMSUB", "c1", "c2"}, []interface{}{"NUMSUB", "t:c1", "t:c2"}}, + {"PUBSUB SHARDCHANNELS", "PUBSUB", []interface{}{"SHARDCHANNELS", "ch*"}, []interface{}{"SHARDCHANNELS", "t:ch*"}}, + {"PUBSUB SHARDNUMSUB", "PUBSUB", []interface{}{"SHARDNUMSUB", "c1"}, []interface{}{"SHARDNUMSUB", "t:c1"}}, + {"PUBSUB NUMPAT has no channels", "PUBSUB", []interface{}{"NUMPAT"}, []interface{}{"NUMPAT"}}, + {"PUBSUB with no subcommand", "PUBSUB", nil, []interface{}{}}, + + // transactions: WATCH keys are prefixed, UNWATCH/MULTI/EXEC are not + {"WATCH", "WATCH", []interface{}{"k1", "k2"}, []interface{}{"t:k1", "t:k2"}}, + + // store variants where every arg is a key + {"ZUNIONSTORE", "ZUNIONSTORE", []interface{}{"dst", "s1", "s2"}, []interface{}{"t:dst", "t:s1", "t:s2"}}, + {"PFMERGE", "PFMERGE", []interface{}{"dst", "s1"}, []interface{}{"t:dst", "t:s1"}}, + + // twoKeys variants + {"LMOVE keeps directions", "LMOVE", []interface{}{"src", "dst", "LEFT", "RIGHT"}, []interface{}{"t:src", "t:dst", "LEFT", "RIGHT"}}, + {"ZRANGESTORE keeps range", "ZRANGESTORE", []interface{}{"dst", "src", 0, -1}, []interface{}{"t:dst", "t:src", 0, -1}}, + {"BRPOPLPUSH keeps timeout", "BRPOPLPUSH", []interface{}{"src", "dst", 5}, []interface{}{"t:src", "t:dst", 5}}, + + // blocking pops + {"BZPOPMIN", "BZPOPMIN", []interface{}{"z1", "z2", 0}, []interface{}{"t:z1", "t:z2", 0}}, + {"BLPOP single key", "BLPOP", []interface{}{"k", 1}, []interface{}{"t:k", 1}}, + + // sort without STORE: only the source key + {"SORT plain", "SORT", []interface{}{"mylist", "LIMIT", 0, 10}, []interface{}{"t:mylist", "LIMIT", 0, 10}}, + + // scripting: numKeys arrives as a string over the wire too + {"EVAL numKeys as string", "EVAL", []interface{}{"script", "2", "k1", "k2", "argv"}, []interface{}{"script", "2", "t:k1", "t:k2", "argv"}}, + {"FCALL", "FCALL", []interface{}{"fn", 1, "k", "argv"}, []interface{}{"fn", 1, "t:k", "argv"}}, + {"EVAL with numKeys beyond args stays in bounds", "EVAL", []interface{}{"script", 5, "k1"}, []interface{}{"script", 5, "t:k1"}}, + {"EVAL single arg untouched", "EVAL", []interface{}{"script"}, []interface{}{"script"}}, + + // object: non-inspecting subcommand untouched + {"OBJECT HELP", "OBJECT", []interface{}{"HELP"}, []interface{}{"HELP"}}, + + // shard pub/sub channels + {"SSUBSCRIBE", "SSUBSCRIBE", []interface{}{"c1"}, []interface{}{"t:c1"}}, + + // empty args on a default-rule command must not panic + {"GET with no args", "GET", nil, []interface{}{}}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := prefixArgs(tc.cmd, tc.args, p) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("prefixArgs(%q, %v) = %v, want %v", tc.cmd, tc.args, got, tc.want) + } + }) + } +} + +func TestPrefixArgs_DoesNotMutateInput(t *testing.T) { + args := []interface{}{"k", "v"} + _ = prefixArgs("SET", args, "t:") + if args[0] != "k" { + t.Errorf("prefixArgs mutated the caller's args slice: %v", args) + } +} + +func TestPrefixOne_OutOfRange(t *testing.T) { + args := []interface{}{"k"} + prefixOne(args, -1, "t:") // must not panic + prefixOne(args, 1, "t:") // must not panic + if args[0] != "k" { + t.Errorf("out-of-range prefixOne mutated args: %v", args) + } +} + +func TestToInt(t *testing.T) { + cases := []struct { + in interface{} + want int + }{ + {2, 2}, + {int64(3), 3}, + {int32(4), 4}, + {uint(5), 5}, + {uint64(6), 6}, + {"7", 7}, + {"not a number", 0}, + {nil, 0}, + {3.9, 0}, // unsupported type: fail closed to 0 keys, not a partial prefix + } + for _, tc := range cases { + if got := toInt(tc.in); got != tc.want { + t.Errorf("toInt(%v) = %d, want %d", tc.in, got, tc.want) + } + } +} + +func TestToString(t *testing.T) { + cases := []struct { + in interface{} + want string + }{ + {"s", "s"}, + {[]byte("b"), "b"}, + {42, "42"}, + } + for _, tc := range cases { + if got := toString(tc.in); got != tc.want { + t.Errorf("toString(%v) = %q, want %q", tc.in, got, tc.want) + } + } +} diff --git a/server/service/host_assignments_openframe_test.go b/server/service/host_assignments_openframe_test.go new file mode 100644 index 00000000000..804fa0918fe --- /dev/null +++ b/server/service/host_assignments_openframe_test.go @@ -0,0 +1,296 @@ +// OPENFRAME(host-assignments): unit tests for the fork-only host-assignment +// service methods — openframe/docs/architecture-host-assignments.md +// +// The 8 svc methods (add/remove/replace/list × policy/query hosts) previously +// had only a manual curl script; this pins down the openframe-mode gate, the +// authorization order, host-existence validation (incl. de-duplication), and +// delegation to the datastore, so an upstream sync cannot silently drop any of +// them. mock.Store only — no external deps. +package service + +import ( + "context" + "errors" + "testing" + + "github.com/fleetdm/fleet/v4/server/contexts/viewer" + "github.com/fleetdm/fleet/v4/server/fleet" + "github.com/fleetdm/fleet/v4/server/mock" + "github.com/stretchr/testify/require" +) + +// newHostAssignmentTestSvc builds a service around a mock store pre-wired with +// a global policy 1 and a global query 1, plus hosts 1..3. The returned +// contexts carry the test license plus an admin/observer viewer. +func newHostAssignmentTestSvc(t *testing.T) (svc fleet.Service, ds *mock.Store, adminCtx, observerCtx context.Context) { + t.Helper() + ds = new(mock.Store) + ds.PolicyFunc = func(ctx context.Context, id uint) (*fleet.Policy, error) { + return &fleet.Policy{PolicyData: fleet.PolicyData{ID: id}}, nil + } + ds.QueryFunc = func(ctx context.Context, id uint) (*fleet.Query, error) { + return &fleet.Query{ID: id}, nil + } + ds.ListHostsLiteByIDsFunc = func(ctx context.Context, ids []uint) ([]*fleet.Host, error) { + var hosts []*fleet.Host + for _, id := range ids { + if id <= 3 { + hosts = append(hosts, &fleet.Host{ID: id}) + } + } + return hosts, nil + } + svc, baseCtx := newTestService(t, ds, nil, nil) + adminCtx = viewer.NewContext(baseCtx, viewer.Viewer{User: &fleet.User{ID: 1, GlobalRole: new(fleet.RoleAdmin)}}) + observerCtx = viewer.NewContext(baseCtx, viewer.Viewer{User: &fleet.User{ID: 2, GlobalRole: new(fleet.RoleObserver)}}) + return svc, ds, adminCtx, observerCtx +} + +// TestOpenframeHostAssignmentsRequireOpenframeMode: every method must reject +// with BadRequest when FLEET_OPENFRAME_MODE is unset, before touching authz or +// the datastore. +func TestOpenframeHostAssignmentsRequireOpenframeMode(t *testing.T) { + t.Setenv("FLEET_OPENFRAME_MODE", "") // hermetic baseline: don't depend on the ambient shell env + svc, ds, ctx, _ := newHostAssignmentTestSvc(t) + + calls := map[string]func() error{ + "AddPolicyHosts": func() error { _, err := svc.AddPolicyHosts(ctx, 1, []uint{1}); return err }, + "RemovePolicyHosts": func() error { _, err := svc.RemovePolicyHosts(ctx, 1, []uint{1}); return err }, + "ReplacePolicyHosts": func() error { return svc.ReplacePolicyHosts(ctx, 1, []uint{1}) }, + "ListPolicyHosts": func() error { _, _, err := svc.ListPolicyHosts(ctx, 1, fleet.ListOptions{}); return err }, + "AddQueryHosts": func() error { _, err := svc.AddQueryHosts(ctx, 1, []uint{1}); return err }, + "RemoveQueryHosts": func() error { _, err := svc.RemoveQueryHosts(ctx, 1, []uint{1}); return err }, + "ReplaceQueryHosts": func() error { return svc.ReplaceQueryHosts(ctx, 1, []uint{1}) }, + "ListQueryHosts": func() error { _, _, err := svc.ListQueryHosts(ctx, 1, fleet.ListOptions{}); return err }, + } + for name, call := range calls { + t.Run(name, func(t *testing.T) { + err := call() + require.Error(t, err) + var br *fleet.BadRequestError + require.ErrorAs(t, err, &br) + }) + } + require.False(t, ds.PolicyFuncInvoked, "the mode gate must fire before any datastore access") + require.False(t, ds.QueryFuncInvoked, "the mode gate must fire before any datastore access") +} + +func TestOpenframePolicyHostAssignments(t *testing.T) { + t.Setenv("FLEET_OPENFRAME_MODE", "1") + + t.Run("add delegates to the datastore", func(t *testing.T) { + svc, ds, adminCtx, _ := newHostAssignmentTestSvc(t) + var gotPolicyID uint + var gotHostIDs []uint + ds.AddPolicyHostsFunc = func(ctx context.Context, policyID uint, hostIDs []uint) (uint, error) { + gotPolicyID, gotHostIDs = policyID, hostIDs + return uint(len(hostIDs)), nil + } + n, err := svc.AddPolicyHosts(adminCtx, 1, []uint{1, 2}) + require.NoError(t, err) + require.Equal(t, uint(2), n) + require.True(t, ds.AddPolicyHostsFuncInvoked) + require.Equal(t, uint(1), gotPolicyID) + require.Equal(t, []uint{1, 2}, gotHostIDs) + }) + + t.Run("add rejects nonexistent hosts", func(t *testing.T) { + svc, ds, adminCtx, _ := newHostAssignmentTestSvc(t) + ds.AddPolicyHostsFunc = func(ctx context.Context, policyID uint, hostIDs []uint) (uint, error) { + return uint(len(hostIDs)), nil + } + _, err := svc.AddPolicyHosts(adminCtx, 1, []uint{1, 99}) + require.Error(t, err, "host 99 does not exist") + require.False(t, ds.AddPolicyHostsFuncInvoked, "validation must run before the write") + }) + + t.Run("add deduplicates host ids before validating", func(t *testing.T) { + svc, ds, adminCtx, _ := newHostAssignmentTestSvc(t) + var validatedIDs []uint + ds.ListHostsLiteByIDsFunc = func(ctx context.Context, ids []uint) ([]*fleet.Host, error) { + validatedIDs = ids + hosts := make([]*fleet.Host, len(ids)) + for i, id := range ids { + hosts[i] = &fleet.Host{ID: id} + } + return hosts, nil + } + ds.AddPolicyHostsFunc = func(ctx context.Context, policyID uint, hostIDs []uint) (uint, error) { + return uint(len(hostIDs)), nil + } + _, err := svc.AddPolicyHosts(adminCtx, 1, []uint{2, 2, 3, 2}) + require.NoError(t, err) + require.Equal(t, []uint{2, 3}, validatedIDs, + "duplicate ids must collapse or the existence count check would spuriously fail") + }) + + t.Run("observer cannot write assignments", func(t *testing.T) { + svc, ds, _, observerCtx := newHostAssignmentTestSvc(t) + _, err := svc.AddPolicyHosts(observerCtx, 1, []uint{1}) + require.Error(t, err) + _, err = svc.RemovePolicyHosts(observerCtx, 1, []uint{1}) + require.Error(t, err) + require.Error(t, svc.ReplacePolicyHosts(observerCtx, 1, []uint{1})) + require.False(t, ds.AddPolicyHostsFuncInvoked) + require.False(t, ds.RemovePolicyHostsFuncInvoked) + require.False(t, ds.ReplacePolicyHostsFuncInvoked) + }) + + t.Run("observer can list assignments", func(t *testing.T) { + svc, ds, _, observerCtx := newHostAssignmentTestSvc(t) + ds.ListPolicyHostsFunc = func(ctx context.Context, policyID uint, opts fleet.ListOptions) ([]fleet.HostIdent, *fleet.PaginationMetadata, error) { + return []fleet.HostIdent{{HostID: 1}}, nil, nil + } + hosts, _, err := svc.ListPolicyHosts(observerCtx, 1, fleet.ListOptions{}) + require.NoError(t, err) + require.Len(t, hosts, 1) + }) + + t.Run("remove delegates without host validation", func(t *testing.T) { + // Removing an already-deleted host must keep working, so Remove does + // not require the hosts to still exist. + svc, ds, adminCtx, _ := newHostAssignmentTestSvc(t) + ds.RemovePolicyHostsFunc = func(ctx context.Context, policyID uint, hostIDs []uint) (uint, error) { + return uint(len(hostIDs)), nil + } + n, err := svc.RemovePolicyHosts(adminCtx, 1, []uint{99}) + require.NoError(t, err) + require.Equal(t, uint(1), n) + require.False(t, ds.ListHostsLiteByIDsFuncInvoked) + }) + + t.Run("replace validates then delegates", func(t *testing.T) { + svc, ds, adminCtx, _ := newHostAssignmentTestSvc(t) + var gotHostIDs []uint + ds.ReplacePolicyHostsFunc = func(ctx context.Context, policyID uint, hostIDs []uint) error { + gotHostIDs = hostIDs + return nil + } + require.NoError(t, svc.ReplacePolicyHosts(adminCtx, 1, []uint{1, 3})) + require.Equal(t, []uint{1, 3}, gotHostIDs) + + require.Error(t, svc.ReplacePolicyHosts(adminCtx, 1, []uint{99})) + }) + + t.Run("missing policy surfaces the datastore error", func(t *testing.T) { + svc, ds, adminCtx, _ := newHostAssignmentTestSvc(t) + ds.PolicyFunc = func(ctx context.Context, id uint) (*fleet.Policy, error) { + return nil, errors.New("policy not found") + } + _, err := svc.AddPolicyHosts(adminCtx, 42, []uint{1}) + require.Error(t, err) + require.False(t, ds.AddPolicyHostsFuncInvoked) + }) +} + +func TestOpenframeQueryHostAssignments(t *testing.T) { + t.Setenv("FLEET_OPENFRAME_MODE", "1") + + t.Run("add validates then delegates", func(t *testing.T) { + svc, ds, adminCtx, _ := newHostAssignmentTestSvc(t) + var gotQueryID uint + var gotHostIDs []uint + ds.AddQueryHostsFunc = func(ctx context.Context, queryID uint, hostIDs []uint) (uint, error) { + gotQueryID, gotHostIDs = queryID, hostIDs + return uint(len(hostIDs)), nil + } + n, err := svc.AddQueryHosts(adminCtx, 1, []uint{2, 3}) + require.NoError(t, err) + require.Equal(t, uint(2), n) + require.Equal(t, uint(1), gotQueryID) + require.Equal(t, []uint{2, 3}, gotHostIDs) + + _, err = svc.AddQueryHosts(adminCtx, 1, []uint{99}) + require.Error(t, err, "host 99 does not exist") + }) + + t.Run("observer cannot write assignments", func(t *testing.T) { + svc, ds, _, observerCtx := newHostAssignmentTestSvc(t) + _, err := svc.AddQueryHosts(observerCtx, 1, []uint{1}) + require.Error(t, err) + _, err = svc.RemoveQueryHosts(observerCtx, 1, []uint{1}) + require.Error(t, err) + require.Error(t, svc.ReplaceQueryHosts(observerCtx, 1, []uint{1})) + require.False(t, ds.AddQueryHostsFuncInvoked) + require.False(t, ds.RemoveQueryHostsFuncInvoked) + require.False(t, ds.ReplaceQueryHostsFuncInvoked) + }) + + t.Run("remove and replace delegate", func(t *testing.T) { + svc, ds, adminCtx, _ := newHostAssignmentTestSvc(t) + ds.RemoveQueryHostsFunc = func(ctx context.Context, queryID uint, hostIDs []uint) (uint, error) { + return uint(len(hostIDs)), nil + } + ds.ReplaceQueryHostsFunc = func(ctx context.Context, queryID uint, hostIDs []uint) error { + return nil + } + n, err := svc.RemoveQueryHosts(adminCtx, 1, []uint{1, 2}) + require.NoError(t, err) + require.Equal(t, uint(2), n) + require.NoError(t, svc.ReplaceQueryHosts(adminCtx, 1, []uint{1})) + require.True(t, ds.RemoveQueryHostsFuncInvoked) + require.True(t, ds.ReplaceQueryHostsFuncInvoked) + }) + + t.Run("list delegates and returns metadata", func(t *testing.T) { + svc, ds, adminCtx, _ := newHostAssignmentTestSvc(t) + meta := &fleet.PaginationMetadata{HasNextResults: true} + ds.ListQueryHostsFunc = func(ctx context.Context, queryID uint, opts fleet.ListOptions) ([]fleet.HostIdent, *fleet.PaginationMetadata, error) { + return []fleet.HostIdent{{HostID: 1}, {HostID: 2}}, meta, nil + } + hosts, gotMeta, err := svc.ListQueryHosts(adminCtx, 1, fleet.ListOptions{}) + require.NoError(t, err) + require.Len(t, hosts, 2) + require.Equal(t, meta, gotMeta) + }) + + t.Run("missing query surfaces the datastore error", func(t *testing.T) { + svc, ds, adminCtx, _ := newHostAssignmentTestSvc(t) + ds.QueryFunc = func(ctx context.Context, id uint) (*fleet.Query, error) { + return nil, errors.New("query not found") + } + _, _, err := svc.ListQueryHosts(adminCtx, 42, fleet.ListOptions{}) + require.Error(t, err) + require.False(t, ds.ListQueryHostsFuncInvoked) + }) +} + +// TestVerifyHostsToAssociate covers the shared host-existence validator used +// by the add/replace paths. +func TestVerifyHostsToAssociate(t *testing.T) { + t.Run("empty list skips the datastore", func(t *testing.T) { + ds := new(mock.Store) + require.NoError(t, verifyHostsToAssociate(t.Context(), ds, nil)) + require.False(t, ds.ListHostsLiteByIDsFuncInvoked) + }) + + t.Run("all hosts exist", func(t *testing.T) { + ds := new(mock.Store) + ds.ListHostsLiteByIDsFunc = func(ctx context.Context, ids []uint) ([]*fleet.Host, error) { + hosts := make([]*fleet.Host, len(ids)) + for i, id := range ids { + hosts[i] = &fleet.Host{ID: id} + } + return hosts, nil + } + require.NoError(t, verifyHostsToAssociate(t.Context(), ds, []uint{1, 2})) + }) + + t.Run("a missing host fails", func(t *testing.T) { + ds := new(mock.Store) + ds.ListHostsLiteByIDsFunc = func(ctx context.Context, ids []uint) ([]*fleet.Host, error) { + return nil, nil + } + require.Error(t, verifyHostsToAssociate(t.Context(), ds, []uint{1})) + }) + + t.Run("datastore error is propagated", func(t *testing.T) { + ds := new(mock.Store) + ds.ListHostsLiteByIDsFunc = func(ctx context.Context, ids []uint) ([]*fleet.Host, error) { + return nil, errors.New("boom") + } + err := verifyHostsToAssociate(t.Context(), ds, []uint{1}) + require.Error(t, err) + require.Contains(t, err.Error(), "boom") + }) +} diff --git a/server/service/openframe/openframe_test.go b/server/service/openframe/openframe_test.go new file mode 100644 index 00000000000..84fcba4bf1a --- /dev/null +++ b/server/service/openframe/openframe_test.go @@ -0,0 +1,260 @@ +// OPENFRAME(agent-openframe-mode): unit tests for the agent token-auth pipeline — +// openframe/docs/agent-openframe-mode.md +// +// The encryption service, token extractor, authorization manager, and refresher +// carry the agent's gateway credentials; a silent regression here (e.g. an +// upstream refactor changing the payload format or dropping the refresh skip +// logic) would break every enrolled agent's authentication. Pure logic + a temp +// dir — no external deps. +package openframe + +import ( + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "encoding/base64" + "os" + "path/filepath" + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +// aes256Key is a 32-byte key matching the AES-256 deployments use. +const aes256Key = "0123456789abcdef0123456789abcdef" + +// encryptForTest produces base64(nonce || AES-GCM ciphertext) — the exact +// payload format Decrypt expects on disk. +func encryptForTest(t *testing.T, key string, plaintext []byte) string { + t.Helper() + block, err := aes.NewCipher([]byte(key)) + require.NoError(t, err) + gcm, err := cipher.NewGCM(block) + require.NoError(t, err) + nonce := make([]byte, gcm.NonceSize()) + _, err = rand.Read(nonce) + require.NoError(t, err) + return base64.StdEncoding.EncodeToString(gcm.Seal(nonce, nonce, plaintext, nil)) +} + +func TestDecryptRoundTrip(t *testing.T) { + // All three AES key sizes must work: the key comes from deployment config, + // not from code, so none of the sizes is "the" supported one. + keys := map[string]string{ + "AES-128": "0123456789abcdef", + "AES-192": "0123456789abcdef01234567", + "AES-256": aes256Key, + } + for name, key := range keys { + t.Run(name, func(t *testing.T) { + es := NewOpenframeEncryptionService(key) + payload := encryptForTest(t, key, []byte("the-token")) + got, err := es.Decrypt(payload) + require.NoError(t, err) + require.Equal(t, []byte("the-token"), got) + }) + } +} + +func TestDecryptEmptyPlaintext(t *testing.T) { + es := NewOpenframeEncryptionService(aes256Key) + got, err := es.Decrypt(encryptForTest(t, aes256Key, []byte{})) + require.NoError(t, err) + require.Empty(t, got) +} + +func TestDecryptInvalidBase64(t *testing.T) { + es := NewOpenframeEncryptionService(aes256Key) + _, err := es.Decrypt("%%% not base64 %%%") + require.Error(t, err) +} + +func TestDecryptBadKeyLength(t *testing.T) { + es := NewOpenframeEncryptionService("short-key") + _, err := es.Decrypt(base64.StdEncoding.EncodeToString([]byte("whatever"))) + require.Error(t, err) +} + +func TestDecryptCiphertextTooShort(t *testing.T) { + // Fewer bytes than the 12-byte GCM nonce must error, not panic on slicing. + es := NewOpenframeEncryptionService(aes256Key) + _, err := es.Decrypt(base64.StdEncoding.EncodeToString([]byte("tiny"))) + require.Error(t, err) + require.Contains(t, err.Error(), "ciphertext too short") +} + +func TestDecryptWrongKey(t *testing.T) { + payload := encryptForTest(t, aes256Key, []byte("the-token")) + es := NewOpenframeEncryptionService("fedcba9876543210fedcba9876543210") + _, err := es.Decrypt(payload) + require.Error(t, err, "GCM must reject a payload sealed with a different key") +} + +func TestDecryptTamperedCiphertext(t *testing.T) { + payload := encryptForTest(t, aes256Key, []byte("the-token")) + raw, err := base64.StdEncoding.DecodeString(payload) + require.NoError(t, err) + raw[len(raw)-1] ^= 0xff + es := NewOpenframeEncryptionService(aes256Key) + _, err = es.Decrypt(base64.StdEncoding.EncodeToString(raw)) + require.Error(t, err, "GCM must reject a tampered ciphertext") +} + +func TestDecryptErrorCountResetsOnSuccess(t *testing.T) { + es := NewOpenframeEncryptionService(aes256Key) + _, err := es.Decrypt("%%% not base64 %%%") + require.Error(t, err) + require.Equal(t, 1, es.decryptErrCount) + + _, err = es.Decrypt(encryptForTest(t, aes256Key, []byte("ok"))) + require.NoError(t, err) + require.Equal(t, 0, es.decryptErrCount, "a successful decrypt must reset the rate-limit counter") +} + +func writeTokenFile(t *testing.T, contents string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "token") + require.NoError(t, os.WriteFile(path, []byte(contents), 0o600)) + return path +} + +func TestExtractToken(t *testing.T) { + es := NewOpenframeEncryptionService(aes256Key) + + t.Run("reads and decrypts the token file", func(t *testing.T) { + path := writeTokenFile(t, encryptForTest(t, aes256Key, []byte("agent-token"))) + te := NewOpenframeTokenExtractor(es, path) + token, err := te.ExtractToken() + require.NoError(t, err) + require.Equal(t, "agent-token", token) + }) + + t.Run("missing file errors and counts", func(t *testing.T) { + te := NewOpenframeTokenExtractor(es, filepath.Join(t.TempDir(), "does-not-exist")) + _, err := te.ExtractToken() + require.Error(t, err) + require.Equal(t, 1, te.readErrCount) + }) + + t.Run("read error count resets on success", func(t *testing.T) { + path := writeTokenFile(t, encryptForTest(t, aes256Key, []byte("agent-token"))) + te := NewOpenframeTokenExtractor(es, path) + te.readErrCount = 3 + _, err := te.ExtractToken() + require.NoError(t, err) + require.Equal(t, 0, te.readErrCount) + }) + + t.Run("undecryptable file contents error", func(t *testing.T) { + path := writeTokenFile(t, "not-even-base64 %%%") + te := NewOpenframeTokenExtractor(es, path) + _, err := te.ExtractToken() + require.Error(t, err) + }) +} + +func TestAuthorizationManager(t *testing.T) { + t.Run("starts empty", func(t *testing.T) { + m := NewOpenFrameAuthorizationManager() + require.Equal(t, "", m.GetToken()) + }) + + t.Run("with-token constructor seeds the token", func(t *testing.T) { + m := NewOpenFrameAuthorizationManagerWithToken("seed") + require.Equal(t, "seed", m.GetToken()) + }) + + t.Run("update replaces the token", func(t *testing.T) { + m := NewOpenFrameAuthorizationManagerWithToken("old") + m.UpdateToken("new") + require.Equal(t, "new", m.GetToken()) + }) + + // The manager is read on every request while the cron refresher writes; + // run with -race to catch a dropped mutex. + t.Run("concurrent readers and writers", func(t *testing.T) { + m := NewOpenFrameAuthorizationManager() + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + wg.Add(2) + go func() { + defer wg.Done() + for j := 0; j < 100; j++ { + m.UpdateToken("tok") + } + }() + go func() { + defer wg.Done() + for j := 0; j < 100; j++ { + _ = m.GetToken() + } + }() + } + wg.Wait() + require.Equal(t, "tok", m.GetToken()) + }) +} + +func TestRefreshToken(t *testing.T) { + es := NewOpenframeEncryptionService(aes256Key) + + newRefresher := func(t *testing.T, tokenPath string) (*OpenframeTokenRefresher, *OpenFrameAuthorizationManager) { + t.Helper() + mgr := NewOpenFrameAuthorizationManager() + te := NewOpenframeTokenExtractor(es, tokenPath) + return NewOpenframeTokenRefresher(te, mgr), mgr + } + + t.Run("stores a freshly extracted token", func(t *testing.T) { + path := writeTokenFile(t, encryptForTest(t, aes256Key, []byte("tok-1"))) + tr, mgr := newRefresher(t, path) + tr.refreshToken() + require.Equal(t, "tok-1", mgr.GetToken()) + }) + + t.Run("picks up a rotated token", func(t *testing.T) { + path := writeTokenFile(t, encryptForTest(t, aes256Key, []byte("tok-1"))) + tr, mgr := newRefresher(t, path) + tr.refreshToken() + require.NoError(t, os.WriteFile(path, []byte(encryptForTest(t, aes256Key, []byte("tok-2"))), 0o600)) + tr.refreshToken() + require.Equal(t, "tok-2", mgr.GetToken()) + }) + + t.Run("empty extracted token is not stored", func(t *testing.T) { + path := writeTokenFile(t, encryptForTest(t, aes256Key, []byte{})) + tr, mgr := newRefresher(t, path) + mgr.UpdateToken("existing") + tr.refreshToken() + require.Equal(t, "existing", mgr.GetToken(), "an empty token must never clobber a working one") + }) + + t.Run("extract failure keeps the current token", func(t *testing.T) { + tr, mgr := newRefresher(t, filepath.Join(t.TempDir(), "gone")) + mgr.UpdateToken("existing") + tr.refreshToken() + require.Equal(t, "existing", mgr.GetToken()) + require.Equal(t, 1, tr.extractErrCount) + }) + + t.Run("error count resets on the next successful extract", func(t *testing.T) { + path := filepath.Join(t.TempDir(), "token") + mgr := NewOpenFrameAuthorizationManager() + tr := NewOpenframeTokenRefresher(NewOpenframeTokenExtractor(es, path), mgr) + tr.refreshToken() // file missing + require.Equal(t, 1, tr.extractErrCount) + require.NoError(t, os.WriteFile(path, []byte(encryptForTest(t, aes256Key, []byte("tok"))), 0o600)) + tr.refreshToken() + require.Equal(t, 0, tr.extractErrCount) + require.Equal(t, "tok", mgr.GetToken()) + }) +} + +func TestRefresherStartStop(t *testing.T) { + path := writeTokenFile(t, encryptForTest(t, aes256Key, []byte("tok"))) + mgr := NewOpenFrameAuthorizationManager() + tr := NewOpenframeTokenRefresher(NewOpenframeTokenExtractor(NewOpenframeEncryptionService(aes256Key), path), mgr) + require.NoError(t, tr.Start()) + tr.Stop() // must not hang waiting for jobs +}