Merge pull request #10 from StudioLambda/fix/critical-audit-findings

Fix all critical audit findings (router, database, events)

Author: ConsoleTVs

framework/database/sql.go

@@ -174,7 +174,10 @@ func (database *SQL) FindNamed(ctx context.Context, query string, dest any, arg
 // original error and any rollback error are joined. If fn succeeds,
 // the transaction is committed. Nested transactions are not supported
 // and return contract.ErrDatabaseNestedTransaction.
-func (database *SQL) WithTransaction(ctx context.Context, fn func(tx contract.Database) error) error {
+//
+// If fn panics, the transaction is rolled back before the panic is
+// re-raised, preventing connection pool leaks.
+func (database *SQL) WithTransaction(ctx context.Context, fn func(tx contract.Database) error) (retErr error) {
 	if _, ok := database.db.(*sqlx.Tx); ok {
 		return contract.ErrDatabaseNestedTransaction
 	}
@@ -185,10 +188,23 @@ func (database *SQL) WithTransaction(ctx context.Context, fn func(tx contract.Da
 		return err
 	}
 
+	defer func() {
+		if p := recover(); p != nil {
+			_ = tx.Rollback()
+			panic(p)
+		}
+
+		if retErr != nil {
+			retErr = errors.Join(retErr, tx.Rollback())
+		}
+	}()
+
 	txWrapper := &SQL{db: tx, raw: database.raw}
 
 	if err := fn(txWrapper); err != nil {
-		return errors.Join(err, tx.Rollback())
+		retErr = err
+
+		return
 	}
 
 	return tx.Commit()

framework/database/sql_test.go

@@ -0,0 +1,150 @@
+//go:build cgo
+
+package database_test
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	_ "github.com/mattn/go-sqlite3"
+	"github.com/stretchr/testify/require"
+	"github.com/studiolambda/cosmos/contract"
+	"github.com/studiolambda/cosmos/framework/database"
+)
+
+func newTestDB(t *testing.T) *database.SQL {
+	t.Helper()
+
+	db, err := database.NewSQL("sqlite3", ":memory:")
+
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		require.NoError(t, db.Close())
+	})
+
+	return db
+}
+
+func TestWithTransactionPanicRollsBack(t *testing.T) {
+	t.Parallel()
+
+	db := newTestDB(t)
+	ctx := context.Background()
+
+	_, err := db.Exec(ctx, "CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)")
+
+	require.NoError(t, err)
+
+	require.Panics(t, func() {
+		_ = db.WithTransaction(ctx, func(tx contract.Database) error {
+			_, err := tx.Exec(ctx, "INSERT INTO items (name) VALUES (?)", "should-be-rolled-back")
+
+			if err != nil {
+				return err
+			}
+
+			panic("unexpected failure")
+		})
+	})
+
+	// The row must not exist because the transaction was rolled back.
+	var count int
+	err = db.Find(ctx, "SELECT COUNT(*) FROM items", &count)
+
+	require.NoError(t, err)
+	require.Equal(t, 0, count)
+}
+
+func TestWithTransactionNestedReturnsError(t *testing.T) {
+	t.Parallel()
+
+	db := newTestDB(t)
+	ctx := context.Background()
+
+	err := db.WithTransaction(ctx, func(tx contract.Database) error {
+		return tx.WithTransaction(ctx, func(inner contract.Database) error {
+			return nil
+		})
+	})
+
+	require.ErrorIs(t, err, contract.ErrDatabaseNestedTransaction)
+}
+
+func TestWithTransactionCommitsOnSuccess(t *testing.T) {
+	t.Parallel()
+
+	db := newTestDB(t)
+	ctx := context.Background()
+
+	_, err := db.Exec(ctx, "CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)")
+
+	require.NoError(t, err)
+
+	err = db.WithTransaction(ctx, func(tx contract.Database) error {
+		_, err := tx.Exec(ctx, "INSERT INTO items (name) VALUES (?)", "committed")
+
+		return err
+	})
+
+	require.NoError(t, err)
+
+	var count int
+	err = db.Find(ctx, "SELECT COUNT(*) FROM items", &count)
+
+	require.NoError(t, err)
+	require.Equal(t, 1, count)
+}
+
+func TestWithTransactionRollsBackOnError(t *testing.T) {
+	t.Parallel()
+
+	db := newTestDB(t)
+	ctx := context.Background()
+
+	_, err := db.Exec(ctx, "CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)")
+
+	require.NoError(t, err)
+
+	callbackErr := errors.New("something went wrong")
+
+	err = db.WithTransaction(ctx, func(tx contract.Database) error {
+		_, err := tx.Exec(ctx, "INSERT INTO items (name) VALUES (?)", "should-be-rolled-back")
+
+		if err != nil {
+			return err
+		}
+
+		return callbackErr
+	})
+
+	require.ErrorIs(t, err, callbackErr)
+
+	var count int
+	err = db.Find(ctx, "SELECT COUNT(*) FROM items", &count)
+
+	require.NoError(t, err)
+	require.Equal(t, 0, count)
+}
+
+func TestWithTransactionPanicPreservesValue(t *testing.T) {
+	t.Parallel()
+
+	db := newTestDB(t)
+	ctx := context.Background()
+
+	var recovered any
+
+	func() {
+		defer func() {
+			recovered = recover()
+		}()
+
+		_ = db.WithTransaction(ctx, func(tx contract.Database) error {
+			panic("test-panic-value")
+		})
+	}()
+
+	require.Equal(t, "test-panic-value", recovered)
+}

framework/event/memory.go

@@ -109,26 +109,31 @@ func (broker *MemoryBroker) Publish(
 	}
 
 	broker.mu.RLock()
-	defer broker.mu.RUnlock()
+
+	var matched []contract.EventHandler
 
 	for pattern, patternHandlers := range broker.handlers {
-		if !matchEvent(pattern, event) {
-			continue
+		if matchEvent(pattern, event) {
+			for _, handler := range patternHandlers {
+				matched = append(matched, handler)
+			}
 		}
+	}
 
-		for _, handler := range patternHandlers {
-			broker.wg.Add(1)
-			broker.sem <- struct{}{}
+	broker.mu.RUnlock()
 
-			go func(h contract.EventHandler) {
-				defer func() {
-					<-broker.sem
-					broker.wg.Done()
-				}()
+	for _, handler := range matched {
+		broker.wg.Add(1)
+		broker.sem <- struct{}{}
 
-				broker.deliverToHandler(h, encoded)
-			}(handler)
-		}
+		go func() {
+			defer func() {
+				<-broker.sem
+				broker.wg.Done()
+			}()
+
+			broker.deliverToHandler(handler, encoded)
+		}()
 	}
 
 	return nil

framework/event/memory_test.go

@@ -84,7 +84,7 @@ func TestMemoryBrokerWildcardStar(t *testing.T) {
 	require.Equal(t, int64(1), atomic.LoadInt64(&received))
 }
 
-func TestMemoryBrokerWildcardStarDoesNotMatchMultipleTokens(t *testing.T) {
+func TestMemoryBrokerPublishDoesNotDeadlockWithSubscribeInHandler(t *testing.T) {
 	t.Parallel()
 
 	ctx := context.Background()
@@ -94,25 +94,32 @@ func TestMemoryBrokerWildcardStarDoesNotMatchMultipleTokens(t *testing.T) {
 		_ = broker.Close()
 	})
 
-	var received int64
+	done := make(chan struct{})
 
 	_, err := broker.Subscribe(
-		ctx, "user.*.created", func(payload contract.EventPayload) {
-			atomic.AddInt64(&received, 1)
+		ctx, "test.event", func(payload contract.EventPayload) {
+			// This Subscribe call requires broker.mu.Lock().
+			// If Publish still holds broker.mu.RLock() during
+			// dispatch, this will deadlock.
+			_, _ = broker.Subscribe(
+				ctx, "other.event", func(payload contract.EventPayload) {},
+			)
+
+			close(done)
 		},
 	)
 
 	require.NoError(t, err)
 
-	err = broker.Publish(
-		ctx, "user.123.456.created", "data",
-	)
-
+	err = broker.Publish(ctx, "test.event", "data")
 	require.NoError(t, err)
 
-	time.Sleep(50 * time.Millisecond)
-
-	require.Equal(t, int64(0), atomic.LoadInt64(&received))
+	select {
+	case <-done:
+		// Handler completed without deadlock.
+	case <-time.After(3 * time.Second):
+		t.Fatal("deadlock detected: handler calling Subscribe blocked for 3 seconds")
+	}
 }
 
 func TestMemoryBrokerWildcardHash(t *testing.T) {

framework/event/mqtt.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"log/slog"
 	"net/url"
 	"strconv"
 	"strings"
@@ -196,11 +197,7 @@ func NewMQTTBrokerWith(options *MQTTBrokerOptions) (*MQTTBroker, error) {
 		ClientConfig: paho.ClientConfig{
 			ClientID: "",
 			OnPublishReceived: []func(paho.PublishReceived) (bool, error){
-				func(pr paho.PublishReceived) (bool, error) {
-					broker.route(pr.Packet)
-
-					return true, nil
-				},
+				broker.HandlePublish,
 			},
 		},
 	}
@@ -225,6 +222,19 @@ func NewMQTTBrokerWith(options *MQTTBrokerOptions) (*MQTTBroker, error) {
 // autopaho ConnectionManager and QoS level. This constructor is
 // useful for advanced scenarios where the user needs full control
 // over the MQTT connection configuration.
+//
+// Because autopaho.ConnectionManager does not allow post-creation
+// configuration changes, the caller must wire up message routing
+// when building the paho config. Use [MQTTBroker.HandlePublish] in
+// the ClientConfig.OnPublishReceived slice:
+//
+//	broker := event.NewMQTTBrokerFrom(nil, 1)
+//	cfg.ClientConfig.OnPublishReceived = []func(paho.PublishReceived) (bool, error){
+//	    broker.HandlePublish,
+//	}
+//	cm, err := autopaho.NewConnection(ctx, cfg)
+//	// then set the client so Publish/Subscribe/Close work:
+//	broker.SetClient(cm)
 func NewMQTTBrokerFrom(
 	client *autopaho.ConnectionManager,
 	qos byte,
@@ -237,25 +247,69 @@ func NewMQTTBrokerFrom(
 	}
 }
 
+// SetClient sets the underlying autopaho ConnectionManager. This is
+// intended for use with [NewMQTTBrokerFrom] when the broker must be
+// created before the ConnectionManager so that [MQTTBroker.HandlePublish]
+// can be wired into the paho config.
+func (broker *MQTTBroker) SetClient(client *autopaho.ConnectionManager) {
+	broker.client = client
+}
+
+// HandlePublish is a paho OnPublishReceived callback that routes
+// incoming MQTT messages to registered handlers. Callers using
+// [NewMQTTBrokerFrom] must include this method in the paho
+// ClientConfig.OnPublishReceived slice so that subscribed handlers
+// receive messages.
+func (broker *MQTTBroker) HandlePublish(pr paho.PublishReceived) (bool, error) {
+	broker.route(pr.Packet)
+
+	return true, nil
+}
+
 // route delivers an incoming MQTT message to all matching
 // handlers based on topic pattern matching. This implements
 // fan-out behavior where multiple handlers can receive the
 // same message if they subscribed to matching patterns.
+//
+// Each handler is called with panic recovery so that a
+// panicking handler does not crash the paho client goroutine
+// and tear down the entire MQTT connection.
 func (broker *MQTTBroker) route(pb *paho.Publish) {
 	broker.mu.RLock()
 	defer broker.mu.RUnlock()
 
 	for pattern, handlers := range broker.handlers {
 		if matchTopic(pattern, pb.Topic) {
 			for _, handler := range handlers {
-				handler(func(dest any) error {
-					return json.Unmarshal(pb.Payload, dest)
-				})
+				broker.deliverToHandler(handler, pb.Topic, pb.Payload)
 			}
 		}
 	}
 }
 
+// deliverToHandler invokes a single handler with panic recovery.
+// Recovered panics are logged via slog so they remain visible for
+// debugging without propagating to the caller.
+func (broker *MQTTBroker) deliverToHandler(
+	handler contract.EventHandler,
+	topic string,
+	payload []byte,
+) {
+	defer func() {
+		if recovered := recover(); recovered != nil {
+			slog.Error(
+				"mqtt event handler panicked",
+				"topic", topic,
+				"error", recovered,
+			)
+		}
+	}()
+
+	handler(func(dest any) error {
+		return json.Unmarshal(payload, dest)
+	})
+}
+
 // Publish sends an event with the given name and payload to all
 // subscribers listening for that event. The payload is serialized
 // to JSON and the event name is converted to MQTT topic format.