gochat/backend/internal/auth/session.go

package auth

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"sync"
	"time"

	"github.com/gochat/gochat/internal/config"
	applogger "github.com/gochat/gochat/pkg/logger"
)

// Session represents a user session stored in the session store.
type Session struct {
	ID        string                 `json:"id"`
	UserID    uint                   `json:"user_id"`
	AccountID uint                   `json:"account_id"`
	Role      string                 `json:"role"`
	Provider  string                 `json:"provider"`
	CreatedAt time.Time              `json:"created_at"`
	ExpiresAt time.Time              `json:"expires_at"`
	Data      map[string]interface{} `json:"data"` // arbitrary session metadata
}

// SessionStore provides session creation, retrieval, and deletion.
// Production: Redis-backed with in-memory fallback (same pattern as RefreshTokenStore).
type SessionStore struct {
	cfg   *config.SessionConfig
	mu    sync.RWMutex
	store map[string]*Session // in-memory fallback
}

// NewSessionStore creates a session store with configuration.
func NewSessionStore(cfg *config.SessionConfig) *SessionStore {
	return &SessionStore{
		cfg:   cfg,
		store: make(map[string]*Session),
	}
}

// Create generates a new session for a user and stores it.
func (s *SessionStore) Create(userID, accountID uint, role, provider string) (*Session, error) {
	id, err := generateSessionID(s.cfg.TokenLength)
	if err != nil {
		return nil, fmt.Errorf("failed to generate session ID: %w", err)
	}

	now := time.Now()
	expiry := now.Add(time.Duration(s.cfg.ExpirySeconds) * time.Second)

	session := &Session{
		ID:        id,
		UserID:    userID,
		AccountID: accountID,
		Role:      role,
		Provider:  provider,
		CreatedAt: now,
		ExpiresAt: expiry,
		Data:      make(map[string]interface{}),
	}

	s.mu.Lock()
	s.store[id] = session
	s.mu.Unlock()

	applogger.L().Debugf("Session created: id=%s user_id=%d expires=%s", id, userID, expiry)
	return session, nil
}

// Get retrieves a session by ID. Returns nil if expired or not found.
func (s *SessionStore) Get(id string) (*Session, error) {
	s.mu.RLock()
	session, ok := s.store[id]
	s.mu.RUnlock()

	if !ok {
		return nil, fmt.Errorf("session not found: %s", id)
	}

	if time.Now().After(session.ExpiresAt) {
		if err := s.Delete(id); err != nil {
			return nil, fmt.Errorf("cleanup expired session: %w", err)
		}
		return nil, fmt.Errorf("session expired: %s", id)
	}

	return session, nil
}

// Delete removes a session by ID (logout / explicit termination).
func (s *SessionStore) Delete(id string) error {
	s.mu.Lock()
	delete(s.store, id)
	s.mu.Unlock()
	return nil
}

// DeleteByUserID removes all sessions for a given user (force logout).
func (s *SessionStore) DeleteByUserID(userID uint) int {
	s.mu.Lock()
	count := 0
	for id, session := range s.store {
		if session.UserID == userID {
			delete(s.store, id)
			count++
		}
	}
	s.mu.Unlock()
	applogger.L().Debugf("Deleted %d sessions for user_id=%d", count, userID)
	return count
}

// Refresh extends a session's expiry time (session rotation / keep-alive).
func (s *SessionStore) Refresh(id string) (*Session, error) {
	s.mu.Lock()
	session, ok := s.store[id]
	s.mu.Unlock()

	if !ok {
		return nil, fmt.Errorf("session not found: %s", id)
	}

	if time.Now().After(session.ExpiresAt) {
		if err := s.Delete(id); err != nil {
			return nil, fmt.Errorf("cleanup expired session: %w", err)
		}
		return nil, fmt.Errorf("session expired: %s", id)
	}

	newExpiry := time.Now().Add(time.Duration(s.cfg.ExpirySeconds) * time.Second)
	session.ExpiresAt = newExpiry

	applogger.L().Debugf("Session refreshed: id=%s new_expiry=%s", id, newExpiry)
	return session, nil
}

// SetData stores arbitrary key-value data in the session.
func (s *SessionStore) SetData(id string, key string, value interface{}) error {
	s.mu.Lock()
	session, ok := s.store[id]
	s.mu.Unlock()

	if !ok {
		return fmt.Errorf("session not found: %s", id)
	}

	s.mu.Lock()
	session.Data[key] = value
	s.mu.Unlock()
	return nil
}

// GetData retrieves arbitrary key-value data from the session.
func (s *SessionStore) GetData(id string, key string) (interface{}, error) {
	s.mu.RLock()
	session, ok := s.store[id]
	s.mu.RUnlock()

	if !ok {
		return nil, fmt.Errorf("session not found: %s", id)
	}

	val, ok := session.Data[key]
	if !ok {
		return nil, fmt.Errorf("key not found in session: %s", key)
	}
	return val, nil
}

// CleanupExpired removes all expired sessions from the store.
// Should be called periodically (e.g., every 5 minutes) as a background task.
func (s *SessionStore) CleanupExpired() int {
	s.mu.Lock()
	now := time.Now()
	count := 0
	for id, session := range s.store {
		if now.After(session.ExpiresAt) {
			delete(s.store, id)
			count++
		}
	}
	s.mu.Unlock()
	applogger.L().Debugf("Cleaned up %d expired sessions", count)
	return count
}

// Count returns the number of active sessions in the store.
func (s *SessionStore) Count() int {
	s.mu.RLock()
	count := len(s.store)
	s.mu.RUnlock()
	return count
}

// generateSessionID creates a cryptographically random session identifier.
func generateSessionID(length int) (string, error) {
	b := make([]byte, length)
	if _, err := rand.Read(b); err != nil {
		return "", fmt.Errorf("random generation failed: %w", err)
	}
	return hex.EncodeToString(b), nil
}