olons_server/internal/auth/auth.go

package auth

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"net/http"
	"strconv"
	"strings"
	"time"

	"git.118im.com/im_team/im_workbench/internal/config"
	"github.com/redis/go-redis/v9"
)

var (
	ErrAuthRequired = errors.New("authentication required")
	ErrForbidden    = errors.New("forbidden")
)

type Principal struct {
	TenantID int64
	MemberID int64
	DeviceID string
}

type Session struct {
	MemberID  int64  `json:"member_id"`
	TenantID  int64  `json:"tenant_id"`
	DeviceID  string `json:"device_id"`
	State     string `json:"state"`
	ExpiresAt string `json:"expires_at"`
	IssuedAt  string `json:"issued_at"`
}

type Authenticator struct {
	cfg   config.AuthConfig
	redis *redis.Client
}

func New(cfg config.AuthConfig, redisClient *redis.Client) *Authenticator {
	return &Authenticator{cfg: cfg, redis: redisClient}
}

func (a *Authenticator) DefaultTenantID() int64 {
	return a.cfg.DefaultTenantID
}

func (a *Authenticator) Authenticate(req *http.Request) (*Principal, error) {
	if err := a.ValidateAPIKey(req); err != nil {
		return nil, err
	}
	if a.cfg.AllowDevHeaderAuth {
		if memberID := parseInt64Header(req, "X-Member-ID"); memberID > 0 {
			return &Principal{TenantID: a.cfg.DefaultTenantID, MemberID: memberID,
				DeviceID: strings.TrimSpace(req.Header.Get("X-Device-ID"))}, nil
		}
	}
	if a.redis == nil {
		return nil, ErrAuthRequired
	}
	token, err := bearerToken(req)
	if err != nil {
		return nil, err
	}
	session, err := a.sessionFromRedis(req.Context(), token)
	if err != nil {
		return nil, err
	}
	if session.MemberID <= 0 || session.TenantID < 0 || session.State != "active" {
		return nil, ErrAuthRequired
	}
	expiresAt, err := time.Parse(time.RFC3339, session.ExpiresAt)
	if err != nil || time.Now().After(expiresAt) {
		return nil, ErrAuthRequired
	}
	if err := a.validateMemberSession(req.Context(), session); err != nil {
		return nil, err
	}
	return &Principal{TenantID: session.TenantID, MemberID: session.MemberID, DeviceID: session.DeviceID}, nil
}

func (a *Authenticator) Logout(req *http.Request) error {
	if _, err := a.Authenticate(req); err != nil {
		return err
	}
	token, err := bearerToken(req)
	if err != nil {
		return err
	}
	key, err := sessionKeyFromToken(a.cfg.SessionKeyPattern, token)
	if err != nil {
		return err
	}
	if a.redis == nil {
		return ErrAuthRequired
	}
	return a.redis.Del(req.Context(), key).Err()
}

func (a *Authenticator) RevokeMemberSessions(ctx context.Context, tenantID, memberID int64) error {
	if a == nil || a.redis == nil {
		return ErrAuthRequired
	}
	key := memberRevokedBeforeKey(tenantID, memberID)
	value := time.Now().UTC().Format(time.RFC3339Nano)
	ttl := time.Duration(a.cfg.SessionTTLSeconds) * time.Second
	return a.redis.Set(ctx, key, value, ttl).Err()
}

func (a *Authenticator) ValidateAPIKey(req *http.Request) error {
	if len(a.cfg.APIKeys) == 0 && !a.cfg.RequiredAPIKey {
		return nil
	}
	got := strings.TrimSpace(req.Header.Get("X-API-Key"))
	if got == "" {
		got = strings.TrimSpace(req.URL.Query().Get("apikey"))
	}
	if got == "" {
		return ErrForbidden
	}
	if len(a.cfg.APIKeys) == 0 {
		return ErrForbidden
	}
	for _, key := range a.cfg.APIKeys {
		if subtle.ConstantTimeCompare([]byte(got), []byte(strings.TrimSpace(key))) == 1 {
			return nil
		}
	}
	return ErrForbidden
}

func (a *Authenticator) sessionFromRedis(ctx context.Context, token string) (*Session, error) {
	key, err := sessionKeyFromToken(a.cfg.SessionKeyPattern, token)
	if err != nil {
		return nil, ErrAuthRequired
	}
	data, err := a.redis.Get(ctx, key).Bytes()
	if errors.Is(err, redis.Nil) {
		return nil, ErrAuthRequired
	}
	if err != nil {
		return nil, err
	}
	var session Session
	if err := json.Unmarshal(data, &session); err != nil {
		return nil, ErrAuthRequired
	}
	return &session, nil
}

func (a *Authenticator) validateMemberSession(ctx context.Context, session *Session) error {
	revokedBefore, err := a.redis.Get(ctx, memberRevokedBeforeKey(session.TenantID, session.MemberID)).Result()
	if errors.Is(err, redis.Nil) {
		return nil
	}
	if err != nil {
		return err
	}
	issuedAt, issuedErr := time.Parse(time.RFC3339Nano, session.IssuedAt)
	revokedAt, revokedErr := time.Parse(time.RFC3339Nano, revokedBefore)
	if issuedErr != nil || revokedErr != nil || !issuedAt.After(revokedAt) {
		return ErrAuthRequired
	}
	return nil
}

func memberRevokedBeforeKey(tenantID, memberID int64) string {
	return "workbench:session:member-revoked-before:" + strconv.FormatInt(tenantID, 10) + ":" + strconv.FormatInt(memberID, 10)
}

func sessionKeyFromToken(pattern, token string) (string, error) {
	raw, err := base64.RawURLEncoding.DecodeString(strings.TrimSpace(token))
	if err != nil || len(raw) != 32 {
		return "", ErrAuthRequired
	}
	digest := sha256.Sum256(raw)
	return strings.Replace(pattern, "%s", hex.EncodeToString(digest[:]), 1), nil
}

func bearerToken(req *http.Request) (string, error) {
	authorization := strings.TrimSpace(req.Header.Get("Authorization"))
	token, ok := strings.CutPrefix(authorization, "Bearer ")
	if !ok || strings.TrimSpace(token) == "" {
		return "", ErrAuthRequired
	}
	return strings.TrimSpace(token), nil
}

func parseInt64Header(req *http.Request, key string) int64 {
	value := strings.TrimSpace(req.Header.Get(key))
	parsed, _ := strconv.ParseInt(value, 10, 64)
	return parsed
}