From f88fdd7fe56127632db48420a5b6e17053432133 Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
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.")
+	}
 }