v0.1.4 (#38)

Reviewed-on: wrenn/wrenn#38
Co-authored-by: pptx704 <rafeed@omukk.dev>
Co-committed-by: pptx704 <rafeed@omukk.dev>

internal/api/middleware_auth.go

@@ -85,15 +85,61 @@ func requireAPIKeyOrJWT(queries *db.Queries, jwtSecret []byte) func(http.Handler
 				return
 			}
 
-			// WebSocket upgrade requests may not carry auth headers (browsers
-			// cannot set custom headers on WS connections). Pass through —
-			// the WS handler authenticates via the first message after upgrade.
-			if isWebSocketUpgrade(r) {
-				next.ServeHTTP(w, r)
-				return
-			}
-
 			writeError(w, http.StatusUnauthorized, "unauthorized", "X-API-Key or Authorization: Bearer <token> required")
 		})
 	}
 }
+
+// optionalAPIKeyOrJWT is like requireAPIKeyOrJWT but does not reject
+// unauthenticated requests. It injects auth context when valid credentials
+// are present (supporting SDK clients that set X-API-Key on WebSocket
+// upgrades) and passes through otherwise so the handler can authenticate
+// after the WebSocket upgrade via the first message.
+func optionalAPIKeyOrJWT(queries *db.Queries, jwtSecret []byte) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// Try API key.
+			if key := r.Header.Get("X-API-Key"); key != "" {
+				hash := auth.HashAPIKey(key)
+				row, err := queries.GetAPIKeyByHash(r.Context(), hash)
+				if err == nil {
+					if err := queries.UpdateAPIKeyLastUsed(r.Context(), row.ID); err != nil {
+						slog.Warn("failed to update api key last_used", "key_id", id.FormatAPIKeyID(row.ID), "error", err)
+					}
+					ctx := auth.WithAuthContext(r.Context(), auth.AuthContext{
+						TeamID:     row.TeamID,
+						APIKeyID:   row.ID,
+						APIKeyName: row.Name,
+					})
+					next.ServeHTTP(w, r.WithContext(ctx))
+					return
+				}
+			}
+
+			// Try JWT bearer token.
+			if header := r.Header.Get("Authorization"); strings.HasPrefix(header, "Bearer ") {
+				tokenStr := strings.TrimPrefix(header, "Bearer ")
+				if claims, err := auth.VerifyJWT(jwtSecret, tokenStr); err == nil {
+					if teamID, err := id.ParseTeamID(claims.TeamID); err == nil {
+						if userID, err := id.ParseUserID(claims.Subject); err == nil {
+							if user, err := queries.GetUserByID(r.Context(), userID); err == nil && user.Status == "active" {
+								ctx := auth.WithAuthContext(r.Context(), auth.AuthContext{
+									TeamID: teamID,
+									UserID: userID,
+									Email:  claims.Email,
+									Name:   claims.Name,
+									Role:   claims.Role,
+								})
+								next.ServeHTTP(w, r.WithContext(ctx))
+								return
+							}
+						}
+					}
+				}
+			}
+
+			// No valid credentials — pass through for handler to authenticate.
+			next.ServeHTTP(w, r)
+		})
+	}
+}