Rogeeandrogee
f719529d66
* fix(security): harden auth and credential handling (HH-444) * fix(security): address HH-444 review blockers * fix(security): close remaining HH-444 review blockers --------- Co-authored-by: Rogee <rogee@ipao.vip>
104 lines
2.5 KiB
Go
104 lines
2.5 KiB
Go
package auth
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
var ErrInvalidWSTicket = errors.New("invalid or expired websocket ticket")
|
|
|
|
type WSTicketClaims struct {
|
|
UserID uint `json:"user_id"`
|
|
AccountID uint `json:"account_id"`
|
|
Role string `json:"role"`
|
|
Provider string `json:"provider"`
|
|
ClientID string `json:"client_id,omitempty"`
|
|
}
|
|
|
|
type wsTicketEntry struct {
|
|
claims WSTicketClaims
|
|
expiresAt time.Time
|
|
}
|
|
|
|
// WSTicketStore exchanges a short-lived opaque ticket for user claims exactly once.
|
|
type WSTicketStore struct {
|
|
rdb *redis.Client
|
|
ttl time.Duration
|
|
mu sync.Mutex
|
|
mem map[string]wsTicketEntry
|
|
}
|
|
|
|
func NewWSTicketStore(rdb *redis.Client, ttl time.Duration) *WSTicketStore {
|
|
if ttl <= 0 {
|
|
ttl = 30 * time.Second
|
|
}
|
|
return &WSTicketStore{rdb: rdb, ttl: ttl, mem: make(map[string]wsTicketEntry)}
|
|
}
|
|
|
|
func (s *WSTicketStore) Issue(ctx context.Context, claims WSTicketClaims) (string, error) {
|
|
random := make([]byte, 32)
|
|
if _, err := rand.Read(random); err != nil {
|
|
return "", fmt.Errorf("generate websocket ticket: %w", err)
|
|
}
|
|
ticket := base64.RawURLEncoding.EncodeToString(random)
|
|
key := wsTicketKey(ticket)
|
|
if s.rdb == nil {
|
|
s.mu.Lock()
|
|
s.mem[key] = wsTicketEntry{claims: claims, expiresAt: time.Now().Add(s.ttl)}
|
|
s.mu.Unlock()
|
|
return ticket, nil
|
|
}
|
|
payload, err := json.Marshal(claims)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if err := s.rdb.Set(ctx, key, payload, s.ttl).Err(); err != nil {
|
|
return "", fmt.Errorf("store websocket ticket: %w", err)
|
|
}
|
|
return ticket, nil
|
|
}
|
|
|
|
func (s *WSTicketStore) Consume(ctx context.Context, ticket string) (*WSTicketClaims, error) {
|
|
if ticket == "" {
|
|
return nil, ErrInvalidWSTicket
|
|
}
|
|
key := wsTicketKey(ticket)
|
|
if s.rdb == nil {
|
|
s.mu.Lock()
|
|
entry, ok := s.mem[key]
|
|
delete(s.mem, key)
|
|
s.mu.Unlock()
|
|
if !ok || time.Now().After(entry.expiresAt) {
|
|
return nil, ErrInvalidWSTicket
|
|
}
|
|
return &entry.claims, nil
|
|
}
|
|
payload, err := s.rdb.GetDel(ctx, key).Bytes()
|
|
if errors.Is(err, redis.Nil) {
|
|
return nil, ErrInvalidWSTicket
|
|
}
|
|
if err != nil {
|
|
return nil, fmt.Errorf("consume websocket ticket: %w", err)
|
|
}
|
|
var claims WSTicketClaims
|
|
if err := json.Unmarshal(payload, &claims); err != nil {
|
|
return nil, ErrInvalidWSTicket
|
|
}
|
|
return &claims, nil
|
|
}
|
|
|
|
func wsTicketKey(ticket string) string {
|
|
digest := sha256.Sum256([]byte(ticket))
|
|
return "gochat:ws_ticket:" + hex.EncodeToString(digest[:])
|
|
}
|