In memory auth context cache

This commit is contained in:
2024-05-17 19:37:58 +08:00
parent 4e4fbb8ba9
commit 8ae6292bf0
12 changed files with 69 additions and 120 deletions

View File

@ -44,9 +44,6 @@ func main() {
} else if err := database.RunMigration(database.C); err != nil {
log.Fatal().Err(err).Msg("An error occurred when running database auto migration.")
}
if err := database.NewBolt(); err != nil {
log.Fatal().Err(err).Msg("An error occurred when init bolt db.")
}
// External
// All the things are optional so when error occurred the server won't crash
@ -70,8 +67,8 @@ func main() {
// Configure timed tasks
quartz := cron.New(cron.WithLogger(cron.VerbosePrintfLogger(&log.Logger)))
quartz.AddFunc("@every 60m", services.DoAutoSignoff)
quartz.AddFunc("@every 60m", services.DoAutoAuthCleanup)
quartz.AddFunc("@every 60m", services.DoAutoDatabaseCleanup)
quartz.AddFunc("@every 60s", services.RecycleAuthContext)
quartz.AddFunc("@every 5m", services.KexCleanup)
quartz.Start()
@ -85,6 +82,4 @@ func main() {
log.Info().Msgf("Passport v%s is quitting...", pkg.AppVersion)
quartz.Stop()
database.B.Close()
}

View File

@ -4,7 +4,6 @@ import (
"github.com/rs/zerolog/log"
"github.com/samber/lo"
"github.com/spf13/viper"
"go.etcd.io/bbolt"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/logger"
@ -27,14 +26,3 @@ func NewGorm() error {
return err
}
var B *bbolt.DB
func NewBolt() error {
var err error
dsn := viper.GetString("database.bolt")
B, err = bbolt.Open(dsn, 0600, nil)
return err
}

View File

@ -58,7 +58,7 @@ func (v AuthTicket) IsAvailable() error {
}
type AuthContext struct {
Ticket AuthTicket `json:"ticket"`
Account Account `json:"account"`
ExpiredAt time.Time `json:"expired_at"`
Ticket AuthTicket `json:"ticket"`
Account Account `json:"account"`
LastUsedAt time.Time `json:"last_used_at"`
}

View File

@ -108,6 +108,8 @@ func editUserinfo(c *fiber.Ctx) error {
return fiber.NewError(fiber.StatusInternalServerError, err.Error())
}
services.InvalidAuthCacheWithUser(account.ID)
return c.SendStatus(fiber.StatusOK)
}

View File

@ -4,6 +4,7 @@ import (
"fmt"
"git.solsynth.dev/hydrogen/passport/pkg/database"
"git.solsynth.dev/hydrogen/passport/pkg/models"
"git.solsynth.dev/hydrogen/passport/pkg/services"
"git.solsynth.dev/hydrogen/passport/pkg/utils"
"github.com/gofiber/fiber/v2"
"github.com/nicksnyder/go-i18n/v2/i18n"
@ -92,6 +93,8 @@ func personalizeAction(c *fiber.Ctx) error {
}).Redirect("/users/me/personalize")
}
services.InvalidAuthCacheWithUser(account.ID)
return flash.WithInfo(c, fiber.Map{
"message": "your account has been personalized",
}).Redirect("/users/me")

View File

@ -116,6 +116,8 @@ func ConfirmAccount(code string) error {
return err
}
InvalidAuthCacheWithUser(user.ID)
return nil
})
}

View File

@ -4,15 +4,12 @@ import (
"fmt"
"time"
"git.solsynth.dev/hydrogen/passport/pkg/database"
"git.solsynth.dev/hydrogen/passport/pkg/models"
"github.com/gofiber/fiber/v2"
jsoniter "github.com/json-iterator/go"
"github.com/rs/zerolog/log"
"go.etcd.io/bbolt"
)
const authContextBucket = "AuthContext"
var authContextCache = make(map[string]models.AuthContext)
func Authenticate(access, refresh string, depth int) (user models.Account, perms map[string]any, newAccess, newRefresh string, err error) {
var claims PayloadClaims
@ -44,7 +41,7 @@ func Authenticate(access, refresh string, depth int) (user models.Account, perms
ctx, err = GrantAuthContext(claims.ID)
if err == nil {
log.Debug().Str("jti", claims.ID).Err(lookupErr).Msg("Missed auth context cache once!")
perms = FilterPermNodes(ctx.Account.PermNodes, ctx.Ticket.Claims)
user = ctx.Account
return
@ -58,26 +55,14 @@ func GetAuthContext(jti string) (models.AuthContext, error) {
var err error
var ctx models.AuthContext
err = database.B.View(func(tx *bbolt.Tx) error {
bucket := tx.Bucket([]byte(authContextBucket))
if bucket == nil {
return fmt.Errorf("unable to find auth context bucket")
}
raw := bucket.Get([]byte(jti))
if raw == nil {
return fmt.Errorf("unable to find auth context")
} else if err := jsoniter.Unmarshal(raw, &ctx); err != nil {
return fmt.Errorf("unable to unmarshal auth context: %v", err)
}
return nil
})
if err == nil && time.Now().Unix() >= ctx.ExpiredAt.Unix() {
_ = RevokeAuthContext(jti)
return ctx, fmt.Errorf("auth context has been expired")
if val, ok := authContextCache[jti]; ok {
ctx = val
ctx.LastUsedAt = time.Now()
authContextCache[jti] = ctx
log.Debug().Str("jti", jti).Msg("Used an auth context cache")
} else {
ctx, err = GrantAuthContext(jti)
log.Debug().Str("jti", jti).Msg("Created a new auth context cache")
}
return ctx, err
@ -99,37 +84,37 @@ func GrantAuthContext(jti string) (models.AuthContext, error) {
return ctx, fmt.Errorf("invalid account: %v", err)
}
// Every context should expire in some while
// Once user update their account info, this will have delay to update
ctx = models.AuthContext{
Ticket: ticket,
Account: user,
ExpiredAt: time.Now().Add(5 * time.Minute),
Ticket: ticket,
Account: user,
LastUsedAt: time.Now(),
}
// Save data into KV cache
return ctx, database.B.Update(func(tx *bbolt.Tx) error {
bucket, err := tx.CreateBucketIfNotExists([]byte(authContextBucket))
if err != nil {
return err
}
// Put the data into memory for cache
authContextCache[jti] = ctx
raw, err := jsoniter.Marshal(ctx)
if err != nil {
return err
}
return bucket.Put([]byte(jti), raw)
})
return ctx, nil
}
func RevokeAuthContext(jti string) error {
return database.B.Update(func(tx *bbolt.Tx) error {
bucket, err := tx.CreateBucketIfNotExists([]byte(authContextBucket))
if err != nil {
return err
}
func RecycleAuthContext() {
if len(authContextCache) == 0 {
return
}
return bucket.Delete([]byte(jti))
})
affected := 0
for key, val := range authContextCache {
if val.LastUsedAt.Add(60*time.Second).Unix() < time.Now().Unix() {
affected++
delete(authContextCache, key)
}
}
log.Debug().Int("affected", affected).Msg("Recycled auth context...")
}
func InvalidAuthCacheWithUser(userId uint) {
for key, val := range authContextCache {
if val.Account.ID == userId {
delete(authContextCache, key)
}
}
}

View File

@ -3,10 +3,8 @@ package services
import (
"git.solsynth.dev/hydrogen/passport/pkg/database"
"git.solsynth.dev/hydrogen/passport/pkg/models"
jsoniter "github.com/json-iterator/go"
"github.com/rs/zerolog/log"
"github.com/spf13/viper"
"go.etcd.io/bbolt"
"time"
)
@ -24,36 +22,3 @@ func DoAutoSignoff() {
log.Debug().Int64("affected", tx.RowsAffected).Msg("Auto sign off accomplished.")
}
}
func DoAutoAuthCleanup() {
log.Debug().Msg("Now 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.")
}
}