sundynix-pets/pets-be/internal/httpapi/refresh.go

package httpapi

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"github.com/sundynix/pets-be/internal/model"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

var errRefreshInvalid = errors.New("refresh credential invalid")

func refreshCredential() (string, error) {
	data := make([]byte, 32)
	if _, err := rand.Read(data); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(data), nil
}
func refreshHash(value string) string {
	hash := sha256.Sum256([]byte(value))
	return hex.EncodeToString(hash[:])
}
func (a *API) sessionPair(id, role, sessionID, refresh string, expires time.Time) (gin.H, error) {
	token, err := a.sessionToken(id, role, sessionID)
	if err != nil {
		return nil, err
	}
	return gin.H{"token": token, "refresh_token": refresh, "expires_in": 7200, "expires_at": time.Now().Add(2 * time.Hour).Unix(), "refresh_expires_at": expires.Unix()}, nil
}
func (a *API) newSession(id, role string) (gin.H, error) {
	raw, err := refreshCredential()
	if err != nil {
		return nil, err
	}
	expires := time.Now().Add(2 * time.Hour)
	sessionID := uuid.NewString()
	pair, err := a.sessionPair(id, role, sessionID, raw, expires)
	if err != nil {
		return nil, err
	}
	// Opportunistic cleanup bounded by indexed expiry, without a background scheduler.
	if err := a.db.Where("expires_at <= ?", time.Now()).Delete(&model.RefreshSession{}).Error; err != nil {
		return nil, err
	}
	err = a.db.Create(&model.RefreshSession{Base: model.Base{ID: sessionID}, ActorID: id, Role: role, TokenHash: refreshHash(raw), ExpiresAt: expires}).Error
	return pair, err
}
func readRefresh(c *gin.Context) (string, bool) {
	var in struct {
		Token string `json:"refresh_token" binding:"required"`
	}
	if c.ShouldBindJSON(&in) != nil || len(in.Token) != 43 {
		fail(c, http.StatusUnauthorized, "登录已过期，请重新登录")
		return "", false
	}
	return in.Token, true
}
func (a *API) refreshToken(role string) gin.HandlerFunc {
	return func(c *gin.Context) {
		raw, valid := readRefresh(c)
		if !valid {
			return
		}
		var pair gin.H
		err := a.db.Transaction(func(tx *gorm.DB) error {
			var session model.RefreshSession
			err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("token_hash = ? AND role = ?", refreshHash(raw), role).First(&session).Error
			if errors.Is(err, gorm.ErrRecordNotFound) {
				return errRefreshInvalid
			}
			if err != nil {
				return err
			}
			if session.Revoked || !session.ExpiresAt.After(time.Now()) {
				return errRefreshInvalid
			}
			if role == "user" {
				var user model.User
				err = tx.Select("id", "disabled").First(&user, "id = ?", session.ActorID).Error
				if errors.Is(err, gorm.ErrRecordNotFound) || (err == nil && user.Disabled) {
					return errRefreshInvalid
				}
				if err != nil {
					return err
				}
			} else {
				var count int64
				if err := tx.Model(&model.Admin{}).Where("id = ?", session.ActorID).Count(&count).Error; err != nil {
					return err
				}
				if count != 1 {
					return errRefreshInvalid
				}
			}
			next, err := refreshCredential()
			if err != nil {
				return err
			}
			nextExpiry := time.Now().Add(2 * time.Hour)
			pair, err = a.sessionPair(session.ActorID, role, session.ID, next, nextExpiry)
			if err != nil {
				return err
			}
			// Rotate inside the lock. The old credential can only succeed once.
			return tx.Model(&session).Updates(map[string]any{"token_hash": refreshHash(next), "expires_at": nextExpiry}).Error
		})
		if errors.Is(err, errRefreshInvalid) {
			fail(c, 401, "登录已过期，请重新登录")
			return
		}
		if err != nil {
			fail(c, 500, "暂时无法刷新登录，请重试")
			return
		}
		ok(c, pair)
	}
}
func (a *API) revokeRefresh(role string) gin.HandlerFunc {
	return func(c *gin.Context) {
		raw, valid := readRefresh(c)
		if !valid {
			return
		}
		if err := a.db.Model(&model.RefreshSession{}).Where("token_hash = ? AND role = ?", refreshHash(raw), role).Update("revoked", true).Error; err != nil {
			fail(c, 500, "退出登录失败")
			return
		}
		ok(c, gin.H{"logged_out": true})
	}
}