pjson-rs 0.6.3

Priority JSON Streaming Protocol - high-performance priority-based JSON streaming (requires nightly Rust)
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
//! HTTP middleware for PJS optimization and monitoring

use axum::{
    extract::{ConnectInfo, Request},
    http::{HeaderMap, HeaderValue, StatusCode, header},
    middleware::Next,
    response::Response,
};
use std::net::SocketAddr;
use std::time::Instant;
use std::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};
use tower::{Layer, Service};

/// Middleware for performance monitoring and optimization
#[derive(Clone)]
pub struct PjsMiddleware {
    enable_compression: bool,
    enable_metrics: bool,
    max_request_size: usize,
}

impl PjsMiddleware {
    /// Construct middleware with default settings (compression and metrics enabled, 10 MiB cap).
    pub fn new() -> Self {
        Self {
            enable_compression: true,
            enable_metrics: true,
            max_request_size: 10 * 1024 * 1024, // 10MB
        }
    }

    /// Toggle the `X-PJS-Compression` advertisement header.
    pub fn with_compression(mut self, enabled: bool) -> Self {
        self.enable_compression = enabled;
        self
    }

    /// Toggle the `X-PJS-Duration-Ms` and `X-PJS-Version` response headers.
    pub fn with_metrics(mut self, enabled: bool) -> Self {
        self.enable_metrics = enabled;
        self
    }

    /// Set the maximum allowed `Content-Length` for incoming requests, in bytes.
    pub fn with_max_request_size(mut self, size: usize) -> Self {
        self.max_request_size = size;
        self
    }
}

impl Default for PjsMiddleware {
    fn default() -> Self {
        Self::new()
    }
}

impl<S> Layer<S> for PjsMiddleware {
    type Service = PjsMiddlewareService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        PjsMiddlewareService {
            inner,
            config: self.clone(),
        }
    }
}

/// Tower service produced by [`PjsMiddleware`].
#[derive(Clone)]
pub struct PjsMiddlewareService<S> {
    inner: S,
    config: PjsMiddleware,
}

impl<S> Service<Request> for PjsMiddlewareService<S>
where
    S: Service<Request, Response = Response> + Clone + Send + 'static,
    S::Future: Send + 'static,
{
    type Response = Response;
    type Error = S::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, request: Request) -> Self::Future {
        let mut inner = self.inner.clone();
        let config = self.config.clone();

        Box::pin(async move {
            let start_time = Instant::now();

            // Check request size
            if let Some(content_length) = request.headers().get(header::CONTENT_LENGTH)
                && let Ok(length_str) = content_length.to_str()
                && let Ok(length) = length_str.parse::<usize>()
                && length > config.max_request_size
            {
                return Ok(Response::builder()
                    .status(StatusCode::PAYLOAD_TOO_LARGE)
                    .body("Request too large".into())
                    .map_err(|_| Response::new("Failed to build error response".into()))
                    .unwrap_or_else(|err_response| err_response));
            }

            // Process request
            let mut response = inner.call(request).await?;

            // Add performance headers
            if config.enable_metrics {
                let duration = start_time.elapsed();
                if let Ok(duration_value) = HeaderValue::from_str(&duration.as_millis().to_string())
                {
                    response
                        .headers_mut()
                        .insert("X-PJS-Duration-Ms", duration_value);
                }

                let version_value = HeaderValue::from_static(env!("CARGO_PKG_VERSION"));
                response
                    .headers_mut()
                    .insert("X-PJS-Version", version_value);
            }

            // Add compression hints
            if config.enable_compression {
                response
                    .headers_mut()
                    .insert("X-PJS-Compression", HeaderValue::from_static("available"));
            }

            Ok(response)
        })
    }
}

/// Opt-in trusted-proxy allowlist for the HTTP rate limiter.
///
/// By default, [`RateLimitMiddleware`] keys rate limiting on the real TCP peer
/// address and never trusts `X-Forwarded-For`/`X-Real-IP` — an unauthenticated
/// client could otherwise send a fresh spoofed value on every request to get a
/// fresh rate-limit bucket, fully bypassing the limiter. Set this only for
/// deployments that sit behind a known reverse proxy or load balancer whose
/// peer address(es) are listed here; requests from any other peer always use
/// the real peer address regardless of these headers.
///
/// # Proxy contract
///
/// `X-Forwarded-For` is read right-to-left and takes precedence over
/// `X-Real-IP` when both are present. The trusted proxy must *append* the
/// address it saw the connection from to `X-Forwarded-For` rather than
/// overwrite it (e.g. nginx's `$proxy_add_x_forwarded_for`, or any proxy that
/// merges into a single header line rather than emitting a new one). Proxies
/// that instead emit `<ip>:<port>` or bracketed IPv6 entries are not
/// supported by this simple allowlist — the walk fails closed on the first
/// unparseable entry (falls back to `X-Real-IP`, then the peer address)
/// rather than skipping it and guessing from what remains. Repeated
/// `X-Forwarded-For` header lines are read and treated as one comma-joined
/// list in line order, per RFC 9110.
#[derive(Debug, Clone, Default)]
pub struct TrustedProxyConfig {
    /// Peer addresses (the proxy's own TCP source address) allowed to supply
    /// `X-Forwarded-For`/`X-Real-IP`.
    pub trusted_proxies: Vec<std::net::IpAddr>,
}

impl TrustedProxyConfig {
    /// Build a trusted-proxy config from an explicit allowlist of proxy addresses.
    pub fn new(trusted_proxies: Vec<std::net::IpAddr>) -> Self {
        Self { trusted_proxies }
    }

    /// Whether `ip` (already canonicalized via [`IpAddr::to_canonical`]) is in
    /// the allowlist. Allowlist entries are canonicalized before comparison so
    /// an IPv4 proxy configured as `10.0.0.1` still matches when it arrives as
    /// the IPv4-mapped IPv6 address `::ffff:10.0.0.1` on a dual-stack listener.
    fn contains(&self, ip: std::net::IpAddr) -> bool {
        self.trusted_proxies.iter().any(|p| p.to_canonical() == ip)
    }
}

/// Rate limiting configuration for HTTP endpoints
#[derive(Debug, Clone)]
pub struct RateLimitConfig {
    /// Maximum requests per time window (default: 100)
    pub max_requests_per_window: u32,
    /// Time window duration (default: 60 seconds)
    pub window_duration: std::time::Duration,
    /// Opt-in trusted-proxy allowlist. `None` (the default) always keys the
    /// rate limiter on the real TCP peer address. See [`TrustedProxyConfig`].
    pub trusted_proxies: Option<TrustedProxyConfig>,
}

impl Default for RateLimitConfig {
    fn default() -> Self {
        Self {
            max_requests_per_window: 100,
            window_duration: std::time::Duration::from_secs(60),
            trusted_proxies: None,
        }
    }
}

impl RateLimitConfig {
    /// Build a per-minute rate limit (`requests_per_minute` requests per 60-second window).
    pub fn new(requests_per_minute: u32) -> Self {
        Self {
            max_requests_per_window: requests_per_minute,
            window_duration: std::time::Duration::from_secs(60),
            trusted_proxies: None,
        }
    }

    /// Override the window duration that `max_requests_per_window` applies to.
    pub fn with_window(mut self, duration: std::time::Duration) -> Self {
        self.window_duration = duration;
        self
    }

    /// Opt in to trusting `X-Forwarded-For`/`X-Real-IP` from the given proxy allowlist.
    pub fn with_trusted_proxies(mut self, config: TrustedProxyConfig) -> Self {
        self.trusted_proxies = Some(config);
        self
    }
}

/// Rate limiting middleware for PJS endpoints
///
/// Uses token bucket algorithm from security::rate_limit module
/// Returns 429 Too Many Requests when limit exceeded
/// Adds X-RateLimit-* headers per RFC 6585
#[derive(Clone)]
pub struct RateLimitMiddleware {
    limiter: std::sync::Arc<crate::security::rate_limit::WebSocketRateLimiter>,
    trusted_proxies: Option<TrustedProxyConfig>,
}

impl RateLimitMiddleware {
    /// Build a fresh middleware with its own internal `WebSocketRateLimiter`.
    ///
    /// Spawns a background task that periodically prunes expired per-IP
    /// entries (see [`crate::security::rate_limit::WebSocketRateLimiter::spawn_cleanup_task`])
    /// if called from within a Tokio runtime; otherwise construction still
    /// succeeds, but periodic pruning is skipped with a logged warning.
    pub fn new(config: RateLimitConfig) -> Self {
        let trusted_proxies = config.trusted_proxies.clone();
        let rate_limit_config = crate::security::rate_limit::RateLimitConfig {
            max_requests_per_window: config.max_requests_per_window,
            window_duration: config.window_duration,
            ..Default::default()
        };

        let limiter = std::sync::Arc::new(crate::security::rate_limit::WebSocketRateLimiter::new(
            rate_limit_config,
        ));
        limiter.spawn_cleanup_task(crate::security::rate_limit::DEFAULT_CLEANUP_INTERVAL);

        Self {
            limiter,
            trusted_proxies,
        }
    }

    /// Wrap an externally constructed `WebSocketRateLimiter` (lets several middlewares share state).
    ///
    /// Spawns a background cleanup task via
    /// [`crate::security::rate_limit::WebSocketRateLimiter::spawn_cleanup_task`],
    /// which is a no-op if one is already running for this limiter (e.g.
    /// because another `RateLimitMiddleware` already wraps the same `Arc`).
    pub fn from_limiter(
        limiter: std::sync::Arc<crate::security::rate_limit::WebSocketRateLimiter>,
    ) -> Self {
        limiter.spawn_cleanup_task(crate::security::rate_limit::DEFAULT_CLEANUP_INTERVAL);

        Self {
            limiter,
            trusted_proxies: None,
        }
    }

    /// Opt in to trusting `X-Forwarded-For`/`X-Real-IP` from the given proxy allowlist.
    pub fn with_trusted_proxies(mut self, config: TrustedProxyConfig) -> Self {
        self.trusted_proxies = Some(config);
        self
    }
}

impl<S> Layer<S> for RateLimitMiddleware {
    type Service = RateLimitService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        RateLimitService {
            inner,
            limiter: self.limiter.clone(),
            trusted_proxies: self.trusted_proxies.clone(),
        }
    }
}

/// Tower service produced by [`RateLimitMiddleware`].
#[derive(Clone)]
pub struct RateLimitService<S> {
    inner: S,
    limiter: std::sync::Arc<crate::security::rate_limit::WebSocketRateLimiter>,
    trusted_proxies: Option<TrustedProxyConfig>,
}

impl<S> Service<Request> for RateLimitService<S>
where
    S: Service<Request, Response = Response> + Clone + Send + 'static,
    S::Future: Send + 'static,
{
    type Response = Response;
    type Error = S::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, request: Request) -> Self::Future {
        let limiter = self.limiter.clone();
        let trusted_proxies = self.trusted_proxies.clone();
        let mut inner = self.inner.clone();

        Box::pin(async move {
            let client_ip = extract_client_ip(&request, trusted_proxies.as_ref());

            // Check rate limit
            match limiter.check_request(client_ip) {
                Ok(()) => {
                    // Rate limit passed - process request
                    let response = inner.call(request).await?;

                    // Add rate limit headers to response
                    let mut response = response;
                    add_rate_limit_headers(&mut response, &limiter, client_ip);

                    Ok(response)
                }
                Err(err) => {
                    let (status, retry_after, error_label) = rate_limit_error_response_parts(&err);

                    let error_body = serde_json::json!({
                        "error": error_label,
                        "message": err.to_string(),
                        "retry_after": retry_after
                    })
                    .to_string();

                    let mut response = Response::builder()
                        .status(status)
                        .header(header::CONTENT_TYPE, "application/json")
                        .header("Retry-After", retry_after.to_string())
                        .body(error_body.into())
                        .unwrap_or_else(|_| Response::new(error_label.into()));

                    // A `CapacityExceeded` rejection has no per-client bucket
                    // to describe — the IP was never admitted into the
                    // tracked-client table — so the X-RateLimit-* quota
                    // headers would misleadingly describe a bucket that
                    // doesn't exist. Only attach them for per-client
                    // rejections, where they're meaningful.
                    if status != StatusCode::SERVICE_UNAVAILABLE {
                        add_rate_limit_headers(&mut response, &limiter, client_ip);
                    }

                    Ok(response)
                }
            }
        })
    }
}

/// Map a [`crate::security::rate_limit::RateLimitError`] to the `(status,
/// retry_after_secs, error_label)` triple used to build the 429/503 response.
///
/// A `CapacityExceeded` rejection is a server-side condition (the
/// tracked-client table is full) rather than "you exceeded your own quota" —
/// folding it into 429 would mislead the caller into backing off their own
/// request rate, which does nothing to free capacity. It gets `503` instead,
/// with a `Retry-After` tied to the cleanup sweep interval since that's the
/// only thing that can free a slot. See `MAX_TRACKED_CLIENTS`'s docs in
/// `security::rate_limit` for the accepted reject-new-clients tradeoff this
/// implies under a sustained capacity attack. Every other variant keeps the
/// existing `429` + fixed 60s hint.
fn rate_limit_error_response_parts(
    err: &crate::security::rate_limit::RateLimitError,
) -> (StatusCode, u64, &'static str) {
    if matches!(
        err,
        crate::security::rate_limit::RateLimitError::CapacityExceeded { .. }
    ) {
        (
            StatusCode::SERVICE_UNAVAILABLE,
            crate::security::rate_limit::DEFAULT_CLEANUP_INTERVAL.as_secs(),
            "Service Unavailable",
        )
    } else {
        (StatusCode::TOO_MANY_REQUESTS, 60, "Too Many Requests")
    }
}

/// Extract the client IP address used as the rate-limit key.
///
/// Always trusts the real TCP peer address first, populated via axum's
/// [`ConnectInfo`] extension — the router must be served with
/// `into_make_service_with_connect_info::<SocketAddr>()`, otherwise no
/// `ConnectInfo` extension is present, every client collapses onto a single
/// shared bucket keyed on `127.0.0.1`, and a one-time warning is logged (see
/// [`warn_missing_connect_info`]).
///
/// `X-Forwarded-For`/`X-Real-IP` are only consulted when `trusted_proxies` is
/// set and the real peer address is in its allowlist. Trusting these headers
/// unconditionally would let any client forge a fresh rate-limit bucket on
/// every request, fully bypassing the limiter. Both the peer address and
/// allowlist entries are compared via [`IpAddr::to_canonical`] so an IPv4
/// proxy is still recognized when it arrives IPv4-mapped on a dual-stack
/// listener.
fn extract_client_ip(
    request: &Request,
    trusted_proxies: Option<&TrustedProxyConfig>,
) -> std::net::IpAddr {
    use std::net::{IpAddr, Ipv4Addr};

    let Some(peer) = request
        .extensions()
        .get::<ConnectInfo<SocketAddr>>()
        .map(|ConnectInfo(addr)| addr.ip().to_canonical())
    else {
        warn_missing_connect_info();
        return IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
    };

    if let Some(proxies) = trusted_proxies
        && proxies.contains(peer)
        && let Some(forwarded_ip) = extract_forwarded_ip(request.headers(), proxies)
    {
        return forwarded_ip;
    }

    peer
}

/// Log once (per process) that a request arrived with no `ConnectInfo`
/// extension, so misconfiguration is loud rather than silently collapsing
/// every client onto one rate-limit bucket.
fn warn_missing_connect_info() {
    static WARNED: std::sync::Once = std::sync::Once::new();
    WARNED.call_once(|| {
        tracing::warn!(
            "RateLimitMiddleware: request has no ConnectInfo<SocketAddr> extension; \
             serve the router with `.into_make_service_with_connect_info::<SocketAddr>()` \
             or every client will share a single rate-limit bucket keyed on 127.0.0.1 \
             (logged once)"
        );
    });
}

/// Parse the client IP from `X-Forwarded-For` or `X-Real-IP`.
///
/// Only called for peers already verified against [`TrustedProxyConfig`] —
/// these headers must never be trusted from an unverified peer.
///
/// `X-Forwarded-For` is walked right-to-left — across all header lines with
/// that name, since `HeaderMap` allows repeats and they are semantically one
/// comma-joined list in line order — skipping any entry that is itself a
/// trusted proxy, and returns the first entry that is not. A well-behaved
/// proxy *appends* the address it saw the connection from, so the rightmost
/// non-trusted entry is the one appended by the closest trusted hop and
/// cannot be forged by the client — taking the leftmost (client-supplied)
/// entry instead would let a client behind a trusted proxy forge a fresh
/// value on every request and reopen the exact bypass this module exists to
/// close.
///
/// The walk **fails closed** on the first unparseable entry: it stops and
/// falls through to `X-Real-IP` rather than skipping past the malformed
/// entry into entries further left, which are progressively more
/// attacker-controlled the further left they sit in the chain.
fn extract_forwarded_ip(
    headers: &HeaderMap,
    trusted_proxies: &TrustedProxyConfig,
) -> Option<std::net::IpAddr> {
    let entries: Vec<&str> = headers
        .get_all("x-forwarded-for")
        .iter()
        .filter_map(|h| h.to_str().ok())
        .flat_map(|s| s.split(','))
        .collect();

    for entry in entries.into_iter().rev() {
        let Ok(ip) = entry.trim().parse::<std::net::IpAddr>() else {
            break;
        };
        let canonical = ip.to_canonical();
        if !trusted_proxies.contains(canonical) {
            return Some(canonical);
        }
    }

    headers
        .get("x-real-ip")
        .and_then(|h| h.to_str().ok())
        .and_then(|s| s.trim().parse::<std::net::IpAddr>().ok())
        .map(|ip| ip.to_canonical())
}

/// Add X-RateLimit-* headers to response per RFC 6585
fn add_rate_limit_headers(
    response: &mut Response,
    limiter: &crate::security::rate_limit::WebSocketRateLimiter,
    client_ip: std::net::IpAddr,
) {
    use std::time::SystemTime;

    // Get stats for the client (we'll need to access internals or add a method)
    // For now, add standard headers with static values
    // TODO: Add method to WebSocketRateLimiter to get current limit status

    response
        .headers_mut()
        .insert("X-RateLimit-Limit", HeaderValue::from_static("100"));

    // Calculate remaining requests (simplified - would need access to client state)
    response
        .headers_mut()
        .insert("X-RateLimit-Remaining", HeaderValue::from_static("99"));

    // Calculate reset time (current time + 60 seconds)
    if let Some(reset_value) = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .ok()
        .map(|d| d.as_secs() + 60)
        .and_then(|time| HeaderValue::from_str(&time.to_string()).ok())
    {
        response
            .headers_mut()
            .insert("X-RateLimit-Reset", reset_value);
    }

    // Suppress unused variable warning
    let _ = (limiter, client_ip);
}

/// Connection upgrade middleware for WebSocket support
pub async fn websocket_upgrade_middleware(
    headers: HeaderMap,
    request: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    // Check if this is a WebSocket upgrade request
    if headers
        .get(header::UPGRADE)
        .and_then(|h| h.to_str().ok())
        .map(|s| s.to_lowercase())
        == Some("websocket".to_string())
    {
        // Handle WebSocket upgrade for PJS streaming
        // This would integrate with the WebSocket handler
        return handle_websocket_upgrade(request).await;
    }

    // Continue with regular HTTP handling
    Ok(next.run(request).await)
}

/// Handle WebSocket upgrade for real-time PJS streaming
async fn handle_websocket_upgrade(_request: Request) -> Result<Response, StatusCode> {
    // Placeholder - would implement actual WebSocket upgrade logic
    // using axum-websocket or similar
    Response::builder()
        .status(StatusCode::NOT_IMPLEMENTED)
        .body("WebSocket support coming soon".into())
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
}

/// Compression middleware for reducing bandwidth
pub async fn compression_middleware(headers: HeaderMap, request: Request, next: Next) -> Response {
    let accepts_compression = headers
        .get(header::ACCEPT_ENCODING)
        .and_then(|h| h.to_str().ok())
        .map(|s| s.contains("gzip") || s.contains("deflate"))
        .unwrap_or(false);

    let mut response = next.run(request).await;

    // Add compression headers if client supports it
    if accepts_compression {
        response.headers_mut().insert(
            "X-PJS-Compression-Available",
            HeaderValue::from_static("gzip,deflate"),
        );

        // In production, would apply actual compression here
        // using tower-http::compression::CompressionLayer
    }

    response
}

/// CORS middleware specifically configured for PJS streaming
pub async fn pjs_cors_middleware(request: Request, next: Next) -> Response {
    let mut response = next.run(request).await;

    // Add CORS headers for streaming endpoints
    let headers = response.headers_mut();
    headers.insert(
        header::ACCESS_CONTROL_ALLOW_ORIGIN,
        HeaderValue::from_static("*"),
    );
    headers.insert(
        header::ACCESS_CONTROL_ALLOW_METHODS,
        HeaderValue::from_static("GET,POST,OPTIONS"),
    );
    headers.insert(
        header::ACCESS_CONTROL_ALLOW_HEADERS,
        HeaderValue::from_static("Content-Type,Authorization,X-PJS-Priority,X-PJS-Format"),
    );
    headers.insert(
        header::ACCESS_CONTROL_EXPOSE_HEADERS,
        HeaderValue::from_static("X-PJS-Duration-Ms,X-PJS-Version,X-PJS-Stream-Id"),
    );

    response
}

/// Security middleware for PJS endpoints
pub async fn security_middleware(request: Request, next: Next) -> Response {
    let mut response = next.run(request).await;

    // Add security headers
    let headers = response.headers_mut();
    headers.insert(
        "X-Content-Type-Options",
        HeaderValue::from_static("nosniff"),
    );
    headers.insert("X-Frame-Options", HeaderValue::from_static("DENY"));
    headers.insert(
        "Content-Security-Policy",
        HeaderValue::from_static("default-src 'self'"),
    );

    response
}

/// Circuit breaker middleware for resilience
#[derive(Clone)]
pub struct CircuitBreakerMiddleware {
    failure_threshold: usize,
    recovery_timeout_seconds: u64,
}

impl CircuitBreakerMiddleware {
    /// Build with default thresholds (5 failures, 30-second recovery).
    pub fn new() -> Self {
        Self {
            failure_threshold: 5,
            recovery_timeout_seconds: 30,
        }
    }

    /// Override the consecutive-failure threshold that opens the circuit.
    pub fn with_failure_threshold(mut self, threshold: usize) -> Self {
        self.failure_threshold = threshold;
        self
    }

    /// Override the recovery (cool-down) duration in seconds.
    pub fn with_recovery_timeout(mut self, seconds: u64) -> Self {
        self.recovery_timeout_seconds = seconds;
        self
    }
}

impl Default for CircuitBreakerMiddleware {
    fn default() -> Self {
        Self::new()
    }
}

/// Health check middleware that monitors PJS service health
pub async fn health_check_middleware(request: Request, next: Next) -> Response {
    // Add health metrics to response headers
    let mut response = next.run(request).await;

    // In production, would check actual service health
    response
        .headers_mut()
        .insert("X-PJS-Health", HeaderValue::from_static("healthy"));

    response
}

/// Content validation middleware configuration
#[derive(Debug, Clone)]
pub struct ContentValidationConfig {
    /// Maximum allowed Content-Length in bytes (default: 10MB)
    pub max_content_length: usize,

    /// Allowed Content-Type values (default: application/json, application/pjs+json)
    pub allowed_content_types: Vec<String>,

    /// Require Content-Type header for POST/PUT/PATCH (default: true)
    pub require_content_type: bool,
}

impl Default for ContentValidationConfig {
    fn default() -> Self {
        Self {
            max_content_length: 10 * 1024 * 1024, // 10MB
            allowed_content_types: vec![
                "application/json".to_string(),
                "application/pjs+json".to_string(),
            ],
            require_content_type: true,
        }
    }
}

/// Content validation middleware handler
///
/// Validates Content-Type and Content-Length headers to prevent:
/// - Unsupported media types (415 error)
/// - Oversized payloads (413 error)
/// - DoS attacks via malformed headers
pub async fn content_validation_middleware(
    config: ContentValidationConfig,
    req: Request,
    next: Next,
) -> Response {
    // Extract method and headers
    let method = req.method().clone();
    let headers = req.headers();

    // Validate Content-Length
    if let Some(content_length_header) = headers.get(header::CONTENT_LENGTH) {
        match content_length_header.to_str() {
            Ok(content_length_str) => match content_length_str.parse::<usize>() {
                Ok(content_length) => {
                    if content_length > config.max_content_length {
                        let error_body = serde_json::json!({
                            "error": "Payload Too Large",
                            "max_size": config.max_content_length,
                            "received_size": content_length
                        })
                        .to_string();

                        return Response::builder()
                            .status(StatusCode::PAYLOAD_TOO_LARGE)
                            .header(header::CONTENT_TYPE, "application/json")
                            .body(error_body.into())
                            .unwrap_or_else(|_| Response::new("Payload Too Large".into()));
                    }
                }
                Err(_) => {
                    let error_body = serde_json::json!({
                        "error": "Invalid Content-Length header"
                    })
                    .to_string();

                    return Response::builder()
                        .status(StatusCode::BAD_REQUEST)
                        .header(header::CONTENT_TYPE, "application/json")
                        .body(error_body.into())
                        .unwrap_or_else(|_| Response::new("Bad Request".into()));
                }
            },
            Err(_) => {
                let error_body = serde_json::json!({
                    "error": "Invalid Content-Length header encoding"
                })
                .to_string();

                return Response::builder()
                    .status(StatusCode::BAD_REQUEST)
                    .header(header::CONTENT_TYPE, "application/json")
                    .body(error_body.into())
                    .unwrap_or_else(|_| Response::new("Bad Request".into()));
            }
        }
    }

    // Validate Content-Type for POST/PUT/PATCH requests
    if config.require_content_type && (method == "POST" || method == "PUT" || method == "PATCH") {
        match headers.get(header::CONTENT_TYPE) {
            Some(content_type_header) => {
                let content_type = content_type_header.to_str().unwrap_or("");

                // Extract base content type (ignore charset and other parameters)
                let base_content_type = content_type.split(';').next().unwrap_or("").trim();

                if !config
                    .allowed_content_types
                    .iter()
                    .any(|allowed| base_content_type.eq_ignore_ascii_case(allowed))
                {
                    let error_body = serde_json::json!({
                        "error": "Unsupported Media Type",
                        "accepted": config.allowed_content_types,
                        "received": content_type
                    })
                    .to_string();

                    return Response::builder()
                        .status(StatusCode::UNSUPPORTED_MEDIA_TYPE)
                        .header(header::CONTENT_TYPE, "application/json")
                        .body(error_body.into())
                        .unwrap_or_else(|_| Response::new("Unsupported Media Type".into()));
                }
            }
            None => {
                let error_body = serde_json::json!({
                    "error": "Unsupported Media Type",
                    "message": "Content-Type header is required for POST/PUT/PATCH requests",
                    "accepted": config.allowed_content_types
                })
                .to_string();

                return Response::builder()
                    .status(StatusCode::UNSUPPORTED_MEDIA_TYPE)
                    .header(header::CONTENT_TYPE, "application/json")
                    .body(error_body.into())
                    .unwrap_or_else(|_| Response::new("Unsupported Media Type".into()));
            }
        }
    }

    // All validations passed, continue to next middleware/handler
    next.run(req).await
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_pjs_middleware_creation() {
        let middleware = PjsMiddleware::new()
            .with_compression(true)
            .with_metrics(true)
            .with_max_request_size(5 * 1024 * 1024);

        assert!(middleware.enable_compression);
        assert!(middleware.enable_metrics);
        assert_eq!(middleware.max_request_size, 5 * 1024 * 1024);
    }

    #[test]
    fn test_rate_limit_config_default() {
        let config = RateLimitConfig::default();
        assert_eq!(config.max_requests_per_window, 100);
        assert_eq!(config.window_duration, std::time::Duration::from_secs(60));
    }

    #[test]
    fn test_rate_limit_config_new() {
        let config = RateLimitConfig::new(50);
        assert_eq!(config.max_requests_per_window, 50);
    }

    #[test]
    fn test_rate_limit_config_with_window() {
        let config = RateLimitConfig::new(100).with_window(std::time::Duration::from_secs(30));
        assert_eq!(config.window_duration, std::time::Duration::from_secs(30));
    }

    #[tokio::test]
    async fn test_rate_limit_middleware_creation() {
        let config = RateLimitConfig::default();
        let _middleware = RateLimitMiddleware::new(config);
    }

    #[tokio::test]
    async fn test_from_limiter_claims_cleanup_spawn() {
        // `RateLimitMiddleware::new`/`from_limiter` always spawn with the
        // production `DEFAULT_CLEANUP_INTERVAL` (300s), which is too slow to
        // wait out in a test — the underlying `spawn_cleanup_task` mechanism
        // (real eviction after a short period, idempotency across repeated
        // calls) is proven directly and quickly in `security::rate_limit`'s
        // own tests. What this test proves instead, deterministically and
        // fast, is that `from_limiter` actually calls it: `from_limiter` is
        // the *first* caller here (not pre-empted by a manual
        // `spawn_cleanup_task` call, which would make this a no-op check).
        let limiter = std::sync::Arc::new(crate::security::rate_limit::WebSocketRateLimiter::new(
            crate::security::rate_limit::RateLimitConfig::default(),
        ));
        assert!(!limiter.is_cleanup_task_spawned());

        let _middleware = RateLimitMiddleware::from_limiter(limiter.clone());

        assert!(
            limiter.is_cleanup_task_spawned(),
            "RateLimitMiddleware::from_limiter must wire up periodic cleanup"
        );
    }

    #[tokio::test]
    async fn test_new_claims_cleanup_spawn() {
        let middleware = RateLimitMiddleware::new(RateLimitConfig::default());

        assert!(
            middleware.limiter.is_cleanup_task_spawned(),
            "RateLimitMiddleware::new must wire up periodic cleanup"
        );
    }

    #[test]
    fn test_capacity_exceeded_maps_to_503_with_sweep_interval_retry_after() {
        let err = crate::security::rate_limit::RateLimitError::CapacityExceeded {
            max: crate::security::rate_limit::MAX_TRACKED_CLIENTS,
        };
        let (status, retry_after, label) = rate_limit_error_response_parts(&err);

        assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
        assert_eq!(
            retry_after,
            crate::security::rate_limit::DEFAULT_CLEANUP_INTERVAL.as_secs()
        );
        assert_eq!(label, "Service Unavailable");
    }

    #[test]
    fn test_other_rate_limit_errors_map_to_429() {
        let err = crate::security::rate_limit::RateLimitError::LimitExceeded {
            limit: 10,
            window: std::time::Duration::from_secs(60),
        };
        let (status, retry_after, label) = rate_limit_error_response_parts(&err);

        assert_eq!(status, StatusCode::TOO_MANY_REQUESTS);
        assert_eq!(retry_after, 60);
        assert_eq!(label, "Too Many Requests");
    }

    #[tokio::test]
    async fn test_from_limiter_on_already_spawned_limiter_does_not_reclaim() {
        // Simulates the scenario `from_limiter`'s docs describe: a limiter
        // whose cleanup was already spawned elsewhere (e.g. by
        // `SecureWebSocketHandler::new` sharing the same `Arc`). Wrapping it
        // again must observe the claim as already made, not attempt (or
        // need) a second spawn.
        let limiter = std::sync::Arc::new(crate::security::rate_limit::WebSocketRateLimiter::new(
            crate::security::rate_limit::RateLimitConfig::default(),
        ));
        limiter.spawn_cleanup_task(crate::security::rate_limit::DEFAULT_CLEANUP_INTERVAL);
        assert!(limiter.is_cleanup_task_spawned());

        let _middleware = RateLimitMiddleware::from_limiter(limiter.clone());

        assert!(limiter.is_cleanup_task_spawned());
    }

    #[test]
    fn test_content_validation_config_default() {
        let config = ContentValidationConfig::default();
        assert_eq!(config.max_content_length, 10 * 1024 * 1024);
        assert_eq!(config.allowed_content_types.len(), 2);
        assert!(config.require_content_type);
    }
}