Merge pull request #30 from tstromberg/main

fix reservation TOCTOU

Author: tstromberg

cmd/server/main.go

@@ -32,7 +32,7 @@ const (
 	minMaskHeaderLength = 20  // Minimum header length before we show full "[REDACTED]"
 )
 
-// getEnvOrDefault returns the value of the environment variable or the default if not set
+// getEnvOrDefault returns the value of the environment variable or the default if not set.
 func getEnvOrDefault(key, defaultValue string) string {
 	if value := os.Getenv(key); value != "" {
 		return value
@@ -232,8 +232,9 @@ func main() {
 			return
 		}
 
-		// Check connection limit before upgrade
-		if !connLimiter.CanAdd(ip) {
+		// Reserve a connection slot before upgrade (prevents TOCTOU race condition)
+		reservationToken := connLimiter.Reserve(ip)
+		if reservationToken == "" {
 			log.Printf("WebSocket 429: connection limit ip=%s", ip)
 			w.WriteHeader(http.StatusTooManyRequests)
 			if _, err := w.Write([]byte("429 Too Many Requests: Connection limit exceeded\n")); err != nil {
@@ -242,6 +243,9 @@ func main() {
 			return
 		}
 
+		// Set reservation token in request context so websocket handler can commit it
+		r = r.WithContext(context.WithValue(r.Context(), "reservation_token", reservationToken))
+
 		// Log successful auth and proceed to upgrade
 		log.Printf("WebSocket UPGRADE: ip=%s duration=%v", ip, time.Since(startTime))
 

pkg/client/client.go

@@ -287,15 +287,18 @@ func (c *Client) connect(ctx context.Context) error {
 		}
 		return fmt.Errorf("dial: %w", err)
 	}
-	c.logger.Info("✓ WebSocket connection ESTABLISHED successfully!")
+	c.logger.Info("========================================")
+	c.logger.Info(fmt.Sprintf("✅ WebSocket ESTABLISHED: %s (org: %s)", c.config.ServerURL, c.config.Organization))
+	c.logger.Info("========================================")
 
 	// Store connection
 	c.mu.Lock()
 	c.ws = ws
 	c.mu.Unlock()
 
 	defer func() {
-		c.logger.Debug("Closing WebSocket connection")
+		c.logger.Info("========================================")
+		c.logger.Info(fmt.Sprintf("❌ WebSocket CLOSING: %s (org: %s)", c.config.ServerURL, c.config.Organization))
 		c.mu.Lock()
 		c.ws = nil
 		c.mu.Unlock()
@@ -304,6 +307,7 @@ func (c *Client) connect(ctx context.Context) error {
 		} else {
 			c.logger.Info("✓ WebSocket connection closed cleanly")
 		}
+		c.logger.Info("========================================")
 	}()
 
 	// Build subscription

pkg/logger/logger.go

@@ -1,42 +1,92 @@
-// Package logger provides structured logging utilities with field support
-// for better debugging and monitoring of webhook sprinkler operations.
+// Package logger provides structured logging using slog with hostname tracking
+// and short source file paths for better debugging across multiple instances.
 package logger
 
 import (
+	"context"
 	"fmt"
-	"log"
-	"sort"
-	"strings"
+	"io"
+	"log/slog"
+	"os"
+	"path/filepath"
+	"runtime"
+	"time"
 )
 
 // Fields represents structured log fields.
 type Fields map[string]any
 
-// WithFieldsf adds structured context to log messages with printf-style formatting.
-func WithFieldsf(fields Fields, format string, args ...any) {
-	// Sort keys for consistent output
-	keys := make([]string, 0, len(fields))
-	for k := range fields {
-		keys = append(keys, k)
-	}
-	sort.Strings(keys)
+var (
+	// defaultLogger is the global logger instance.
+	defaultLogger *slog.Logger
+	// hostname is cached on init for performance.
+	hostname string
+)
 
-	var parts []string
-	for _, k := range keys {
-		parts = append(parts, fmt.Sprintf("%s=%v", k, fields[k]))
+func init() {
+	var err error
+	hostname, err = os.Hostname()
+	if err != nil {
+		hostname = "unknown"
 	}
 
-	msg := fmt.Sprintf(format, args...)
-	if len(parts) > 0 {
-		log.Printf("%s [%s]", msg, strings.Join(parts, " "))
-	} else {
-		log.Print(msg)
+	// Initialize with default text handler
+	defaultLogger = New(os.Stderr)
+}
+
+// New creates a new slog logger with hostname and short source paths.
+func New(w io.Writer) *slog.Logger {
+	opts := &slog.HandlerOptions{
+		AddSource: true,
+		Level:     slog.LevelInfo,
+		ReplaceAttr: func(_ []string, a slog.Attr) slog.Attr {
+			// Shorten source file paths to just basename:line
+			if a.Key == slog.SourceKey {
+				if source, ok := a.Value.Any().(*slog.Source); ok {
+					source.File = filepath.Base(source.File)
+					// Remove function name to keep it concise
+					source.Function = ""
+				}
+			}
+			return a
+		},
 	}
+
+	handler := slog.NewTextHandler(w, opts)
+	logger := slog.New(handler)
+
+	// Add hostname to all log messages
+	return logger.With("instance", hostname)
+}
+
+// SetDefault sets the default logger.
+func SetDefault(l *slog.Logger) {
+	defaultLogger = l
+}
+
+// SetLogger sets the default logger (alias for SetDefault).
+func SetLogger(l *slog.Logger) {
+	defaultLogger = l
+}
+
+// Default returns the default logger.
+func Default() *slog.Logger {
+	return defaultLogger
+}
+
+// Hostname returns the cached hostname.
+func Hostname() string {
+	return hostname
 }
 
 // Info logs an info message with optional fields.
 func Info(msg string, fields Fields) {
-	WithFieldsf(fields, "%s", msg)
+	defaultLogger.LogAttrs(context.Background(), slog.LevelInfo, msg, attrsFromFields(fields)...)
+}
+
+// Warn logs a warning message with optional fields.
+func Warn(msg string, fields Fields) {
+	defaultLogger.LogAttrs(context.Background(), slog.LevelWarn, msg, attrsFromFields(fields)...)
 }
 
 // Error logs an error message with optional fields.
@@ -45,10 +95,44 @@ func Error(msg string, err error, fields Fields) {
 		fields = Fields{}
 	}
 	fields["error"] = err.Error()
-	WithFieldsf(fields, "ERROR: %s", msg)
+	defaultLogger.LogAttrs(context.Background(), slog.LevelError, msg, attrsFromFields(fields)...)
 }
 
-// Warn logs a warning message with optional fields.
-func Warn(msg string, fields Fields) {
-	WithFieldsf(fields, "WARNING: %s", msg)
+// Debug logs a debug message with optional fields.
+func Debug(msg string, fields Fields) {
+	defaultLogger.LogAttrs(context.Background(), slog.LevelDebug, msg, attrsFromFields(fields)...)
+}
+
+// attrsFromFields converts Fields to slog.Attr slice.
+func attrsFromFields(fields Fields) []slog.Attr {
+	if fields == nil {
+		return nil
+	}
+	attrs := make([]slog.Attr, 0, len(fields))
+	for k, v := range fields {
+		attrs = append(attrs, slog.Any(k, v))
+	}
+	return attrs
+}
+
+// LogAt logs a message at the specified level with source information.
+// This is useful when you want to override the default source location.
+func LogAt(level slog.Level, skip int, msg string, fields Fields) {
+	var pcs [1]uintptr
+	runtime.Callers(skip+2, pcs[:])
+	r := slog.NewRecord(
+		time.Now(),
+		level,
+		msg,
+		pcs[0],
+	)
+	r.AddAttrs(attrsFromFields(fields)...)
+	_ = defaultLogger.Handler().Handle(context.Background(), r) //nolint:errcheck // Best effort logging
+}
+
+// WithFieldsf provides backward compatibility for tests.
+// Deprecated: Use Info/Warn/Error with Fields instead.
+func WithFieldsf(fields Fields, format string, args ...any) {
+	msg := fmt.Sprintf(format, args...)
+	Info(msg, fields)
 }