goodie/pkg/server/session.go

package server

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"log"
	"sync"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
)

// Session represents an authenticated user session.
type Session struct {
	ID                 string
	UID                int64
	CompanyID          int64
	AllowedCompanyIDs  []int64
	Login              string
	CSRFToken          string
	CreatedAt          time.Time
	LastActivity       time.Time
}

// SessionStore is a session store with an in-memory cache backed by PostgreSQL.
// Mirrors: odoo/http.py OpenERPSession
type SessionStore struct {
	mu       sync.RWMutex
	sessions map[string]*Session
	ttl      time.Duration
	pool     *pgxpool.Pool
}

// NewSessionStore creates a new session store with the given TTL and DB pool.
func NewSessionStore(ttl time.Duration, pool *pgxpool.Pool) *SessionStore {
	return &SessionStore{
		sessions: make(map[string]*Session),
		ttl:      ttl,
		pool:     pool,
	}
}

// InitSessionTable creates the sessions table if it does not exist.
func InitSessionTable(ctx context.Context, pool *pgxpool.Pool) error {
	_, err := pool.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS sessions (
			id VARCHAR(64) PRIMARY KEY,
			uid INT8 NOT NULL,
			company_id INT8 NOT NULL,
			login VARCHAR(255),
			csrf_token VARCHAR(64) DEFAULT '',
			created_at TIMESTAMP DEFAULT NOW(),
			last_seen TIMESTAMP DEFAULT NOW()
		)
	`)
	if err == nil {
		// Add csrf_token column if table already exists without it
		pool.Exec(ctx, `ALTER TABLE sessions ADD COLUMN IF NOT EXISTS csrf_token VARCHAR(64) DEFAULT ''`)
	}
	if err != nil {
		return err
	}
	log.Println("odoo: sessions table ready")
	return nil
}

// New creates a new session, stores it in memory and PostgreSQL, and returns it.
func (s *SessionStore) New(uid, companyID int64, login string) *Session {
	token := generateToken()
	now := time.Now()
	sess := &Session{
		ID:           token,
		UID:          uid,
		CompanyID:    companyID,
		Login:        login,
		CSRFToken:    generateToken(),
		CreatedAt:    now,
		LastActivity: now,
	}

	// Store in memory cache
	s.mu.Lock()
	s.sessions[token] = sess
	s.mu.Unlock()

	// Persist to PostgreSQL
	if s.pool != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_, err := s.pool.Exec(ctx,
			`INSERT INTO sessions (id, uid, company_id, login, csrf_token, created_at, last_seen)
			 VALUES ($1, $2, $3, $4, $5, $6, $7)
			 ON CONFLICT (id) DO NOTHING`,
			token, uid, companyID, login, sess.CSRFToken, now, now)
		if err != nil {
			log.Printf("session: failed to persist session to DB: %v", err)
		}
	}

	return sess
}

// Get retrieves a session by ID. Checks in-memory cache first, falls back to DB.
// Returns nil if not found or expired.
func (s *SessionStore) Get(id string) *Session {
	// Check memory cache first
	s.mu.RLock()
	sess, ok := s.sessions[id]
	s.mu.RUnlock()

	if ok {
		if time.Since(sess.LastActivity) > s.ttl {
			s.Delete(id)
			return nil
		}

		now := time.Now()
		needsDBUpdate := time.Since(sess.LastActivity) > 30*time.Second

		// Update last activity in memory
		s.mu.Lock()
		sess.LastActivity = now
		s.mu.Unlock()

		// Throttle DB writes: only persist every 30s to avoid per-request overhead
		if needsDBUpdate && s.pool != nil {
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			if _, err := s.pool.Exec(ctx,
				`UPDATE sessions SET last_seen = $1 WHERE id = $2`, now, id); err != nil {
				log.Printf("session: failed to update last_seen in DB: %v", err)
			}
		}

		return sess
	}

	// Fallback to DB
	if s.pool == nil {
		return nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	sess = &Session{}
	var csrfToken string
	err := s.pool.QueryRow(ctx,
		`SELECT id, uid, company_id, login, COALESCE(csrf_token, ''), created_at, last_seen
		 FROM sessions WHERE id = $1`, id).Scan(
		&sess.ID, &sess.UID, &sess.CompanyID, &sess.Login, &csrfToken,
		&sess.CreatedAt, &sess.LastActivity)
	if err != nil {
		return nil
	}
	if csrfToken != "" {
		sess.CSRFToken = csrfToken
	} else {
		sess.CSRFToken = generateToken()
	}

	// Check TTL
	if time.Since(sess.LastActivity) > s.ttl {
		s.Delete(id)
		return nil
	}

	// Update last activity and add to memory cache
	now := time.Now()
	s.mu.Lock()
	sess.LastActivity = now
	s.sessions[id] = sess
	s.mu.Unlock()

	// Update last_seen in DB
	if _, err := s.pool.Exec(ctx,
		`UPDATE sessions SET last_seen = $1 WHERE id = $2`, now, id); err != nil {
		log.Printf("session: failed to update last_seen in DB: %v", err)
	}

	return sess
}

// Delete removes a session from memory and DB.
func (s *SessionStore) Delete(id string) {
	s.mu.Lock()
	delete(s.sessions, id)
	s.mu.Unlock()

	if s.pool != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_, err := s.pool.Exec(ctx, `DELETE FROM sessions WHERE id = $1`, id)
		if err != nil {
			log.Printf("session: failed to delete session from DB: %v", err)
		}
	}
}

func generateToken() string {
	b := make([]byte, 32)
	rand.Read(b)
	return hex.EncodeToString(b)
}