From b3f00f4efba7efeb6645872a98e97587622cb300 Mon Sep 17 00:00:00 2001 From: Kian Parvin Date: Wed, 8 Jan 2025 17:57:07 +0200 Subject: [PATCH] chore: fix cleanup audit log test --- internal/jimm/auditlog/auditlog.go | 19 +++++++- internal/jimm/auditlog/auditlog_test.go | 2 +- internal/jimm/auditlog/cleanup.go | 50 +++++++------------ internal/jimm/auditlog/cleanup_test.go | 64 +++++++++++++++---------- internal/jimm/auditlog/export_test.go | 7 ++- internal/jimm/jimm.go | 2 +- 6 files changed, 82 insertions(+), 62 deletions(-) diff --git a/internal/jimm/auditlog/auditlog.go b/internal/jimm/auditlog/auditlog.go index dce721b7a..551cd3033 100644 --- a/internal/jimm/auditlog/auditlog.go +++ b/internal/jimm/auditlog/auditlog.go @@ -19,17 +19,25 @@ import ( ofganames "github.com/canonical/jimm/v3/internal/openfga/names" ) +// PollTimeOfDay holds the time hour, minutes and seconds to poll for cleanup. +type PollTimeOfDay struct { + Hours int + Minutes int + Seconds int +} + // auditLogManager provides a means to manage audit logs within JIMM. type auditLogManager struct { store *db.Database authSvc *openfga.OFGAClient jimmTag names.ControllerTag retentionPeriodInDays int + pollTimeOfDay PollTimeOfDay } // NewAuditLogManager returns a new auditLog manager that provides audit Log // creation, and removal. -func NewAuditLogManager(store *db.Database, authSvc *openfga.OFGAClient, jimmTag names.ControllerTag, retentionDays int) (*auditLogManager, error) { +func NewAuditLogManager(store *db.Database, authSvc *openfga.OFGAClient, jimmTag names.ControllerTag, retentionDays int, pollTime *PollTimeOfDay) (*auditLogManager, error) { if store == nil { return nil, errors.E("auditlog store cannot be nil") } @@ -39,7 +47,14 @@ func NewAuditLogManager(store *db.Database, authSvc *openfga.OFGAClient, jimmTag if jimmTag.String() == "" { return nil, errors.E("auditlog jimm tag cannot be empty") } - return &auditLogManager{store, authSvc, jimmTag, retentionDays}, nil + // By default we poll at 9 AM. + defaultPollTime := PollTimeOfDay{ + Hours: 9, + } + if pollTime == nil { + pollTime = &defaultPollTime + } + return &auditLogManager{store, authSvc, jimmTag, retentionDays, *pollTime}, nil } // addAuditLogEntry causes an entry to be added the the audit log. diff --git a/internal/jimm/auditlog/auditlog_test.go b/internal/jimm/auditlog/auditlog_test.go index a9b5aec97..230d98189 100644 --- a/internal/jimm/auditlog/auditlog_test.go +++ b/internal/jimm/auditlog/auditlog_test.go @@ -45,7 +45,7 @@ func (s *auditLogManagerSuite) Init(c *qt.C) { s.jimmTag = names.NewControllerTag("foo") - s.manager, err = auditlog.NewAuditLogManager(db, ofgaClient, s.jimmTag, 1) + s.manager, err = auditlog.NewAuditLogManager(db, ofgaClient, s.jimmTag, 1, nil) c.Assert(err, qt.IsNil) // Create test identity diff --git a/internal/jimm/auditlog/cleanup.go b/internal/jimm/auditlog/cleanup.go index ddd83e9b2..75b8e6820 100644 --- a/internal/jimm/auditlog/cleanup.go +++ b/internal/jimm/auditlog/cleanup.go @@ -9,41 +9,18 @@ import ( "go.uber.org/zap" ) -// pollTimeOfDay holds the time hour, minutes and seconds to poll at. -type pollTimeOfDay struct { - Hours int - Minutes int - Seconds int -} - -var pollDuration = pollTimeOfDay{ - Hours: 9, -} - -// StartCleanup starts a routine which checks daily for any logs -// needed to be cleaned up. +// StartCleanup loop forever and checks daily for any logs +// that need to be cleaned up. This method should be run +// in a separate Go routine to avoid blocking, it will terminate +// when the provided context is cancelled. func (j *auditLogManager) StartCleanup(ctx context.Context) { if j.retentionPeriodInDays == 0 { return } - go j.poll(ctx) -} - -// poll is designed to be run in a routine where it can be cancelled safely -// from the service's context. It calculates the poll duration at 9am each day -// UTC. -func (j *auditLogManager) poll(ctx context.Context) { - for { select { - case <-time.After(calculateNextPollDuration(time.Now().UTC())): - retentionDate := time.Now().AddDate(0, 0, -(j.retentionPeriodInDays)) - deleted, err := j.store.DeleteAuditLogsBefore(ctx, retentionDate) - if err != nil { - zapctx.Error(ctx, "failed to cleanup audit logs", zap.Error(err)) - continue - } - zapctx.Debug(ctx, "audit log cleanup run successfully", zap.Int64("count", deleted)) + case <-time.After(calculateNextPollDuration(j.pollTimeOfDay, time.Now().UTC())): + j.cleanup(ctx) case <-ctx.Done(): zapctx.Debug(ctx, "exiting audit log cleanup polling") return @@ -51,13 +28,22 @@ func (j *auditLogManager) poll(ctx context.Context) { } } +func (j *auditLogManager) cleanup(ctx context.Context) { + retentionDate := time.Now().AddDate(0, 0, -(j.retentionPeriodInDays)) + deleted, err := j.store.DeleteAuditLogsBefore(ctx, retentionDate) + if err != nil { + zapctx.Error(ctx, "failed to cleanup audit logs", zap.Error(err)) + } + zapctx.Debug(ctx, "audit log cleanup run successfully", zap.Int64("count", deleted)) +} + // calculateNextPollDuration returns the next duration to poll on. // We recalculate each time and not rely on running every 24 hours // for absolute consistency within ns apart. -func calculateNextPollDuration(startingTime time.Time) time.Duration { +func calculateNextPollDuration(pollTime PollTimeOfDay, startingTime time.Time) time.Duration { now := startingTime - pollTime := time.Date(now.Year(), now.Month(), now.Day(), pollDuration.Hours, pollDuration.Minutes, pollDuration.Seconds, 0, time.UTC) - tillNextPoll := pollTime.Sub(now) + pollTimeToday := time.Date(now.Year(), now.Month(), now.Day(), pollTime.Hours, pollTime.Minutes, pollTime.Seconds, 0, time.UTC) + tillNextPoll := pollTimeToday.Sub(now) var d time.Duration // If the next poll time is behind the current time if tillNextPoll < 0 { diff --git a/internal/jimm/auditlog/cleanup_test.go b/internal/jimm/auditlog/cleanup_test.go index f8441aaed..bb08ad532 100644 --- a/internal/jimm/auditlog/cleanup_test.go +++ b/internal/jimm/auditlog/cleanup_test.go @@ -1,4 +1,5 @@ // Copyright 2025 Canonical. + package auditlog_test import ( @@ -7,66 +8,79 @@ import ( "time" qt "github.com/frankban/quicktest" + "github.com/juju/names/v5" + "github.com/canonical/jimm/v3/internal/db" "github.com/canonical/jimm/v3/internal/dbmodel" - "github.com/canonical/jimm/v3/internal/errors" "github.com/canonical/jimm/v3/internal/jimm/auditlog" + "github.com/canonical/jimm/v3/internal/testutils/jimmtest" ) -func (s *auditLogManagerSuite) TestAuditLogCleanupServicePurgesLogs(c *qt.C) { +func TestAuditLogCleanupServicePurgesLogs(t *testing.T) { + c := qt.New(t) c.Parallel() + ctx := context.Background() - now := time.Now().UTC().Round(time.Millisecond) - err := s.db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{ - Time: now.AddDate(0, 0, -1), - }) - c.Check(errors.ErrorCode(err), qt.Equals, errors.CodeUpgradeInProgress) + db := &db.Database{ + DB: jimmtest.PostgresDB(c, time.Now), + } + err := db.Migrate(context.Background()) + c.Assert(err, qt.IsNil) + + ofgaClient, _, _, err := jimmtest.SetupTestOFGAClient(c.Name()) + c.Assert(err, qt.IsNil) + + jimmTag := names.NewControllerTag("foo") + + manager, err := auditlog.NewAuditLogManager(db, ofgaClient, jimmTag, 1, nil) + c.Assert(err, qt.IsNil) + + now := time.Now().UTC() + + // A log from today + c.Assert(db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{ + Time: now.AddDate(0, 0, 0), + }), qt.IsNil) // A log from 1 day ago - c.Assert(s.db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{ + c.Assert(db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{ Time: now.AddDate(0, 0, -1), }), qt.IsNil) // A log from 2 days ago - c.Assert(s.db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{ + c.Assert(db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{ Time: now.AddDate(0, 0, -2), }), qt.IsNil) - // A log from 3 days ago - c.Assert(s.db.AddAuditLogEntry(ctx, &dbmodel.AuditLogEntry{ - Time: now.AddDate(0, 0, -3), - }), qt.IsNil) - // Check 3 created logs := make([]dbmodel.AuditLogEntry, 0) - err = s.db.DB.Find(&logs).Error + err = db.DB.Find(&logs).Error c.Assert(err, qt.IsNil) c.Assert(logs, qt.HasLen, 3) - auditlog.PollDuration.Hours = now.Hour() - auditlog.PollDuration.Minutes = now.Minute() - auditlog.PollDuration.Seconds = now.Second() + 2 - // Suite is setup to clean logs more than 1 day old. - s.manager.StartCleanup(ctx) + // Manager is setup above to remove logs older than 1 day. + manager.Cleanup(ctx) // Check 2 were purged logs = make([]dbmodel.AuditLogEntry, 0) - err = s.db.DB.Find(&logs).Error + err = db.DB.Find(&logs).Error c.Assert(err, qt.IsNil) - c.Assert(logs, qt.HasLen, 3) + c.Assert(logs, qt.HasLen, 1) } func TestCalculateNextPollDuration(t *testing.T) { c := qt.New(t) + pollTime := auditlog.PollTimeOfDay{Hours: 9} + // Test where 9am is behind 12pm startingTime := time.Date(2023, 1, 1, 12, 0, 0, 0, time.UTC) - d := auditlog.CalculateNextPollDuration(startingTime) + d := auditlog.CalculateNextPollDuration(pollTime, startingTime) c.Assert(d, qt.Equals, time.Hour*21) - // Test where 9am is ahead of 7pm + // Test where 9am is ahead of 7am startingTime = time.Date(2023, 1, 1, 7, 0, 0, 0, time.UTC) - d = auditlog.CalculateNextPollDuration(startingTime) + d = auditlog.CalculateNextPollDuration(pollTime, startingTime) c.Assert(d, qt.Equals, time.Hour*2) } diff --git a/internal/jimm/auditlog/export_test.go b/internal/jimm/auditlog/export_test.go index c7413a44f..c7ef5ddbe 100644 --- a/internal/jimm/auditlog/export_test.go +++ b/internal/jimm/auditlog/export_test.go @@ -2,10 +2,15 @@ package auditlog +import "context" + // AuditLogManager is a type alias to export auditLogManager for use in tests. type AuditLogManager = auditLogManager var ( - PollDuration = pollDuration CalculateNextPollDuration = calculateNextPollDuration ) + +func (j *auditLogManager) Cleanup(ctx context.Context) { + j.cleanup(ctx) +} diff --git a/internal/jimm/jimm.go b/internal/jimm/jimm.go index 35094d6c3..3a921adde 100644 --- a/internal/jimm/jimm.go +++ b/internal/jimm/jimm.go @@ -300,7 +300,7 @@ func New(p Parameters) (*JIMM, error) { } j.loginManager = loginManager - auditLogManager, err := auditlog.NewAuditLogManager(j.Database, j.OpenFGAClient, j.ResourceTag(), p.AuditLogRetentionDays) + auditLogManager, err := auditlog.NewAuditLogManager(j.Database, j.OpenFGAClient, j.ResourceTag(), p.AuditLogRetentionDays, nil) if err != nil { return nil, err }