From f88fdd7fe56127632db48420a5b6e17053432133 Mon Sep 17 00:00:00 2001 From: LittleSheep Date: Sat, 23 Mar 2024 13:04:25 +0800 Subject: [PATCH] :zap: Auto cleanning up auth context cache --- pkg/cmd/main.go | 10 ++------ pkg/services/sessions.go | 53 +++++++++++++++++++++++++++++++++++++--- 2 files changed, 51 insertions(+), 12 deletions(-) diff --git a/pkg/cmd/main.go b/pkg/cmd/main.go index d317fa7..d3d08f4 100644 --- a/pkg/cmd/main.go +++ b/pkg/cmd/main.go @@ -66,14 +66,8 @@ func main() { // Configure timed tasks quartz := cron.New(cron.WithLogger(cron.VerbosePrintfLogger(&log.Logger))) - quartz.AddFunc("@every 60m", func() { - log.Info().Msg("Running auto sign off...") - if tx := services.PerformAutoSignoff(); tx.Error != nil { - log.Error().Err(tx.Error).Msg("An error occurred when running auto sign off...") - } else { - log.Info().Int64("affected", tx.RowsAffected).Msg("Auto sign off accomplished.") - } - }) + quartz.AddFunc("@every 60m", services.DoAutoSignoff) + quartz.AddFunc("@every 60m", services.DoAutoAuthCleanup) quartz.Run() // Messages diff --git a/pkg/services/sessions.go b/pkg/services/sessions.go index 8d366f7..5795887 100644 --- a/pkg/services/sessions.go +++ b/pkg/services/sessions.go @@ -5,8 +5,10 @@ import ( "git.solsynth.dev/hydrogen/identity/pkg/database" "git.solsynth.dev/hydrogen/identity/pkg/models" + jsoniter "github.com/json-iterator/go" + "github.com/rs/zerolog/log" "github.com/spf13/viper" - "gorm.io/gorm" + "go.etcd.io/bbolt" ) func LookupSessionWithToken(tokenId string) (models.AuthSession, error) { @@ -21,7 +23,50 @@ func LookupSessionWithToken(tokenId string) (models.AuthSession, error) { return session, nil } -func PerformAutoSignoff() *gorm.DB { - signoffDuration := time.Duration(viper.GetInt64("security.auto_signoff_duration")) * time.Second - return database.C.Where("last_grant_at < ?", time.Now().Add(-signoffDuration)).Delete(&models.AuthSession{}) +func DoAutoSignoff() { + duration := time.Duration(viper.GetInt64("security.auto_signoff_duration")) * time.Second + divider := time.Now().Add(-duration) + + log.Debug().Time("before", divider).Msg("Now auto signing off sessions...") + + if tx := database.C. + Where("last_grant_at < ?", divider). + Delete(&models.AuthSession{}); tx.Error != nil { + log.Error().Err(tx.Error).Msg("An error occurred when running auto sign off...") + } else { + log.Debug().Int64("affected", tx.RowsAffected).Msg("Auto sign off accomplished.") + } +} + +func DoAutoAuthCleanup() { + log.Debug().Msg("Now auto cleaning up cached auth context...") + + count := 0 + err := database.B.Batch(func(tx *bbolt.Tx) error { + bucket := tx.Bucket([]byte(authContextBucket)) + if bucket == nil { + return nil + } + + cursor := bucket.Cursor() + + var ctx models.AuthContext + for key, val := cursor.First(); key != nil; key, val = cursor.Next() { + if err := jsoniter.Unmarshal(val, &ctx); err != nil { + bucket.Delete(key) + count++ + } else if time.Now().Unix() >= ctx.ExpiredAt.Unix() { + bucket.Delete(key) + count++ + } + } + + return nil + }) + + if err != nil { + log.Error().Err(err).Msg("An error occurred when running auth context cleanup...") + } else { + log.Debug().Int("affected", count).Msg("Clean up auth context accomplished.") + } }