fix: address PR review — nil sessStore, error masking, WS origin check

1. Nil session store: runServe now falls back to an in-memory SQLite DB
   for session tracking when memory is disabled, preventing nil panics.

2. GetOrCreate hides errors: Added ErrSessionNotFound sentinel. The
   session store's Get method now returns it for sql.ErrNoRows and
   surfaces real DB errors. GetOrCreate only creates a new session on
   ErrSessionNotFound, propagating other failures to the caller.

3. WS origin unrestricted: Replaced the permissive CheckOrigin with a
   localhost-only check when no API key is configured, preventing
   cross-site WebSocket hijacking of a local gateway. When an API key
   is set, all origins are allowed since the hello frame authenticates.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: hackertron

cmd/yantra/main.go

@@ -357,6 +357,17 @@ func runServe(cmd *cobra.Command, args []string) error {
 		defer memDB.Close()
 	}
 
+	// Sessions are required for the gateway even when memory is disabled.
+	// Fall back to an in-memory SQLite DB for session tracking only.
+	if sessStore == nil {
+		sessDB, err := memory.OpenDB(":memory:")
+		if err != nil {
+			return fmt.Errorf("opening session DB: %w", err)
+		}
+		defer sessDB.Close()
+		sessStore = memory.NewSessionStore(sessDB)
+	}
+
 	policy := tool.NewWorkspacePolicy(cfg.Tools.Shell)
 	reg := tool.NewRegistry(policy)
 	if err := tool.RegisterBuiltins(reg, cfg.Tools, mem); err != nil {

internal/gateway/session.go

@@ -2,6 +2,7 @@ package gateway
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sync"
 	"time"
@@ -74,6 +75,9 @@ func (sm *SessionManager) GetOrCreate(ctx context.Context, sessionID string) (*M
 	// Look up or create the persistent session record.
 	rec, err := sm.server.sessStore.Get(ctx, sessionID)
 	if err != nil {
+		if !errors.Is(err, types.ErrSessionNotFound) {
+			return nil, fmt.Errorf("looking up session: %w", err)
+		}
 		// Not found — create a new record.
 		rec, err = sm.server.sessStore.Create(ctx, "gateway-session")
 		if err != nil {

internal/gateway/ws.go

@@ -5,18 +5,13 @@ import (
 	"encoding/json"
 	"log/slog"
 	"net/http"
+	"net/url"
 	"sync"
 
 	"github.com/gorilla/websocket"
 	"github.com/hackertron/Yantra/internal/types"
 )
 
-var upgrader = websocket.Upgrader{
-	ReadBufferSize:  4096,
-	WriteBufferSize: 4096,
-	CheckOrigin:     func(r *http.Request) bool { return true }, // permissive for dev
-}
-
 // wsConn wraps a single WebSocket connection.
 type wsConn struct {
 	conn      *websocket.Conn
@@ -26,8 +21,39 @@ type wsConn struct {
 	writeMu   sync.Mutex // gorilla writes are not concurrent-safe
 }
 
+// newUpgrader creates a WebSocket upgrader with origin checking appropriate
+// to the server's security configuration. When an API key is set, all
+// origins are allowed (the hello frame authenticates). When no API key is
+// set (dev/local mode), only localhost origins are permitted to prevent
+// cross-site WebSocket hijacking of a local gateway.
+func (s *Server) newUpgrader() websocket.Upgrader {
+	return websocket.Upgrader{
+		ReadBufferSize:  4096,
+		WriteBufferSize: 4096,
+		CheckOrigin: func(r *http.Request) bool {
+			// If an API key is required, WS auth via the hello frame
+			// protects the endpoint; allow any origin.
+			if s.cfg.APIKey != "" {
+				return true
+			}
+			// Dev mode: restrict to localhost origins.
+			origin := r.Header.Get("Origin")
+			if origin == "" {
+				return true // non-browser clients (curl, wscat)
+			}
+			u, err := url.Parse(origin)
+			if err != nil {
+				return false
+			}
+			host := u.Hostname()
+			return host == "localhost" || host == "127.0.0.1" || host == "::1"
+		},
+	}
+}
+
 func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
-	conn, err := upgrader.Upgrade(w, r, nil)
+	up := s.newUpgrader()
+	conn, err := up.Upgrade(w, r, nil)
 	if err != nil {
 		slog.Error("websocket upgrade failed", "error", err)
 		return

internal/memory/session_store.go

@@ -3,7 +3,9 @@ package memory
 import (
 	"context"
 	"crypto/rand"
+	"database/sql"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"time"
 
@@ -51,7 +53,10 @@ func (s *SQLiteSessionStore) Get(ctx context.Context, id string) (*types.Session
 		`SELECT name, created_at, updated_at, message_count, archived FROM sessions WHERE id = ?`, id).
 		Scan(&name, &createdAt, &updatedAt, &msgCount, &archived)
 	if err != nil {
-		return nil, &types.MemoryError{Op: "session_get", Message: "not found", Err: err}
+		if errors.Is(err, sql.ErrNoRows) {
+			return nil, types.ErrSessionNotFound
+		}
+		return nil, &types.MemoryError{Op: "session_get", Message: "query failed", Err: err}
 	}
 
 	rec := &types.SessionRecord{

internal/types/errors.go

@@ -4,10 +4,11 @@ import "errors"
 
 // Sentinel errors for the runtime.
 var (
-	ErrCancelled      = errors.New("turn cancelled")
-	ErrBudgetExceeded = errors.New("budget exceeded")
-	ErrMaxTurns       = errors.New("max turns reached")
-	ErrTimeout        = errors.New("turn timed out")
+	ErrCancelled        = errors.New("turn cancelled")
+	ErrBudgetExceeded   = errors.New("budget exceeded")
+	ErrMaxTurns         = errors.New("max turns reached")
+	ErrTimeout          = errors.New("turn timed out")
+	ErrSessionNotFound  = errors.New("session not found")
 )
 
 // ProviderError represents an error from an LLM provider.