fix panic

Author: Thomas Stromberg

cmd/server/main.go

@@ -32,13 +32,10 @@ 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.
-func getEnvOrDefault(key, defaultValue string) string {
-	if value := os.Getenv(key); value != "" {
-		return value
-	}
-	return defaultValue
-}
+// contextKey is a custom type for context keys to avoid collisions.
+type contextKey string
+
+const reservationTokenKey contextKey = "reservation_token"
 
 var (
 	webhookSecret = flag.String("webhook-secret", os.Getenv("GITHUB_WEBHOOK_SECRET"), "GitHub webhook secret for signature verification")
@@ -50,8 +47,12 @@ var (
 	maxConnsPerIP = flag.Int("max-conns-per-ip", 10, "Maximum WebSocket connections per IP")
 	maxConnsTotal = flag.Int("max-conns-total", 1000, "Maximum total WebSocket connections")
 	rateLimit     = flag.Int("rate-limit", 100, "Maximum requests per minute per IP")
-	allowedEvents = flag.String("allowed-events", getEnvOrDefault("ALLOWED_WEBHOOK_EVENTS", "*"),
-		"Comma-separated list of allowed webhook event types (use '*' for all, default: '*')")
+	allowedEvents = flag.String("allowed-events", func() string {
+		if value := os.Getenv("ALLOWED_WEBHOOK_EVENTS"); value != "" {
+			return value
+		}
+		return "*"
+	}(), "Comma-separated list of allowed webhook event types (use '*' for all, default: '*')")
 	debugHeaders = flag.Bool("debug-headers", false, "Log request headers for debugging (security warning: may log sensitive data)")
 )
 
@@ -244,7 +245,7 @@ func main() {
 		}
 
 		// Set reservation token in request context so websocket handler can commit it
-		r = r.WithContext(context.WithValue(r.Context(), "reservation_token", reservationToken))
+		r = r.WithContext(context.WithValue(r.Context(), reservationTokenKey, reservationToken))
 
 		// Log successful auth and proceed to upgrade
 		log.Printf("WebSocket UPGRADE: ip=%s duration=%v", ip, time.Since(startTime))

pkg/client/client.go

@@ -85,7 +85,8 @@ type Client struct {
 	ws         *websocket.Conn
 	stopCh     chan struct{}
 	stoppedCh  chan struct{}
-	writeCh    chan any // Channel for serializing all writes
+	stopOnce   sync.Once // Ensures Stop() is only executed once
+	writeCh    chan any  // Channel for serializing all writes
 	eventCount int
 	retries    int
 }
@@ -220,16 +221,27 @@ func (c *Client) Start(ctx context.Context) error {
 }
 
 // Stop gracefully stops the client.
+// Safe to call multiple times - only the first call will take effect.
+// Also safe to call before Start() or if Start() was never called.
 func (c *Client) Stop() {
-	close(c.stopCh)
-	c.mu.Lock()
-	if c.ws != nil {
-		if closeErr := c.ws.Close(); closeErr != nil {
-			c.logger.Error("Error closing websocket on shutdown", "error", closeErr)
+	c.stopOnce.Do(func() {
+		close(c.stopCh)
+		c.mu.Lock()
+		if c.ws != nil {
+			if closeErr := c.ws.Close(); closeErr != nil {
+				c.logger.Error("Error closing websocket on shutdown", "error", closeErr)
+			}
 		}
-	}
-	c.mu.Unlock()
-	<-c.stoppedCh
+		c.mu.Unlock()
+
+		// Wait for Start() to finish, but with timeout in case Start() was never called
+		select {
+		case <-c.stoppedCh:
+			// Start() completed normally
+		case <-time.After(100 * time.Millisecond):
+			// Start() was never called or hasn't started yet - that's ok
+		}
+	})
 }
 
 // connect establishes a WebSocket connection and handles events.

pkg/security/connlimiter.go

@@ -25,8 +25,8 @@ type connectionInfo struct {
 
 // reservation represents a reserved connection slot.
 type reservation struct {
-	ip        string
 	createdAt time.Time
+	ip        string
 }
 
 // ConnectionLimiter tracks connections per IP and total.
@@ -185,7 +185,7 @@ func (cl *ConnectionLimiter) CancelReservation(token string) {
 }
 
 // CanAdd checks if a connection can be added for the given IP without actually adding it.
-// DEPRECATED: This method has a TOCTOU race condition. Use Reserve() instead, which
+// Deprecated: This method has a TOCTOU race condition. Use Reserve() instead, which
 // atomically checks and reserves a slot, preventing the race.
 //
 // This method is kept for backward compatibility and testing only.

pkg/security/race_test.go

@@ -114,9 +114,9 @@ func TestConnectionLimiterTOCTOU_Documentation(t *testing.T) {
 	ip := "192.168.1.1"
 
 	var wg sync.WaitGroup
-	var canAddSuccess int32  // How many times CanAdd returned true
-	var addSuccess int32     // How many times Add actually succeeded
-	var addFailed int32      // How many times Add failed despite CanAdd=true
+	var canAddSuccess int32 // How many times CanAdd returned true
+	var addSuccess int32    // How many times Add actually succeeded
+	var addFailed int32     // How many times Add failed despite CanAdd=true
 
 	// Launch many goroutines simultaneously trying to add connections
 	// This simulates multiple HTTP handlers racing to add connections

pkg/srv/websocket.go

@@ -150,14 +150,6 @@ func (h *WebSocketHandler) extractGitHubToken(ws *websocket.Conn, ip string) (st
 	return githubToken, true
 }
 
-// tokenDebugInfo extracts token prefix for debug logging.
-func tokenDebugInfo(token string) string {
-	if len(token) >= tokenPrefixLength {
-		return token[:tokenPrefixLength]
-	}
-	return ""
-}
-
 // errorInfo holds error response details.
 type errorInfo struct {
 	code    string
@@ -281,7 +273,10 @@ func (*WebSocketHandler) handleAuthError(
 	logContext string,
 ) error {
 	errInfo := determineErrorInfo(err, username, orgName, userOrgs)
-	tokenPrefix := tokenDebugInfo(githubToken)
+	tokenPrefix := ""
+	if len(githubToken) >= tokenPrefixLength {
+		tokenPrefix = githubToken[:tokenPrefixLength]
+	}
 
 	logger.Error(logContext, err, logger.Fields{
 		"ip":           ip,
@@ -411,11 +406,6 @@ type wsCloser struct {
 	mu        sync.Mutex
 }
 
-// newWSCloser creates a new WebSocket closer wrapper.
-func newWSCloser(ws *websocket.Conn) *wsCloser {
-	return &wsCloser{ws: ws}
-}
-
 // Close closes the WebSocket connection exactly once.
 func (wc *wsCloser) Close() error {
 	var err error
@@ -491,7 +481,7 @@ func (h *WebSocketHandler) Handle(ws *websocket.Conn) {
 	log.Printf("WebSocket Handle() got IP: %s", ip)
 
 	// Wrap WebSocket with sync.Once closer to prevent double-close
-	wc := newWSCloser(ws)
+	wc := &wsCloser{ws: ws}
 
 	// Ensure WebSocket is properly closed (client will be set later if connection succeeds)
 	var client *Client
@@ -508,7 +498,8 @@ func (h *WebSocketHandler) Handle(ws *websocket.Conn) {
 	})
 
 	// Get reservation token from context (set by main.go before upgrade)
-	reservationToken, _ := ws.Request().Context().Value("reservation_token").(string)
+	// Context key is a string type for package boundary crossing
+	reservationToken, _ := ws.Request().Context().Value("reservation_token").(string) //nolint:errcheck // Type assertion intentionally unchecked - empty string is valid default
 	if reservationToken == "" {
 		// No reservation token - this should not happen in production
 		// (main.go always sets it), but handle gracefully for tests

pkg/webhook/handler.go

@@ -182,7 +182,12 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		// GitHub webhooks can fire before the pull_requests array is populated
 		commitSHA := extractCommitSHA(eventType, payload)
 		// Extract repo URL as fallback for org-based matching
-		repoURL := extractRepoURL(payload)
+		repoURL := ""
+		if repo, ok := payload["repository"].(map[string]any); ok {
+			if htmlURL, ok := repo["html_url"].(string); ok {
+				repoURL = htmlURL
+			}
+		}
 
 		// If we can't extract repo URL, drop the event
 		if repoURL == "" {
@@ -299,11 +304,15 @@ func ExtractPRURL(eventType string, payload map[string]any) string {
 			}
 		}
 		// Log when we can't extract PR URL from check event
+		payloadKeys := make([]string, 0, len(payload))
+		for k := range payload {
+			payloadKeys = append(payloadKeys, k)
+		}
 		logger.Warn("no PR URL found in check event", logger.Fields{
 			"event_type":      eventType,
 			"has_check_run":   payload["check_run"] != nil,
 			"has_check_suite": payload["check_suite"] != nil,
-			"payload_keys":    getPayloadKeys(payload),
+			"payload_keys":    payloadKeys,
 		})
 	default:
 		// For other event types, no PR URL can be extracted
@@ -378,15 +387,6 @@ func extractPRFromCheckEvent(checkEvent map[string]any, payload map[string]any,
 	return constructedURL
 }
 
-// getPayloadKeys returns the keys from a payload map for logging.
-func getPayloadKeys(payload map[string]any) []string {
-	keys := make([]string, 0, len(payload))
-	for k := range payload {
-		keys = append(keys, k)
-	}
-	return keys
-}
-
 // getMapKeys returns the keys from a map for logging.
 func getMapKeys(m map[string]any) []string {
 	keys := make([]string, 0, len(m))
@@ -417,13 +417,3 @@ func extractCommitSHA(eventType string, payload map[string]any) string {
 	return ""
 }
 
-// extractRepoURL extracts the repository HTML URL from the payload.
-// This is used as a fallback when PR URL cannot be extracted (e.g., check event race condition).
-func extractRepoURL(payload map[string]any) string {
-	if repo, ok := payload["repository"].(map[string]any); ok {
-		if htmlURL, ok := repo["html_url"].(string); ok {
-			return htmlURL
-		}
-	}
-	return ""
-}