1use std::sync::Arc;
2
3use anyhow::Result;
4use bytes::Bytes;
5use tokio::sync::Mutex;
6use tokio::time::{
7 Duration,
8 Instant,
9};
10use tracing::{
11 debug,
12 error,
13 trace,
14};
15
16fn find_headers_end(data: &[u8]) -> Option<usize> {
17 data.windows(4).position(|w| w == b"\r\n\r\n")
18}
19
20use crate::HttpLogger;
21use crate::http_response_analyzer::{
22 HttpResponseAnalyzer,
23 ResponseAnalyzerConfig,
24};
25
26pub const DEFAULT_MIN_LOG_SYNC_MS: u64 = 50;
27
28pub struct ResponseLoggingState {
29 pub complete_response: Vec<u8>,
30 pub already_logged: bool,
31 pub logging_enabled: bool,
32 pub is_chunked: bool,
33}
34
35impl Default for ResponseLoggingState {
36 fn default() -> Self {
37 Self::new()
38 }
39}
40
41impl ResponseLoggingState {
42 pub fn new() -> Self {
43 Self {
44 complete_response: Vec::new(),
45 already_logged: false,
46 logging_enabled: false,
47 is_chunked: false,
48 }
49 }
50}
51
52pub struct ResponseChunkContext {
53 pub complete_response: Vec<u8>,
54 pub is_chunked: bool,
55 pub found_end_marker: bool,
56 pub total_chunks_received: usize,
57 pub response_logged: Arc<Mutex<bool>>,
58 pub request_id: Arc<Mutex<Option<String>>>,
59}
60
61impl ResponseChunkContext {
62 pub fn new(response_logged: Arc<Mutex<bool>>, request_id: Arc<Mutex<Option<String>>>) -> Self {
63 Self {
64 complete_response: Vec::new(),
65 is_chunked: false,
66 found_end_marker: false,
67 total_chunks_received: 0,
68 response_logged,
69 request_id,
70 }
71 }
72}
73
74pub struct ResponseLoggingContext {
75 pub complete_response: Vec<u8>,
76 pub is_chunked: bool,
77 pub found_end_marker: bool,
78 pub response_logged: Arc<Mutex<bool>>,
79 pub request_id: Arc<Mutex<Option<String>>>,
80 pub first_chunk_time: Option<Instant>,
81}
82
83#[derive(Debug, Clone)]
84pub struct ResponseHandlerConfig {
85 pub analyzer_config: ResponseAnalyzerConfig,
86 pub min_log_sync_ms: u64,
87}
88
89impl Default for ResponseHandlerConfig {
90 fn default() -> Self {
91 Self {
92 analyzer_config: ResponseAnalyzerConfig::default(),
93 min_log_sync_ms: DEFAULT_MIN_LOG_SYNC_MS,
94 }
95 }
96}
97
98#[derive(Debug, Clone)]
99pub struct HttpResponseHandler {
100 config_id: i64,
101 analyzer: HttpResponseAnalyzer,
102 config: ResponseHandlerConfig,
103}
104
105impl HttpResponseHandler {
106 pub fn new(config_id: i64) -> Self {
107 Self::with_config(config_id, ResponseHandlerConfig::default())
108 }
109
110 pub fn with_config(config_id: i64, config: ResponseHandlerConfig) -> Self {
111 Self {
112 config_id,
113 analyzer: HttpResponseAnalyzer::new(config.analyzer_config.clone()),
114 config,
115 }
116 }
117
118 pub fn config(&self) -> &ResponseHandlerConfig {
119 &self.config
120 }
121
122 pub async fn check_response_logging_status(
123 &self, buffer: &[u8], n: usize, state: &mut ResponseLoggingState,
124 ) -> Result<()> {
125 let is_enabled =
126 match kftray_commons::utils::http_logs_config::get_http_logs_config(self.config_id)
127 .await
128 {
129 Ok(config) => config.enabled,
130 Err(_) => false,
131 };
132
133 if is_enabled {
134 state.logging_enabled = true;
135 state.already_logged = true;
136
137 if n > 0 {
138 state.complete_response.extend_from_slice(&buffer[..n]);
139
140 if let Some(headers_end) = find_headers_end(&state.complete_response) {
141 let headers_data = &state.complete_response[..headers_end];
142 state.is_chunked = Self::is_chunked_transfer(headers_data);
143 }
144 }
145 } else {
146 state.already_logged = true;
147 }
148
149 Ok(())
150 }
151
152 pub async fn process_response_chunk(
153 &self, buffer: &[u8], n: usize, context: &mut ResponseChunkContext,
154 logger: &Option<HttpLogger>,
155 ) -> Result<()> {
156 context.complete_response.extend_from_slice(&buffer[..n]);
157 context.total_chunks_received += 1;
158
159 let response_size = context.complete_response.len();
160 trace!(
161 "Processing response chunk #{}: {}B, total: {}B",
162 context.total_chunks_received, n, response_size
163 );
164
165 if !context.is_chunked
166 && let Some(headers_end) = find_headers_end(&context.complete_response)
167 {
168 let headers_data = &context.complete_response[..headers_end];
169 context.is_chunked = Self::is_chunked_transfer(headers_data);
170 }
171
172 if context.is_chunked && !context.found_end_marker {
173 context.found_end_marker =
174 HttpResponseAnalyzer::has_chunked_end_marker(&context.complete_response);
175 if context.found_end_marker {
176 debug!(
177 "Found end marker in chunked response after {} chunks, {}B",
178 context.total_chunks_received, response_size
179 );
180
181 let mut logging_context = ResponseLoggingContext {
182 complete_response: std::mem::take(&mut context.complete_response),
183 is_chunked: context.is_chunked,
184 found_end_marker: context.found_end_marker,
185 response_logged: context.response_logged.clone(),
186 request_id: context.request_id.clone(),
187 first_chunk_time: None,
188 };
189
190 let result = self
191 .check_and_log_complete_response(&mut logging_context, logger)
192 .await;
193
194 context.complete_response = logging_context.complete_response;
196
197 result?;
198 }
199 }
200
201 Ok(())
202 }
203
204 pub async fn check_and_log_complete_response(
205 &self, context: &mut ResponseLoggingContext, logger: &Option<HttpLogger>,
206 ) -> Result<()> {
207 let mut response_logged_guard = context.response_logged.lock().await;
208
209 if !context.complete_response.is_empty() && !*response_logged_guard {
210 let should_log = if context.is_chunked {
211 context.found_end_marker
212 || self.analyzer.is_ready_for_logging_with_config(
213 &context.complete_response,
214 context.is_chunked,
215 context.found_end_marker,
216 )
217 } else if let Some(start_time) = context.first_chunk_time {
218 self.check_time_based_logging(&context.complete_response, start_time)
219 } else {
220 self.analyzer.appears_complete(
221 &context.complete_response,
222 context.is_chunked,
223 context.found_end_marker,
224 )
225 };
226
227 if should_log {
228 debug!(
229 "Logging complete response: {}B, chunked: {}, end marker: {}",
230 context.complete_response.len(),
231 context.is_chunked,
232 context.found_end_marker
233 );
234
235 self.log_response(
236 &mut context.complete_response,
237 &context.response_logged,
238 &context.request_id,
239 logger,
240 )
241 .await?;
242
243 *response_logged_guard = true;
244 }
245 }
246
247 Ok(())
248 }
249
250 fn check_time_based_logging(&self, response_data: &[u8], start_time: Instant) -> bool {
251 if let Some(headers_end) = find_headers_end(response_data) {
252 let headers = &response_data[..headers_end];
253
254 if response_data.len() > headers_end + 4
255 && let Ok(headers_str) = std::str::from_utf8(headers)
256 {
257 let h_lower = headers_str.to_lowercase();
258 if h_lower.contains("upgrade: websocket") && h_lower.contains("connection: upgrade")
259 {
260 if h_lower.contains("sec-websocket-accept:") {
261 debug!(
262 "Logging WebSocket upgrade response immediately - complete with all required headers"
263 );
264 return true;
265 } else {
266 debug!(
267 "Potential WebSocket upgrade response detected but missing accept header"
268 );
269 }
270 }
271 }
272
273 if let Ok(headers_str) = std::str::from_utf8(headers) {
274 let first_line = headers_str.lines().next().unwrap_or("");
275 let parts: Vec<&str> = first_line.split_whitespace().collect();
276 if parts.len() >= 2
277 && let Ok(status) = parts[1].parse::<u16>()
278 && ((100..200).contains(&status) || status == 204 || status == 304)
279 {
280 debug!("Logging status {} response which never has a body", status);
281 return true;
282 }
283 }
284
285 if std::str::from_utf8(headers)
286 .map(|h| h.to_lowercase().contains("connection: close"))
287 .unwrap_or(false)
288 {
289 debug!("Connection: close header found, logging response at connection end");
290 return true;
291 }
292
293 if let Ok(headers_str) = std::str::from_utf8(headers)
294 && headers_str.starts_with("HTTP/1.0")
295 && !headers_str.to_lowercase().contains("content-length:")
296 && !headers_str
297 .to_lowercase()
298 .contains("transfer-encoding: chunked")
299 {
300 debug!("HTTP/1.0 response without length indicators - logging at connection end");
301 return true;
302 }
303
304 let elapsed = start_time.elapsed();
305 if elapsed.as_secs() >= 10 && response_data.len() > 5000 {
306 debug!(
307 "Long-lived connection detected ({}s, {}B) - logging accumulated data to prevent being stuck",
308 elapsed.as_secs(),
309 response_data.len()
310 );
311 return true;
312 }
313
314 if response_data.len() > 1_000_000 {
315 debug!(
316 "Large response detected ({}B) - logging to prevent memory issues",
317 response_data.len()
318 );
319 return true;
320 }
321
322 if response_data.len() > 100_000 && elapsed.as_secs() > 30 {
323 debug!(
324 "Long-running HTTP stream detected ({}s, {}B) - forcing log",
325 elapsed.as_secs(),
326 response_data.len()
327 );
328 return true;
329 }
330 }
331
332 false
333 }
334
335 async fn format_response_for_logging(&self, response_data: &[u8]) -> Result<String> {
336 debug!(
337 "Formatting response for logging, size: {} bytes",
338 response_data.len()
339 );
340
341 if let Some(headers_end) = find_headers_end(response_data) {
342 debug!("Found headers end at position: {}", headers_end);
343
344 let headers_bytes = &response_data[..headers_end];
345 let body_start = headers_end + 4;
346 let body = &response_data[body_start..];
347
348 debug!(
349 "Headers size: {} bytes, Body size: {} bytes",
350 headers_bytes.len(),
351 body.len()
352 );
353
354 let (status, headers) = match crate::parser::ResponseParser::parse(headers_bytes) {
355 Ok((status, headers)) => {
356 debug!(
357 "Successfully parsed headers, status: {:?}, headers count: {}",
358 status,
359 headers.len()
360 );
361 (status, headers)
362 }
363 Err(e) => {
364 error!("Failed to parse headers: {:?}", e);
365 return Err(e);
366 }
367 };
368
369 let mut formatted = String::new();
370 if let Some(_status_code) = status {
371 let status_line = match std::str::from_utf8(
372 &headers_bytes[..headers_bytes
373 .windows(2)
374 .position(|w| w == b"\r\n")
375 .unwrap_or(headers_bytes.len())],
376 ) {
377 Ok(line) => {
378 debug!("Status line: {}", line);
379 line
380 }
381 Err(e) => {
382 error!("Failed to convert status line to UTF-8: {:?}", e);
383 "HTTP/1.1 200 OK"
384 }
385 };
386 formatted.push_str(status_line);
387 formatted.push('\n');
388 } else {
389 debug!("No status line found, adding default");
390 formatted.push_str("HTTP/1.1 200 OK\n");
391 }
392
393 for header in &headers {
394 if let (Ok(name), Ok(value)) = (
395 std::str::from_utf8(header.name.as_bytes()),
396 std::str::from_utf8(header.value),
397 ) {
398 formatted.push_str(&format!("{name}: {value}\n"));
399 }
400 }
401 formatted.push('\n');
402
403 if !body.is_empty() {
404 debug!("Processing response body with size: {} bytes", body.len());
405
406 let is_chunked = crate::parser::RequestParser::is_chunked_transfer(&headers);
407 let is_gzipped = crate::parser::RequestParser::is_gzip_encoded(&headers);
408 let is_brotli = crate::parser::RequestParser::is_brotli_encoded(&headers);
409
410 if is_chunked {
411 debug!("Response is chunked encoded");
412 }
413 if is_gzipped {
414 debug!("Response is gzip compressed");
415 }
416 if is_brotli {
417 debug!("Response is brotli compressed");
418 }
419
420 let processed_body =
421 match crate::parser::BodyParser::process_response_body(body, &headers).await {
422 Ok(processed) => {
423 debug!(
424 "Successfully processed response body: {} bytes -> {} bytes",
425 body.len(),
426 processed.len()
427 );
428 processed
429 }
430 Err(e) => {
431 debug!(
432 "Error processing response body: {:?}, using original body",
433 e
434 );
435 body.to_vec()
436 }
437 };
438
439 let content_type = crate::parser::BodyParser::get_content_type(&headers);
440 debug!("Content type for body formatting: {:?}", content_type);
441
442 let body_formatted = match crate::parser::BodyParser::format_body(
443 &processed_body,
444 content_type,
445 ) {
446 Ok(formatted) => {
447 debug!(
448 "Successfully formatted response body: {} bytes",
449 formatted.len()
450 );
451 formatted
452 }
453 Err(e) => {
454 debug!(
455 "Error formatting response body: {:?}, trying direct string conversion",
456 e
457 );
458
459 if let Ok(text) = std::str::from_utf8(&processed_body) {
460 debug!("Direct string conversion succeeded");
461 text.to_string()
462 } else {
463 debug!("Direct string conversion failed, using lossy conversion");
464 String::from_utf8_lossy(&processed_body).to_string()
465 }
466 }
467 };
468
469 formatted.push_str(&body_formatted);
470 } else {
471 formatted.push_str("#<empty body>");
472 }
473
474 debug!("Formatted response size: {} bytes", formatted.len());
475 Ok(formatted)
476 } else {
477 debug!("Could not find headers end, returning raw data as string");
478 Ok(String::from_utf8_lossy(response_data).to_string())
479 }
480 }
481
482 pub async fn log_response(
483 &self, complete_response: &mut [u8], _response_logged: &Arc<Mutex<bool>>,
484 request_id: &Arc<Mutex<Option<String>>>, logger: &Option<HttpLogger>,
485 ) -> Result<()> {
486 if let Some(logger) = logger {
487 let req_id_guard = request_id.lock().await;
488 if let Some(req_id) = &*req_id_guard {
489 debug!(
490 "Logging response for request ID: {}, response size: {} bytes",
491 req_id,
492 complete_response.len()
493 );
494
495 let formatted_response =
496 match self.format_response_for_logging(complete_response).await {
497 Ok(formatted) => {
498 debug!(
499 "Successfully formatted response for logging, size: {} bytes",
500 formatted.len()
501 );
502
503 if formatted.starts_with("HTTP/") {
504 formatted
505 } else {
506 debug!("Response doesn't start with HTTP/, adding basic header");
507 let mut enhanced = String::from("HTTP/1.1 200 OK\n\n");
508 enhanced.push_str(&formatted);
509 enhanced
510 }
511 }
512 Err(e) => {
513 error!("Failed to format response: {:?}", e);
514 let mut basic_response = String::from("HTTP/1.1 200 OK\n");
515 basic_response.push_str(&format!("X-Formatting-Error: {e}\n\n"));
516 basic_response.push_str("# Failed to format response properly\n");
517 basic_response.push_str(&format!("# Error: {e}\n"));
518 basic_response.push_str("# Raw content follows (first 1000 bytes):\n");
519
520 let preview_size = std::cmp::min(complete_response.len(), 1000);
521 basic_response.push_str(&String::from_utf8_lossy(
522 &complete_response[..preview_size],
523 ));
524 basic_response
525 }
526 };
527
528 let bytes = Bytes::from(formatted_response.into_bytes());
529 debug!("Sending response to logger, size: {} bytes", bytes.len());
530
531 logger.log_response(bytes, req_id.clone()).await;
532 debug!("Response successfully logged for request ID: {}", req_id);
533 } else {
534 error!("No request ID available for logging response");
535 }
536 drop(req_id_guard);
537 } else {
538 debug!("No logger available for logging response");
539 }
540
541 Ok(())
542 }
543
544 pub async fn handle_remaining_response_data(
545 &self, complete_response: &[u8], logging_enabled: bool, response_logged: &Arc<Mutex<bool>>,
546 request_id: &Arc<Mutex<Option<String>>>, logger: &Option<HttpLogger>,
547 ) -> Result<()> {
548 if complete_response.is_empty() {
549 debug!("No response data to log at connection end");
550 return Ok(());
551 }
552
553 let mut response_logged_guard = response_logged.lock().await;
554 let already_logged = *response_logged_guard;
555
556 if logging_enabled && !already_logged {
557 *response_logged_guard = true;
558 debug!(
559 "Marking response as logged (size: {}B)",
560 complete_response.len()
561 );
562 drop(response_logged_guard);
563
564 let req_id_guard = request_id.lock().await;
565 if let Some(req_id) = &*req_id_guard {
566 debug!(
567 "Connection closing with response data (size: {} bytes) for request ID: {}",
568 complete_response.len(),
569 req_id
570 );
571
572 if let Some(log_instance) = logger.as_ref() {
573 let req_clone = req_id.clone();
574 let response_data = Bytes::copy_from_slice(complete_response);
575 let response_size = response_data.len();
576
577 drop(req_id_guard);
578
579 debug!(
580 "Logging HTTP response at connection close - this ensures all data is captured"
581 );
582
583 let timeout_duration = Duration::from_secs(5);
584 match tokio::time::timeout(
585 timeout_duration,
586 log_instance.log_response(response_data, req_clone),
587 )
588 .await
589 {
590 Ok(_) => {
591 debug!(
592 "Final response ({}B) successfully logged (connection end)",
593 response_size
594 );
595 }
596 Err(_) => {
597 error!("Response logging timed out after 5 seconds");
598 }
599 }
600 } else {
601 drop(req_id_guard);
602 error!("No logger instance available for logging response");
603 }
604
605 let wait_time = Duration::from_millis(self.config.min_log_sync_ms * 2);
606 tokio::time::sleep(wait_time).await;
607 debug!("Completed final response logging before connection close");
608 } else {
609 drop(req_id_guard);
610 debug!("No request ID found for response at connection end");
611 }
612 } else {
613 drop(response_logged_guard);
614
615 if logging_enabled && already_logged {
616 debug!("Response already logged earlier - connection closing");
617 } else if !logging_enabled {
618 debug!("HTTP logging disabled for this connection");
619 }
620 }
621
622 Ok(())
623 }
624
625 fn is_chunked_transfer(headers_data: &[u8]) -> bool {
626 if let Ok(headers_str) = std::str::from_utf8(headers_data) {
627 let h_lower = headers_str.to_lowercase();
628 if h_lower.contains("transfer-encoding: chunked") {
629 trace!("Detected chunked transfer encoding in response - will wait for all chunks");
630 return true;
631 }
632
633 if h_lower.contains("upgrade: websocket") && h_lower.contains("connection: upgrade") {
634 debug!("Initial detection of WebSocket upgrade response");
635 return false;
636 }
637 }
638 false
639 }
640}
641
642pub use DEFAULT_MIN_LOG_SYNC_MS as MIN_LOG_SYNC_MS;
643
644#[cfg(test)]
645mod tests {
646 use super::*;
647
648 #[tokio::test]
649 async fn test_check_response_logging_status() {
650 let _db = kftray_commons::test_utils::test_db().await;
651 let handler = HttpResponseHandler::new(123);
652
653 let buffer = b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\ndata";
654 let mut state = ResponseLoggingState::new();
655
656 handler
657 .check_response_logging_status(buffer, buffer.len(), &mut state)
658 .await
659 .unwrap();
660
661 assert!(state.already_logged);
662 assert!(!state.logging_enabled);
664 }
665
666 #[tokio::test]
667 async fn test_custom_config() {
668 let config = ResponseHandlerConfig {
669 min_log_sync_ms: 25,
670 ..Default::default()
671 };
672
673 let handler = HttpResponseHandler::with_config(123, config);
674
675 assert_eq!(handler.config().min_log_sync_ms, 25);
676
677 let response_data = b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\nSome data";
678
679 assert!(handler.check_time_based_logging(response_data, Instant::now()));
680
681 let default_handler = HttpResponseHandler::new(123);
682
683 let regular_response = b"HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\nSome data";
684
685 assert!(!default_handler.check_time_based_logging(regular_response, Instant::now()));
686 }
687
688 #[test]
689 fn test_find_headers_end() {
690 let response_with_headers =
691 b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n\r\nBody content";
692 let response_without_end = b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n";
693
694 let header_end_position = response_with_headers
695 .windows(4)
696 .position(|w| w == b"\r\n\r\n")
697 .unwrap();
698
699 assert_eq!(
700 find_headers_end(response_with_headers),
701 Some(header_end_position)
702 );
703 assert_eq!(find_headers_end(response_without_end), None);
704 }
705
706 #[test]
707 fn test_response_logging_state_new() {
708 let state = ResponseLoggingState::new();
709
710 assert!(state.complete_response.is_empty());
711 assert!(!state.already_logged);
712 assert!(!state.logging_enabled);
713 assert!(!state.is_chunked);
714 }
715
716 #[tokio::test]
717 async fn test_creating_response_chunk_context() {
718 let response_logged = Arc::new(Mutex::new(false));
719 let request_id = Arc::new(Mutex::new(Some("test-id".to_string())));
720
721 let mut context = ResponseChunkContext::new(response_logged, request_id);
722
723 let first_chunk = b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n";
724 context.complete_response.extend_from_slice(first_chunk);
725 context.total_chunks_received += 1;
726
727 assert_eq!(context.total_chunks_received, 1);
728 assert_eq!(context.complete_response, first_chunk);
729 assert!(!context.is_chunked);
730 assert!(!context.found_end_marker);
731 }
732
733 #[test]
734 fn test_handler_with_config() {
735 let default_handler = HttpResponseHandler::new(456);
736 assert_eq!(default_handler.config_id, 456);
737 assert_eq!(
738 default_handler.config.min_log_sync_ms,
739 DEFAULT_MIN_LOG_SYNC_MS
740 );
741
742 let custom_config = ResponseHandlerConfig {
743 min_log_sync_ms: 100,
744 analyzer_config: ResponseAnalyzerConfig {
745 min_headers_size: 32,
746 },
747 };
748
749 let custom_handler = HttpResponseHandler::with_config(789, custom_config);
750 assert_eq!(custom_handler.config_id, 789);
751 assert_eq!(custom_handler.config.min_log_sync_ms, 100);
752 assert_eq!(custom_handler.config.analyzer_config.min_headers_size, 32);
753 }
754
755 #[tokio::test]
756 async fn test_creating_logging_context() {
757 let _handler = HttpResponseHandler::new(123);
758
759 let response_logged = Arc::new(Mutex::new(false));
760 let request_id = Arc::new(Mutex::new(Some("test-id".to_string())));
761
762 let buffer = b"HTTP/1.1 200 OK\r\nContent-Length: 11\r\n\r\nHello World";
763
764 let context = ResponseLoggingContext {
765 complete_response: buffer.to_vec(),
766 is_chunked: false,
767 found_end_marker: false,
768 response_logged,
769 request_id,
770 first_chunk_time: Some(Instant::now()),
771 };
772
773 assert_eq!(context.complete_response, buffer);
774 assert!(!context.is_chunked);
775 assert!(!context.found_end_marker);
776 }
777
778 #[test]
779 fn test_response_handler_config_default() {
780 let config = ResponseHandlerConfig::default();
781
782 assert_eq!(config.min_log_sync_ms, DEFAULT_MIN_LOG_SYNC_MS);
783 assert_eq!(
784 config.analyzer_config.min_headers_size,
785 crate::http_response_analyzer::DEFAULT_MIN_VALID_HEADERS_SIZE
786 );
787 }
788
789 #[tokio::test]
790 async fn test_check_time_based_logging_websocket() {
791 let handler = HttpResponseHandler::new(123);
792
793 let websocket_response = b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n";
794 assert!(handler.check_time_based_logging(websocket_response, Instant::now()));
795 }
796
797 #[tokio::test]
798 async fn test_check_time_based_logging_regular_response_no_time_passed() {
799 let handler = HttpResponseHandler::new(123);
800
801 let regular_response = b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n\r\nContent";
802 let now = Instant::now();
803
804 assert!(!handler.check_time_based_logging(regular_response, now));
805 }
806
807 #[tokio::test]
808 async fn test_check_time_based_logging_regular_response_time_passed() {
809 let handler = HttpResponseHandler::new(123);
810
811 let very_old_time = Instant::now() - Duration::from_secs(15);
812
813 let mut large_response =
814 Vec::from(&b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n\r\n"[..]);
815 large_response.extend(vec![b'A'; 6000]);
816
817 assert!(
818 handler.check_time_based_logging(&large_response, very_old_time),
819 "Time-based logging should be triggered for large response with sufficient time elapsed"
820 );
821 }
822
823 #[tokio::test]
824 async fn test_check_time_based_logging_unfinished_response() {
825 let handler = HttpResponseHandler::new(123);
826
827 let unfinished_response = b"HTTP/1.1 200 OK\r\nContent-Type: text/plain";
828
829 let past_time = Instant::now() - Duration::from_millis(DEFAULT_MIN_LOG_SYNC_MS + 10);
830
831 assert!(!handler.check_time_based_logging(unfinished_response, past_time));
832 }
833
834 #[derive(Clone)]
835 struct MockLogger {
836 logs: Arc<Mutex<Vec<String>>>,
837 }
838
839 impl MockLogger {
840 fn new() -> Self {
841 Self {
842 logs: Arc::new(Mutex::new(Vec::new())),
843 }
844 }
845
846 async fn log_count(&self) -> usize {
847 self.logs.lock().await.len()
848 }
849
850 async fn last_log(&self) -> Option<String> {
851 let logs = self.logs.lock().await;
852 logs.last().cloned()
853 }
854 }
855
856 #[tokio::test]
857 async fn test_chunked_transfer_detection() {
858 let chunked_headers = b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\ndata";
859 assert!(HttpResponseHandler::is_chunked_transfer(chunked_headers));
860
861 let regular_headers = b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n\r\nContent";
862 assert!(!HttpResponseHandler::is_chunked_transfer(regular_headers));
863
864 let non_chunked_transfer = b"HTTP/1.1 200 OK\r\nTransfer-Encoding: gzip\r\n\r\ndata";
865 assert!(!HttpResponseHandler::is_chunked_transfer(
866 non_chunked_transfer
867 ));
868 }
869
870 #[tokio::test]
871 async fn test_websocket_upgrade_detection() {
872 use crate::HttpResponseAnalyzer;
873
874 let websocket_data = b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n\r\n";
875
876 println!("Testing websocket detection with complete headers");
877 let result = HttpResponseAnalyzer::is_websocket_upgrade(websocket_data);
878 println!("Result: {result}");
879
880 assert!(result, "Should detect complete websocket upgrade response");
881
882 let regular_data = b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n\r\nHello";
883 assert!(
884 !HttpResponseAnalyzer::is_websocket_upgrade(regular_data),
885 "Should not detect regular HTTP response as websocket"
886 );
887
888 let partial_upgrade = b"HTTP/1.1 101 Switching Protocols\r\nConnection: Upgrade\r\nUpgrade: websocket\r\n\r\n";
889 let partial_result = HttpResponseAnalyzer::is_websocket_upgrade(partial_upgrade);
890 println!("Partial result: {partial_result}");
891
892 assert!(
893 !partial_result,
894 "Should NOT detect WebSocket upgrade without Sec-WebSocket-Accept header"
895 );
896 }
897
898 #[tokio::test]
899 async fn test_format_response_for_logging() {
900 let handler = HttpResponseHandler::new(123);
901 let response_data =
902 b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nServer: test\r\n\r\nHello World";
903
904 let result = handler.format_response_for_logging(response_data).await;
905
906 assert!(result.is_ok(), "format_response_for_logging should succeed");
907 let formatted = result.unwrap();
908
909 println!("Formatted response:\n{formatted}");
910 println!(
911 "Original response: {:?}",
912 std::str::from_utf8(response_data).unwrap()
913 );
914
915 assert!(
916 formatted.contains("HTTP/1.1 200 OK"),
917 "Should contain status line"
918 );
919 assert!(
920 formatted.contains("Hello World"),
921 "Should contain body content"
922 );
923 }
924
925 #[tokio::test]
926 async fn test_mock_logging() {
927 let logger = Arc::new(MockLogger::new());
928
929 {
930 let mut logs = logger.logs.lock().await;
931 logs.push("Test log entry 1".to_string());
932 logs.push("Test log entry 2".to_string());
933 }
934
935 assert_eq!(logger.log_count().await, 2);
936
937 assert_eq!(logger.last_log().await.unwrap(), "Test log entry 2");
938 }
939
940 #[tokio::test]
941 async fn test_format_response_for_logging_full() {
942 let handler = HttpResponseHandler::new(123);
943
944 let response_data =
945 b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: 9\r\n\r\nTest body";
946
947 let result = handler.format_response_for_logging(response_data).await;
948
949 assert!(result.is_ok());
950 let formatted = result.unwrap();
951
952 println!("Formatted response: {formatted}");
953
954 assert!(formatted.contains("HTTP/1.1 200 OK"));
955
956 assert!(formatted.contains("Test body"));
957 }
958
959 #[tokio::test]
960 async fn test_format_response_with_empty_body() {
961 let handler = HttpResponseHandler::new(123);
962
963 let response_data = b"HTTP/1.1 204 No Content\r\nServer: Test\r\n\r\n";
964
965 let result = handler.format_response_for_logging(response_data).await;
966
967 assert!(result.is_ok());
968 let formatted = result.unwrap();
969
970 println!("Formatted empty body response: {formatted}");
971
972 assert!(!formatted.is_empty());
973
974 assert!(formatted.contains("empty") || formatted.contains("#") || !formatted.is_empty());
975 }
976}