NextWks/src/core/auth/session.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"fmt"
	"net/http"
	"time"
)

// contextKey is used for storing values in request context.
type contextKey string

const (
	ContextUserID contextKey = "user_id"
	ContextRole   contextKey = "role"
)

// SessionStore manages user sessions backed by SQLite.
type SessionStore struct {
	db     *sql.DB
	roleDB *sql.DB // Optional: same DB, used for role lookups
}

// NewSessionStore creates a session store.
func NewSessionStore(db *sql.DB) *SessionStore {
	return &SessionStore{db: db, roleDB: db}
}

// Session represents an authenticated user session.
type Session struct {
	ID        string
	UserID    string
	CreatedAt time.Time
	ExpiresAt time.Time
}

// CreateSession generates a new session for a user and returns the token.
func (s *SessionStore) CreateSession(userID string, expiryMinutes int) (string, error) {
	token := generateToken(32)
	tokenHash := hashToken(token)

	_, err := s.db.Exec(
		`INSERT INTO sessions (id, user_id, token_hash, created_at, expires_at)
		 VALUES (?, ?, ?, datetime('now'), datetime('now', '+' || ? || ' minutes'))`,
		token[:16], userID, tokenHash, expiryMinutes,
	)
	if err != nil {
		return "", fmt.Errorf("create session: %w", err)
	}

	return token, nil
}

// ValidateSession checks if a session token is valid and returns the session.
func (s *SessionStore) ValidateSession(token string) (*Session, error) {
	tokenHash := hashToken(token)

	var sess Session
	var createdAt, expiresAt string
	err := s.db.QueryRow(
		`SELECT id, user_id, created_at, expires_at
		 FROM sessions
		 WHERE token_hash = ? AND expires_at > datetime('now')`,
		tokenHash,
	).Scan(&sess.ID, &sess.UserID, &createdAt, &expiresAt)

	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("validate session: %w", err)
	}

	sess.CreatedAt, _ = time.Parse("2006-01-02 15:04:05", createdAt)
	sess.ExpiresAt, _ = time.Parse("2006-01-02 15:04:05", expiresAt)

	return &sess, nil
}

// DeleteSession removes a session (logout).
func (s *SessionStore) DeleteSession(token string) error {
	tokenHash := hashToken(token)
	_, err := s.db.Exec("DELETE FROM sessions WHERE token_hash = ?", tokenHash)
	return err
}

// CleanExpired removes all expired sessions.
func (s *SessionStore) CleanExpired() error {
	_, err := s.db.Exec("DELETE FROM sessions WHERE expires_at <= datetime('now')")
	return err
}

// SessionMiddleware returns an HTTP middleware that validates session cookies.
// If valid, the user_id is stored in the request context.
func (s *SessionStore) SessionMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		cookie, err := r.Cookie("nextwks_session")
		if err != nil {
			// No cookie — pass through without session
			next.ServeHTTP(w, r)
			return
		}

		session, err := s.ValidateSession(cookie.Value)
		if err != nil || session == nil {
			// Invalid or expired — clear cookie and continue
			http.SetCookie(w, &http.Cookie{
				Name:     "nextwks_session",
				Value:    "",
				Path:     "/",
				MaxAge:   -1,
				HttpOnly: true,
				SameSite: http.SameSiteStrictMode,
			})
			next.ServeHTTP(w, r)
			return
		}

		// Set user_id in context
		// Set user_id and role in context
		ctx := context.WithValue(r.Context(), ContextUserID, session.UserID)
		// Look up role from database
		var role string
		s.roleDB.QueryRow("SELECT role FROM users WHERE username = ?", session.UserID).Scan(&role)
		if role == "" {
			role = "user"
		}
		ctx = context.WithValue(ctx, ContextRole, role)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// GetUserID retrieves the authenticated user ID from the request context.
func GetUserID(r *http.Request) (string, bool) {
	uid, ok := r.Context().Value(ContextUserID).(string)
	return uid, ok
}

// generateToken creates a cryptographically secure random hex token.
func generateToken(length int) string {
	b := make([]byte, length)
	rand.Read(b)
	return hex.EncodeToString(b)
}

// hashToken creates a SHA-256 hash of a token for storage.
func hashToken(token string) string {
	h := sha256.Sum256([]byte(token))
	return hex.EncodeToString(h[:])
}