diff --git a/cmd/api/auth.go b/cmd/api/auth.go index a529b9d9..320a2ccd 100644 --- a/cmd/api/auth.go +++ b/cmd/api/auth.go @@ -5,6 +5,7 @@ import ( "crypto/subtle" "net/http" "strings" + "time" "github.com/jmpsec/osctrl/cmd/api/handlers" "github.com/jmpsec/osctrl/pkg/config" @@ -156,7 +157,7 @@ func handlerAuthCheck(h http.Handler, auth, jwtSecret string) http.Handler { return } // Update metadata for the user - if err := apiUsers.UpdateTokenIPAddress(utils.GetIP(r), claims.Username); err != nil { + if err := apiUsers.RecordTokenUse(user, utils.GetIP(r), time.Now()); err != nil { log.Err(err).Msgf("error updating token for user %s", claims.Username) } // Set middleware values diff --git a/cmd/api/handlers/stats.go b/cmd/api/handlers/stats.go index 1d3010e0..42a3dd83 100644 --- a/cmd/api/handlers/stats.go +++ b/cmd/api/handlers/stats.go @@ -15,7 +15,6 @@ import ( "github.com/jmpsec/osctrl/pkg/dbutil" "github.com/jmpsec/osctrl/pkg/logging" "github.com/jmpsec/osctrl/pkg/nodes" - "github.com/jmpsec/osctrl/pkg/queries" "github.com/jmpsec/osctrl/pkg/settings" "github.com/jmpsec/osctrl/pkg/users" "github.com/jmpsec/osctrl/pkg/utils" @@ -107,11 +106,22 @@ func (h *HandlersApi) StatsHandler(w http.ResponseWriter, r *http.Request) { Environments: make([]EnvStats, 0, len(allEnvs)), } + // Active query and carve totals for every environment in one grouped + // count. If it fails, every environment is left out, as the per-env + // lookups it replaces skipped any environment whose count failed. + activeCounts, countErr := h.Queries.ActiveCountsByEnvironment() + if countErr != nil { + log.Warn().Err(countErr).Msg("stats: failed to count active queries and carves, skipping environments") + } + for _, e := range allEnvs { // Filter to envs the user can actually see. if !h.Users.CheckPermissions(user, users.UserLevel, e.UUID) { continue } + if countErr != nil { + continue + } envHours := h.Settings.InactiveHours(e.ID) ns, err := h.Nodes.GetStatsByEnv(e.Name, envHours) @@ -129,19 +139,7 @@ func (h *HandlersApi) StatsHandler(w http.ResponseWriter, r *http.Request) { log.Warn().Err(err).Str("env", e.Name).Msg("stats: failed to get platform counts, defaulting to zeros") } - // Use type-specific methods to avoid double-counting: - // GetQueries returns StandardQueryType active items only. - // GetCarves returns CarveQueryType active items only. - activeQ, err := h.Queries.GetQueries(queries.TargetActive, e.ID) - if err != nil { - log.Warn().Err(err).Str("env", e.Name).Msg("stats: failed to count active queries, skipping env") - continue - } - activeC, err := h.Queries.GetCarves(queries.TargetActive, e.ID) - if err != nil { - log.Warn().Err(err).Str("env", e.Name).Msg("stats: failed to count active carves, skipping env") - continue - } + active := activeCounts[e.ID] row := EnvStats{ InactiveHours: envHours, @@ -150,16 +148,16 @@ func (h *HandlersApi) StatsHandler(w http.ResponseWriter, r *http.Request) { Active: ns.Active, Inactive: ns.Inactive, Total: ns.Total, - ActiveQueries: len(activeQ), - ActiveCarves: len(activeC), + ActiveQueries: active.Queries, + ActiveCarves: active.Carves, PlatformCounts: platCounts, } out.Environments = append(out.Environments, row) out.ActiveNodes += ns.Active out.InactiveNodes += ns.Inactive out.TotalNodes += ns.Total - out.TotalActiveQueries += len(activeQ) - out.TotalActiveCarves += len(activeC) + out.TotalActiveQueries += active.Queries + out.TotalActiveCarves += active.Carves // Aggregate cross-env platform totals. out.PlatformCounts.Linux += platCounts.Linux out.PlatformCounts.Darwin += platCounts.Darwin @@ -923,18 +921,17 @@ func (h *HandlersApi) NodeActivityTilesBatchHandler(w http.ResponseWriter, r *ht return } - // Only include nodes that actually belong to this env. + // Only include nodes that actually belong to this env. A lookup failure + // leaves nodes out rather than failing the poll, as the per-node lookups + // it replaces did. out := make(map[string]activity.NodeTileSeries, len(uuids)) - for _, u := range uuids { - node, err := h.Nodes.GetByUUID(u) - if err != nil { - continue - } - if !strings.EqualFold(node.Environment, env.Name) { - continue - } - if s, ok := series[node.UUID]; ok { - out[node.UUID] = s + inEnv, err := h.Nodes.UUIDsInEnvironment(uuids, env.Name) + if err != nil { + log.Warn().Err(err).Str("env", env.Name).Msg("activity tiles: node membership lookup failed") + } + for u := range inEnv { + if s, ok := series[u]; ok { + out[u] = s } } diff --git a/cmd/api/handlers/stats_counts_test.go b/cmd/api/handlers/stats_counts_test.go new file mode 100644 index 00000000..6ca25cbf --- /dev/null +++ b/cmd/api/handlers/stats_counts_test.go @@ -0,0 +1,54 @@ +package handlers + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/jmpsec/osctrl/pkg/auditlog" + "github.com/jmpsec/osctrl/pkg/config" + "github.com/jmpsec/osctrl/pkg/queries" + "github.com/jmpsec/osctrl/pkg/settings" +) + +// The stats endpoint reports active queries and carves per environment from +// one grouped count. Console and file-explorer rows share the table and must +// not be counted, nor finished queries. +func TestStatsHandlerActiveQueryAndCarveCounts(t *testing.T) { + db, h, env, _ := setupConsoleHandlers(t) + h.DebugHTTPConfig = &config.YAMLConfigurationDebug{} + var err error + h.AuditLog, err = auditlog.CreateAuditLogManager(db, config.ServiceAPI, false) + require.NoError(t, err) + require.NoError(t, h.Settings.NewIntegerValue(config.ServiceAPI, settings.InactiveHours, 24, 0)) + + for i, dq := range []queries.DistributedQuery{ + {Type: queries.StandardQueryType, Active: true}, + {Type: queries.StandardQueryType, Active: true}, + {Type: queries.StandardQueryType, Active: true, Completed: true}, + {Type: queries.CarveQueryType, Active: true}, + {Type: queries.ConsoleQueryType, Active: true}, + {Type: queries.FileExplorerQueryType, Active: true}, + } { + dq.Name = fmt.Sprintf("stats_q%d", i) + dq.Query = "SELECT 1;" + dq.EnvironmentID = env.ID + require.NoError(t, db.Create(&dq).Error) + } + + rr := httptest.NewRecorder() + h.StatsHandler(rr, consoleRequest(http.MethodGet, "/stats", nil, "alice")) + require.Equal(t, http.StatusOK, rr.Code, rr.Body.String()) + + var stats StatsResponse + require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &stats)) + require.Len(t, stats.Environments, 1) + require.Equal(t, 2, stats.Environments[0].ActiveQueries) + require.Equal(t, 1, stats.Environments[0].ActiveCarves) + require.Equal(t, 2, stats.TotalActiveQueries) + require.Equal(t, 1, stats.TotalActiveCarves) +} diff --git a/pkg/logging/db.go b/pkg/logging/db.go index 58a71de5..b5e9d3d8 100644 --- a/pkg/logging/db.go +++ b/pkg/logging/db.go @@ -148,9 +148,11 @@ func (logDB *LoggerDB) Status(data []byte, environment, uuid string, debug bool) if err := json.Unmarshal(data, &logs); err != nil { log.Err(err).Msgf("error parsing logs %s %v", string(data), err) } - // Iterate and insert in DB + stamp := logDB.rowStamper() + entries := make([]OsqueryStatusData, 0, len(logs)) for _, l := range logs { - entry := OsqueryStatusData{ + entries = append(entries, OsqueryStatusData{ + Model: stamp(), UUID: strings.ToUpper(l.HostIdentifier), Environment: environment, Line: strconv.Itoa(int(l.Line)), @@ -158,11 +160,9 @@ func (logDB *LoggerDB) Status(data []byte, environment, uuid string, debug bool) Version: l.Version, Filename: l.Filename, Severity: strconv.Itoa(int(l.Severity)), - } - if err := logDB.Database.Conn.Create(&entry).Error; err != nil { - log.Err(err).Msg("Error creating status log entry") - } + }) } + insertLogRows(logDB.Database.Conn, entries, "status") } // Result - Function that sends JSON result logs to the configured DB @@ -172,9 +172,11 @@ func (logDB *LoggerDB) Result(data []byte, environment, uuid string, debug bool) if err != nil { log.Err(err).Msgf("error parsing logs %s", string(data)) } - // Iterate and insert in DB + stamp := logDB.rowStamper() + entries := make([]OsqueryResultData, 0, len(logs)) for _, l := range logs { - entry := OsqueryResultData{ + entries = append(entries, OsqueryResultData{ + Model: stamp(), UUID: strings.ToUpper(l.HostIdentifier), Environment: environment, Name: l.Name, @@ -182,10 +184,70 @@ func (logDB *LoggerDB) Result(data []byte, environment, uuid string, debug bool) Epoch: l.Epoch, Columns: string(l.Columns), Counter: l.Counter, + }) + } + insertLogRows(logDB.Database.Conn, entries, "result") +} + +// logInsertBatchSize caps rows per multi-row INSERT. Both log models bind 10 +// columns per row (the gorm.Model timestamps plus seven fields), so 50 rows is +// 500 bind parameters: under the 999 limit of older SQLite builds and far +// under the 65535 of Postgres and MySQL. +const logInsertBatchSize = 50 + +// insertLogRows writes one osquery log batch with multi-row INSERTs inside a +// single transaction. +// +// Row-by-row inserts each ran in GORM's implicit per-statement transaction, +// so a batch of N lines cost N commits — N WAL flushes on Postgres — on the +// busiest write path in osctrl-tls. One transaction makes that one commit. +// +// The batch is all-or-nothing, which row-by-row was not: one unstorable row +// (Postgres rejects NUL bytes in text, for instance) used to cost only itself. +// So on failure the rolled-back batch is retried row by row from a pristine +// copy, keeping the old partial-success behaviour without duplicating +// anything — the transaction guarantees none of the batch was kept. +func insertLogRows[T any](db *gorm.DB, entries []T, kind string) { + if len(entries) == 0 { + return + } + // CreateInBatches writes generated IDs back into the rows it inserted, + // even in a batch that is later rolled back. Retrying those rows would + // insert explicit IDs, so the fallback works from an untouched copy. + pristine := append([]T(nil), entries...) + err := db.Transaction(func(tx *gorm.DB) error { + return tx.CreateInBatches(entries, logInsertBatchSize).Error + }) + if err == nil { + return + } + log.Warn().Err(err).Int("rows", len(entries)).Msgf("batched %s log insert failed, retrying row by row", kind) + for i := range pristine { + if err := db.Create(&pristine[i]).Error; err != nil { + log.Err(err).Msgf("Error creating %s log entry", kind) } - if err := logDB.Database.Conn.Create(&entry).Error; err != nil { - log.Err(err).Msg("Error creating result log entry") - } + } +} + +// rowStamper returns a function handing out strictly increasing timestamps, +// one per row of a log batch. +// +// GORM stamps every row of a multi-row INSERT with the same created_at, but +// the log readers order by created_at alone, so ties would let the lines of +// one batch come back shuffled. Row-by-row inserts never tied in practice; +// stepping by the backend's storage precision (MySQL keeps milliseconds, +// Postgres microseconds) keeps every row distinct after the database rounds +// it, and preserves the order osquery sent the lines in. +func (logDB *LoggerDB) rowStamper() func() gorm.Model { + step := time.Microsecond + if logDB.Database.Conn.Name() == "mysql" { + step = time.Millisecond + } + next := time.Now() + return func() gorm.Model { + ts := next + next = next.Add(step) + return gorm.Model{CreatedAt: ts, UpdatedAt: ts} } } diff --git a/pkg/logging/db_insert_test.go b/pkg/logging/db_insert_test.go new file mode 100644 index 00000000..8b532848 --- /dev/null +++ b/pkg/logging/db_insert_test.go @@ -0,0 +1,186 @@ +package logging + +import ( + "encoding/json" + "fmt" + "strings" + "testing" + + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/jmpsec/osctrl/pkg/backend" +) + +func newTestLoggerDB(t *testing.T) *LoggerDB { + t.Helper() + dsn := "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared" + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + logDB, err := CreateLoggerDB(&backend.DBManager{Conn: db}) + if err != nil { + t.Fatalf("CreateLoggerDB: %v", err) + } + return logDB +} + +func statusBatch(t *testing.T, messages []string) []byte { + t.Helper() + type line struct { + HostIdentifier string `json:"hostIdentifier"` + Message string `json:"message"` + Line int `json:"line"` + } + lines := make([]line, len(messages)) + for i, m := range messages { + lines[i] = line{HostIdentifier: "node-a", Message: m, Line: i} + } + b, err := json.Marshal(lines) + if err != nil { + t.Fatalf("marshal: %v", err) + } + return b +} + +func numbered(prefix string, n int) []string { + out := make([]string, n) + for i := range out { + out[i] = fmt.Sprintf("%s-%03d", prefix, i) + } + return out +} + +// A batch larger than one INSERT must land completely, across several +// multi-row statements. +func TestStatusLogsSpanningSeveralInsertBatches(t *testing.T) { + logDB := newTestLoggerDB(t) + want := numbered("line", 2*logInsertBatchSize+7) + + logDB.Status(statusBatch(t, want), "prod", "node-a", false) + + var count int64 + logDB.Database.Conn.Model(&OsqueryStatusData{}).Count(&count) + if count != int64(len(want)) { + t.Fatalf("stored %d status rows, want %d", count, len(want)) + } +} + +// Readers order by created_at alone. A multi-row INSERT would give every row +// the same timestamp and let the lines of one batch come back shuffled; each +// row must get its own, in the order osquery sent them. +func TestStatusLogsKeepSendOrder(t *testing.T) { + logDB := newTestLoggerDB(t) + want := numbered("line", 3*logInsertBatchSize) + + logDB.Status(statusBatch(t, want), "prod", "node-a", false) + + var rows []OsqueryStatusData + if err := logDB.Database.Conn.Order("created_at").Find(&rows).Error; err != nil { + t.Fatalf("read back: %v", err) + } + if len(rows) != len(want) { + t.Fatalf("read %d rows, want %d", len(rows), len(want)) + } + for i, r := range rows { + if r.Message != want[i] { + t.Fatalf("row %d = %q, want %q: batch order not preserved", i, r.Message, want[i]) + } + if i > 0 && !r.CreatedAt.After(rows[i-1].CreatedAt) { + t.Fatalf("row %d created_at %v not after row %d's %v", i, r.CreatedAt, i-1, rows[i-1].CreatedAt) + } + } + // The newest-first reader the UI uses must see the same order reversed. + got, err := logDB.StatusLogsLimit("node-a", "prod", 5) + if err != nil { + t.Fatalf("StatusLogsLimit: %v", err) + } + for i, r := range got { + if r.Message != want[len(want)-1-i] { + t.Fatalf("newest-first row %d = %q, want %q", i, r.Message, want[len(want)-1-i]) + } + } +} + +// The batch is one transaction, so one unstorable row would otherwise drop +// every line sent with it. Row-by-row inserts only ever lost the bad row, and +// that must still hold — without the rest being written twice. +func TestStatusLogsBadRowDoesNotDropTheBatch(t *testing.T) { + logDB := newTestLoggerDB(t) + // SQLite stores almost anything, so a trigger stands in for a value the + // production database would refuse. + if err := logDB.Database.Conn.Exec(`CREATE TRIGGER reject_poison BEFORE INSERT ON osquery_status_data + WHEN NEW.message = 'poison' BEGIN SELECT RAISE(ABORT, 'poison row'); END`).Error; err != nil { + t.Fatalf("create trigger: %v", err) + } + // Put the bad row in a later INSERT, so earlier ones had already + // succeeded inside the transaction when it rolled back. + msgs := numbered("line", logInsertBatchSize+10) + msgs[logInsertBatchSize+3] = "poison" + + logDB.Status(statusBatch(t, msgs), "prod", "node-a", false) + + var rows []OsqueryStatusData + logDB.Database.Conn.Order("created_at").Find(&rows) + if len(rows) != len(msgs)-1 { + t.Fatalf("stored %d rows, want %d: everything except the poison row", len(rows), len(msgs)-1) + } + seen := make(map[string]bool, len(rows)) + for _, r := range rows { + if r.Message == "poison" { + t.Fatal("poison row was stored") + } + if seen[r.Message] { + t.Fatalf("row %q stored twice: the rolled-back batch leaked", r.Message) + } + seen[r.Message] = true + } +} + +func TestResultLogsAreBatched(t *testing.T) { + logDB := newTestLoggerDB(t) + type result struct { + Name string `json:"name"` + HostIdentifier string `json:"hostIdentifier"` + Columns map[string]string `json:"columns"` + } + n := logInsertBatchSize + 1 + batch := make([]result, n) + for i := range batch { + batch[i] = result{Name: "pack_uptime", HostIdentifier: "node-a", Columns: map[string]string{"seq": fmt.Sprint(i)}} + } + data, err := json.Marshal(batch) + if err != nil { + t.Fatalf("marshal: %v", err) + } + + logDB.Result(data, "prod", "node-a", false) + + var rows []OsqueryResultData + logDB.Database.Conn.Order("created_at").Find(&rows) + if len(rows) != n { + t.Fatalf("stored %d result rows, want %d", len(rows), n) + } + for i, r := range rows { + if want := fmt.Sprintf(`{"seq":"%d"}`, i); r.Columns != want { + t.Fatalf("row %d columns = %s, want %s", i, r.Columns, want) + } + if r.UUID != "NODE-A" || r.Environment != "prod" { + t.Fatalf("row %d = %+v, fields not carried over", i, r) + } + } +} + +func TestEmptyLogBatchIsANoop(t *testing.T) { + logDB := newTestLoggerDB(t) + logDB.Status([]byte(`[]`), "prod", "node-a", false) + logDB.Result([]byte(`[]`), "prod", "node-a", false) + + var status, result int64 + logDB.Database.Conn.Model(&OsqueryStatusData{}).Count(&status) + logDB.Database.Conn.Model(&OsqueryResultData{}).Count(&result) + if status != 0 || result != 0 { + t.Fatalf("empty batches stored %d status / %d result rows", status, result) + } +} diff --git a/pkg/logging/dispatch.go b/pkg/logging/dispatch.go index 178505ce..52c218f9 100644 --- a/pkg/logging/dispatch.go +++ b/pkg/logging/dispatch.go @@ -15,7 +15,7 @@ import ( // etc.). func (l *LoggerTLS) DispatchLogs(data []byte, uuid, logType string, envID uint, environment string, metadata nodes.NodeMetadata, debug bool) { // Use metadata to update record - if err := l.Nodes.UpdateMetadataByUUID(uuid, metadata); err != nil { + if err := l.Nodes.UpdateMetadataByUUID(uuid, envID, metadata); err != nil { log.Err(err).Msg("error updating metadata") } if debug { diff --git a/pkg/nodes/metadata_test.go b/pkg/nodes/metadata_test.go new file mode 100644 index 00000000..2d1ec234 --- /dev/null +++ b/pkg/nodes/metadata_test.go @@ -0,0 +1,164 @@ +package nodes + +import ( + "errors" + "fmt" + "os" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + "gorm.io/driver/mysql" + "gorm.io/driver/postgres" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + "gorm.io/gorm/schema" +) + +func TestUpdateMetadataByUUID(t *testing.T) { + db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + testUpdateMetadataByUUID(t, db) +} + +// Optional real-engine checks; each uses and removes its own prefixed table. +func TestUpdateMetadataByUUIDExternalBackend(t *testing.T) { + for _, backend := range []string{"POSTGRES", "MYSQL"} { + t.Run(backend, func(t *testing.T) { + dsn := os.Getenv("OSCTRL_TEST_" + backend + "_DSN") + if dsn == "" { + t.Skip("set OSCTRL_TEST_" + backend + "_DSN to test this engine") + } + dialect := postgres.Open(dsn) + if backend == "MYSQL" { + dialect = mysql.Open(dsn) + } + db, err := gorm.Open(dialect, &gorm.Config{NamingStrategy: schema.NamingStrategy{ + TablePrefix: fmt.Sprintf("metadata_%d_", time.Now().UnixNano()), + }}) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable(&OsqueryNode{})) }) + testUpdateMetadataByUUID(t, db) + }) + } +} + +func testUpdateMetadataByUUID(t *testing.T, db *gorm.DB) { + t.Helper() + require.NoError(t, db.AutoMigrate(&OsqueryNode{})) + repo := NodeManager{DB: db} + + // The same UUID enrolled in two environments has a live row in each. + rows := []OsqueryNode{ + {UUID: "NODE-A", EnvironmentID: 1, Hostname: "old-host", Username: "keep-me", BytesReceived: 100}, + {UUID: "NODE-A", EnvironmentID: 2, Hostname: "other-env", BytesReceived: 7}, + {UUID: "NODE-B", EnvironmentID: 1, Hostname: "unrelated", BytesReceived: 50}, + } + require.NoError(t, db.Create(&rows).Error) + + var selects, updates int + require.NoError(t, db.Callback().Query().After("gorm:query").Register("count_metadata_selects", func(*gorm.DB) { selects++ })) + require.NoError(t, db.Callback().Update().After("gorm:update").Register("count_metadata_updates", func(*gorm.DB) { updates++ })) + + // Lower-case UUID, as host identifiers can arrive; empty fields must not + // blank out what is stored. + require.NoError(t, repo.UpdateMetadataByUUID("node-a", 1, NodeMetadata{ + Hostname: "new-host", + ConfigHash: "abc123", + OsqueryVersion: "5.23.1", + BytesReceived: 25, + })) + require.Equal(t, 0, selects, "the per-log-batch hot path must not read the node first") + require.Equal(t, 1, updates, "metadata must be a single UPDATE") + + got := map[string]OsqueryNode{} + var all []OsqueryNode + require.NoError(t, db.Order("id").Find(&all).Error) + for _, n := range all { + got[fmt.Sprintf("%s/%d", n.UUID, n.EnvironmentID)] = n + } + target := got["NODE-A/1"] + require.Equal(t, "new-host", target.Hostname) + require.Equal(t, "abc123", target.ConfigHash) + require.Equal(t, "5.23.1", target.OsqueryVersion) + require.Equal(t, "keep-me", target.Username, "an empty field must leave the stored value alone") + require.Equal(t, 125, target.BytesReceived) + + // Only the row of the authenticated environment is touched. + require.Equal(t, "other-env", got["NODE-A/2"].Hostname, "same UUID in another environment must be untouched") + require.Equal(t, 7, got["NODE-A/2"].BytesReceived) + require.Equal(t, "unrelated", got["NODE-B/1"].Hostname) + require.Equal(t, 50, got["NODE-B/1"].BytesReceived) + + // Unknown nodes still report not-found, as the read-first version did. + err := repo.UpdateMetadataByUUID("NODE-MISSING", 1, NodeMetadata{BytesReceived: 1}) + require.Error(t, err) + require.True(t, errors.Is(err, gorm.ErrRecordNotFound), "got %v", err) + err = repo.UpdateMetadataByUUID("NODE-A", 99, NodeMetadata{BytesReceived: 1}) + require.True(t, errors.Is(err, gorm.ErrRecordNotFound), "wrong environment must not match: got %v", err) +} + +// Log batches from one node are processed in parallel goroutines. Reading +// bytes_received and writing back a sum lost increments whenever two +// overlapped; the SQL-side increment must account for every byte. +func TestUpdateMetadataByUUIDConcurrentBytes(t *testing.T) { + dir := t.TempDir() + // A file database so connections share state and SQLite serialises + // writers for real, rather than one in-memory connection doing it all. + db, err := gorm.Open(sqlite.Open(dir+"/nodes.db?_busy_timeout=5000"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&OsqueryNode{})) + repo := NodeManager{DB: db} + require.NoError(t, db.Create(&OsqueryNode{UUID: "NODE-A", EnvironmentID: 1}).Error) + + const writers, each = 8, 25 + var wg sync.WaitGroup + for range writers { + wg.Go(func() { + for range each { + if err := repo.UpdateMetadataByUUID("NODE-A", 1, NodeMetadata{BytesReceived: 3}); err != nil { + t.Errorf("update: %v", err) + return + } + } + }) + } + wg.Wait() + + var n OsqueryNode + require.NoError(t, db.Where("uuid = ?", "NODE-A").First(&n).Error) + require.Equal(t, writers*each*3, n.BytesReceived, "concurrent batches lost increments") +} + +func TestUUIDsInEnvironment(t *testing.T) { + db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&OsqueryNode{})) + repo := NodeManager{DB: db} + + rows := []OsqueryNode{ + {UUID: "IN-PROD", Environment: "prod"}, + {UUID: "IN-STAGING", Environment: "staging"}, + {UUID: "DELETED", Environment: "prod"}, + // Enrolled in both: membership must not depend on which row is first. + {UUID: "BOTH", Environment: "staging"}, + {UUID: "BOTH", Environment: "prod"}, + } + require.NoError(t, db.Create(&rows).Error) + require.NoError(t, db.Delete(&rows[2]).Error) + + var selects int + require.NoError(t, db.Callback().Query().After("gorm:query").Register("count_membership_selects", func(*gorm.DB) { selects++ })) + + got, err := repo.UUIDsInEnvironment([]string{"in-prod", "IN-STAGING", "DELETED", "BOTH", "UNKNOWN"}, "PROD") + require.NoError(t, err) + require.Equal(t, 1, selects, "membership for a whole page must be one query") + require.Equal(t, map[string]struct{}{"IN-PROD": {}, "BOTH": {}}, got, + "only live nodes enrolled in the environment, keyed by stored UUID") + + empty, err := repo.UUIDsInEnvironment(nil, "prod") + require.NoError(t, err) + require.Empty(t, empty) + require.Equal(t, 1, selects, "an empty set must not query") +} diff --git a/pkg/nodes/nodes.go b/pkg/nodes/nodes.go index c7b99120..35181c2c 100644 --- a/pkg/nodes/nodes.go +++ b/pkg/nodes/nodes.go @@ -137,6 +137,38 @@ func (n *NodeManager) GetByUUID(uuid string) (OsqueryNode, error) { return node, nil } +// UUIDsInEnvironment reports which of uuids belong to a live node enrolled in +// environment env, keyed by the stored (upper-case) UUID. The environment name +// is matched case-insensitively, as the handlers compare it elsewhere. +// +// One query for the whole set, reading only the two columns it needs, in place +// of a full-row lookup per UUID: the rows carry the raw enrollment blob, and +// callers check membership for a whole page of nodes at a time. +func (n *NodeManager) UUIDsInEnvironment(uuids []string, env string) (map[string]struct{}, error) { + out := make(map[string]struct{}, len(uuids)) + if len(uuids) == 0 { + return out, nil + } + upper := make([]string, len(uuids)) + for i, u := range uuids { + upper[i] = strings.ToUpper(u) + } + var rows []struct { + UUID string + Environment string + } + if err := n.DB.Model(&OsqueryNode{}).Select("uuid", "environment"). + Where("uuid IN ?", upper).Find(&rows).Error; err != nil { + return nil, err + } + for _, r := range rows { + if strings.EqualFold(r.Environment, env) { + out[r.UUID] = struct{}{} + } + } + return out, nil +} + // GetByUUIDEnv to retrieve full node object from DB, by uuid and environment ID // UUID is expected uppercase func (n *NodeManager) GetByUUIDEnv(uuid string, envid uint) (OsqueryNode, error) { @@ -352,48 +384,45 @@ func (n *NodeManager) GetOsqueryVersionCounts() ([]OsqueryVersionCount, error) { return GetOsqueryVersionCounts(n.DB) } -// UpdateMetadataByUUID to update node metadata by UUID -func (n *NodeManager) UpdateMetadataByUUID(uuid string, metadata NodeMetadata) error { - // Retrieve node - node, err := n.GetByUUID(uuid) - if err != nil { - return fmt.Errorf("getNodeByUUID %w", err) - } - // Prepare metadata updates +// UpdateMetadataByUUID records the metadata osquery decorates its log batches +// with on the node of that UUID enrolled in envID. +// +// It runs on every log batch every node sends, so it is a single UPDATE. +// bytes_received is incremented in SQL instead of being read and written back, +// which also stops two log batches from the same node, processed concurrently, +// from losing each other's bytes. The other columns are assigned whenever the +// batch carries a value: storing an unchanged value leaves the row as it was, +// so there is no need to read the row first to compare. +// +// Scoping by environment matters because a UUID enrolled in two environments +// has a live row in each; the caller has already authenticated the batch +// against envID, so that is the row it describes. +func (n *NodeManager) UpdateMetadataByUUID(uuid string, envID uint, metadata NodeMetadata) error { updates := map[string]interface{}{ - "bytes_received": node.BytesReceived + metadata.BytesReceived, - } - // Record username - if metadata.Username != node.Username && metadata.Username != "" { - updates["username"] = metadata.Username - } - // Record hostname - if metadata.Hostname != node.Hostname && metadata.Hostname != "" { - updates["hostname"] = metadata.Hostname - } - // Record localname - if metadata.Localname != node.Localname && metadata.Localname != "" { - updates["localname"] = metadata.Localname - } - // Record IP address - if metadata.IPAddress != node.IPAddress && metadata.IPAddress != "" { - updates["ip_address"] = metadata.IPAddress - } - // Configuration and daemon hash and osquery version update, if different - if metadata.ConfigHash != node.ConfigHash && metadata.ConfigHash != "" { - updates["config_hash"] = metadata.ConfigHash - } - if metadata.DaemonHash != node.DaemonHash && metadata.DaemonHash != "" { - updates["daemon_hash"] = metadata.DaemonHash - } - if metadata.OsqueryVersion != node.OsqueryVersion && metadata.OsqueryVersion != "" { - updates["osquery_version"] = metadata.OsqueryVersion - } - if metadata.OsqueryUser != node.OsqueryUser && metadata.OsqueryUser != "" { - updates["osquery_user"] = metadata.OsqueryUser - } - if err := n.MetadataRefresh(node, updates); err != nil { - return fmt.Errorf("MetadataRefresh %w", err) + "bytes_received": gorm.Expr("bytes_received + ?", metadata.BytesReceived), + } + for column, value := range map[string]string{ + "username": metadata.Username, + "hostname": metadata.Hostname, + "localname": metadata.Localname, + "ip_address": metadata.IPAddress, + "config_hash": metadata.ConfigHash, + "daemon_hash": metadata.DaemonHash, + "osquery_version": metadata.OsqueryVersion, + "osquery_user": metadata.OsqueryUser, + } { + if value != "" { + updates[column] = value + } + } + result := n.DB.Model(&OsqueryNode{}). + Where("uuid = ? AND environment_id = ?", strings.ToUpper(uuid), envID). + Updates(updates) + if result.Error != nil { + return fmt.Errorf("update metadata %w", result.Error) + } + if result.RowsAffected == 0 { + return fmt.Errorf("update metadata for node %s: %w", uuid, gorm.ErrRecordNotFound) } return nil } @@ -507,11 +536,6 @@ func (n *NodeManager) UpdateIP(nodeID uint, ip string) error { } -// MetadataRefresh to perform all needed update operations per node to keep metadata refreshed -func (n *NodeManager) MetadataRefresh(node OsqueryNode, updates map[string]interface{}) error { - return n.DB.Model(&node).Updates(updates).Error -} - // SortableColumns is the closed set of columns that may be ordered by external // callers. Enforced in GetByEnvPaged so the allowlist is part of the data layer, // not just the HTTP handler. Resolves audit finding U-DB-1. diff --git a/pkg/queries/queries.go b/pkg/queries/queries.go index 7b9cd48e..4c28a8c7 100644 --- a/pkg/queries/queries.go +++ b/pkg/queries/queries.go @@ -335,6 +335,45 @@ func (q *Queries) GetCarves(target string, envid uint) ([]DistributedQuery, erro return q.Gets(target, CarveQueryType, envid) } +// ActiveCounts is how many active standard queries and carves an environment +// has. +type ActiveCounts struct { + Queries int + Carves int +} + +// ActiveCountsByEnvironment counts active standard queries and carves for +// every environment in one grouped query, applying exactly the TargetActive +// filter of GetQueries and GetCarves. For callers that only need the totals: +// those methods load every matching row, query text included, to be counted. +func (q *Queries) ActiveCountsByEnvironment() (map[uint]ActiveCounts, error) { + var rows []struct { + EnvironmentID uint + Type string + Total int + } + if err := q.DB.Model(&DistributedQuery{}). + Select("environment_id, type, COUNT(*) AS total"). + Where("active = ? AND completed = ? AND deleted = ? AND expired = ? AND type IN ?", + true, false, false, false, []string{StandardQueryType, CarveQueryType}). + Group("environment_id, type"). + Find(&rows).Error; err != nil { + return nil, err + } + out := make(map[uint]ActiveCounts) + for _, r := range rows { + c := out[r.EnvironmentID] + switch r.Type { + case StandardQueryType: + c.Queries = r.Total + case CarveQueryType: + c.Carves = r.Total + } + out[r.EnvironmentID] = c + } + return out, nil +} + // Get to get a query by name func (q *Queries) Get(name string, envid uint) (DistributedQuery, error) { var query DistributedQuery @@ -592,26 +631,34 @@ func (q *Queries) UpdateQueryStatus(queryName string, nodeID uint, statusCode in return fmt.Errorf("error getting query id: %w", err) } - var nodeQuery NodeQuery - - if err := q.DB.Where("node_id = ? AND query_id = ?", nodeID, query.ID).Find(&nodeQuery).Error; err != nil { - return err + // One UPDATE rather than SELECT-then-UPDATE. No matching row means the + // node was never targeted by this query, which stays an error, as it was + // when the empty SELECT left nothing to update. + updated := q.DB.Model(&NodeQuery{}). + Where("node_id = ? AND query_id = ?", nodeID, query.ID). + Updates(map[string]interface{}{"status": result}) + if updated.Error != nil { + return updated.Error } - if err := q.DB.Model(&nodeQuery).Updates(map[string]interface{}{"status": result}).Error; err != nil { - return err + if updated.RowsAffected == 0 { + return fmt.Errorf("no node query for node %d and query %s", nodeID, queryName) } - var pending int64 + // Completion only needs to know whether any target is still pending, so + // look for one rather than counting them. status is not indexed: COUNT + // read every node_query row of the query on every result, O(targets²) + // over a query's life, while LIMIT 1 stops at the first pending row. + var stillPending []uint if err := q.DB.Model(&NodeQuery{}). Where("query_id = ? AND status = ?", query.ID, DistributedQueryStatusPending). - Count(&pending).Error; err != nil { + Limit(1).Pluck("id", &stillPending).Error; err != nil { return err } // Standard distributed queries are complete once every targeted node has // reached a terminal node_query status. Carves use the same delivery // mechanism, but the actual file transfer continues after query delivery, so // they must not be auto-completed here. - if pending == 0 && query.Type != CarveQueryType { + if len(stillPending) == 0 && query.Type != CarveQueryType { if err := q.DB.Model(&query).Updates(map[string]interface{}{"completed": true, "active": false}).Error; err != nil { return err } diff --git a/pkg/queries/queries_test.go b/pkg/queries/queries_test.go index 8cd4b0c7..df6b327f 100644 --- a/pkg/queries/queries_test.go +++ b/pkg/queries/queries_test.go @@ -1,6 +1,7 @@ package queries_test import ( + "fmt" "testing" "time" @@ -477,3 +478,129 @@ func TestSetNodeQueriesAsExpired(t *testing.T) { }) } + +// A node with no node_query row was never targeted. Reporting it is an error, +// and it must not touch the query's state or anyone else's status. +func TestUpdateQueryStatusRejectsUntargetedNode(t *testing.T) { + db := testDB(t) + q, nodes, query := setupTestData(t, db) + require.NoError(t, db.Create(&queries.NodeQuery{NodeID: nodes[0].ID, QueryID: query.ID, Status: queries.DistributedQueryStatusPending}).Error) + + err := q.UpdateQueryStatus(query.Name, nodes[2].ID, 0) + require.Error(t, err, "a node the query never targeted must be reported") + + var target queries.NodeQuery + require.NoError(t, db.Where("node_id = ? AND query_id = ?", nodes[0].ID, query.ID).First(&target).Error) + assert.Equal(t, queries.DistributedQueryStatusPending, target.Status, "the real target's status must be untouched") + var unchanged queries.DistributedQuery + require.NoError(t, db.First(&unchanged, query.ID).Error) + assert.False(t, unchanged.Completed, "an untargeted report must not complete the query") +} + +// An errored result is terminal too: the query completes once every target +// has either answered or failed. +func TestUpdateQueryStatusErrorsCountAsTerminal(t *testing.T) { + db := testDB(t) + q, nodes, query := setupTestData(t, db) + for _, n := range nodes[:2] { + require.NoError(t, db.Create(&queries.NodeQuery{NodeID: n.ID, QueryID: query.ID, Status: queries.DistributedQueryStatusPending}).Error) + } + + require.NoError(t, q.UpdateQueryStatus(query.Name, nodes[0].ID, 1)) + require.NoError(t, q.UpdateQueryStatus(query.Name, nodes[1].ID, 0)) + + var failed queries.NodeQuery + require.NoError(t, db.Where("node_id = ? AND query_id = ?", nodes[0].ID, query.ID).First(&failed).Error) + assert.Equal(t, queries.DistributedQueryStatusError, failed.Status) + var done queries.DistributedQuery + require.NoError(t, db.First(&done, query.ID).Error) + assert.True(t, done.Completed, "a query whose targets all finished, one with an error, is complete") +} + +// Completion looks for any remaining pending target rather than counting +// them. With many targets, the last pending row can sit behind every finished +// one; the query must stay open until exactly that row resolves. +func TestUpdateQueryStatusCompletesOnlyAfterLastPendingTarget(t *testing.T) { + db := testDB(t) + q := queries.CreateQueries(db) + query := &queries.DistributedQuery{ + Name: "fanout", Query: "SELECT 1;", Active: true, EnvironmentID: 1, + Expiration: time.Now().Add(time.Hour), + } + require.NoError(t, db.Create(query).Error) + + const targets = 40 + ids := make([]uint, targets) + for i := range ids { + n := nodes.OsqueryNode{UUID: fmt.Sprintf("NODE-%02d", i)} + require.NoError(t, db.Create(&n).Error) + ids[i] = n.ID + require.NoError(t, db.Create(&queries.NodeQuery{NodeID: n.ID, QueryID: query.ID, Status: queries.DistributedQueryStatusPending}).Error) + } + + // Answer in reverse insertion order, leaving the first-inserted row last. + for i := targets - 1; i >= 1; i-- { + require.NoError(t, q.UpdateQueryStatus(query.Name, ids[i], 0)) + var open queries.DistributedQuery + require.NoError(t, db.First(&open, query.ID).Error) + require.Falsef(t, open.Completed, "completed with %d target(s) still pending", i) + } + require.NoError(t, q.UpdateQueryStatus(query.Name, ids[0], 0)) + var done queries.DistributedQuery + require.NoError(t, db.First(&done, query.ID).Error) + assert.True(t, done.Completed, "query must complete when its last target answers") + assert.False(t, done.Active) +} + +// The grouped count must agree with GetQueries/GetCarves(TargetActive) for +// every environment, since StatsHandler swapped one for the other. +func TestActiveCountsByEnvironmentMatchesTargetActive(t *testing.T) { + db := testDB(t) + q := queries.CreateQueries(db) + + type spec struct { + env uint + typ string + active, completed, deleted, expired bool + hidden, softDeleted bool + } + specs := []spec{ + {env: 1, typ: queries.StandardQueryType, active: true}, + {env: 1, typ: queries.StandardQueryType, active: true, hidden: true}, // TargetActive does not filter hidden + {env: 1, typ: queries.StandardQueryType, active: true, completed: true}, + {env: 1, typ: queries.StandardQueryType, active: true, deleted: true}, + {env: 1, typ: queries.StandardQueryType, active: true, expired: true}, + {env: 1, typ: queries.StandardQueryType, active: true, softDeleted: true}, + {env: 1, typ: queries.StandardQueryType}, + {env: 1, typ: queries.CarveQueryType, active: true}, + {env: 1, typ: queries.ConsoleQueryType, active: true}, + {env: 1, typ: queries.FileExplorerQueryType, active: true}, + {env: 2, typ: queries.CarveQueryType, active: true}, + {env: 2, typ: queries.CarveQueryType, active: true}, + {env: 3, typ: queries.StandardQueryType, completed: true}, + } + for i, s := range specs { + dq := queries.DistributedQuery{ + Name: fmt.Sprintf("q%02d", i), Query: "SELECT 1;", Type: s.typ, EnvironmentID: s.env, + Active: s.active, Completed: s.completed, Deleted: s.deleted, Expired: s.expired, Hidden: s.hidden, + } + require.NoError(t, db.Create(&dq).Error) + if s.softDeleted { + require.NoError(t, db.Delete(&dq).Error) + } + } + + got, err := q.ActiveCountsByEnvironment() + require.NoError(t, err) + for _, env := range []uint{1, 2, 3, 4} { + wantQ, err := q.GetQueries(queries.TargetActive, env) + require.NoError(t, err) + wantC, err := q.GetCarves(queries.TargetActive, env) + require.NoError(t, err) + assert.Equalf(t, queries.ActiveCounts{Queries: len(wantQ), Carves: len(wantC)}, got[env], "environment %d", env) + } + assert.Equal(t, queries.ActiveCounts{Queries: 2, Carves: 1}, got[1]) + assert.Equal(t, queries.ActiveCounts{Carves: 2}, got[2]) + _, hasEnv3 := got[3] + assert.False(t, hasEnv3, "an environment with nothing active has no entry") +} diff --git a/pkg/users/users.go b/pkg/users/users.go index cbc1081b..0ace0c78 100644 --- a/pkg/users/users.go +++ b/pkg/users/users.go @@ -595,16 +595,24 @@ func (m *UserManager) UpdateMetadata(ipaddress, useragent, username, csrftoken s return nil } -// UpdateTokenIPAddress updates IP and Last Access for a user's token -func (m *UserManager) UpdateTokenIPAddress(ipaddress, username string) error { - user, err := m.Get(username) - if err != nil { - return fmt.Errorf("error getting user %w", err) +// tokenUseResolution is how stale LastTokenUse may get before a request +// refreshes it. LastTokenUse is informational — nothing authorizes on it — so +// minute resolution costs nothing visible, while recording every request had +// each SPA poll write the same admin_users row. +const tokenUseResolution = time.Minute + +// RecordTokenUse stamps the IP and time of an authenticated API request on +// user, the record the auth middleware has just loaded to validate the token, +// so no second lookup is needed. The write is skipped when the IP is unchanged +// and the stored use is within tokenUseResolution; a new IP always records. +func (m *UserManager) RecordTokenUse(user AdminUser, ipaddress string, now time.Time) error { + if user.LastIPAddress == ipaddress && now.Sub(user.LastTokenUse) < tokenUseResolution { + return nil } if err := m.DB.Model(&user).Updates( AdminUser{ LastIPAddress: ipaddress, - LastTokenUse: time.Now(), + LastTokenUse: now, }).Error; err != nil { return fmt.Errorf("update %w", err) } diff --git a/pkg/users/users_test.go b/pkg/users/users_test.go index 42c6117f..b27ed396 100644 --- a/pkg/users/users_test.go +++ b/pkg/users/users_test.go @@ -394,13 +394,10 @@ func TestUserChangePreferredLanguageClear(t *testing.T) { assert.NoError(t, err) } -func TestUpdateTokenIPAddress(t *testing.T) { +// The auth middleware already loaded the user to validate the token, so +// recording the request's use must not look the user up again. +func TestRecordTokenUseWritesWithoutLookup(t *testing.T) { manager, mock := setupTestManager(t) - mock.ExpectQuery( - regexp.QuoteMeta(`SELECT * FROM "admin_users" WHERE username = $1 AND "admin_users"."deleted_at" IS NULL ORDER BY "admin_users"."id" LIMIT $2`)). - WithArgs("testUser", 1). - WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1)) - mock.ExpectBegin() mock.ExpectExec( regexp.QuoteMeta(`UPDATE "admin_users" SET "updated_at"=$1,"last_ip_address"=$2,"last_token_use"=$3 WHERE "admin_users"."deleted_at" IS NULL AND "id" = $4`)). @@ -408,9 +405,35 @@ func TestUpdateTokenIPAddress(t *testing.T) { WillReturnResult(sqlmock.NewResult(1, 1)) mock.ExpectCommit() - err := manager.UpdateTokenIPAddress("1.2.3.4", "testUser") + user := AdminUser{Model: gorm.Model{ID: 1}, Username: "testUser", LastIPAddress: "1.2.3.4", LastTokenUse: time.Now().Add(-2 * tokenUseResolution)} + assert.NoError(t, manager.RecordTokenUse(user, "1.2.3.4", time.Now())) + assert.NoError(t, mock.ExpectationsWereMet()) +} - assert.NoError(t, err) +func TestRecordTokenUseSkipsRecentSameIP(t *testing.T) { + manager, mock := setupTestManager(t) + now := time.Now() + user := AdminUser{Model: gorm.Model{ID: 1}, Username: "testUser", LastIPAddress: "1.2.3.4", LastTokenUse: now.Add(-tokenUseResolution / 2)} + + // No expectations: any statement fails the test. + assert.NoError(t, manager.RecordTokenUse(user, "1.2.3.4", now)) + assert.NoError(t, mock.ExpectationsWereMet()) +} + +// A request from a new address is recorded immediately, however recent the +// last one was. +func TestRecordTokenUseRecordsNewIPImmediately(t *testing.T) { + manager, mock := setupTestManager(t) + mock.ExpectBegin() + mock.ExpectExec(regexp.QuoteMeta(`UPDATE "admin_users" SET`)). + WithArgs(sqlmock.AnyArg(), "5.6.7.8", sqlmock.AnyArg(), 1). + WillReturnResult(sqlmock.NewResult(1, 1)) + mock.ExpectCommit() + + now := time.Now() + user := AdminUser{Model: gorm.Model{ID: 1}, Username: "testUser", LastIPAddress: "1.2.3.4", LastTokenUse: now.Add(-time.Second)} + assert.NoError(t, manager.RecordTokenUse(user, "5.6.7.8", now)) + assert.NoError(t, mock.ExpectationsWereMet()) } func TestUpdateMetadata(t *testing.T) {