@@ -10,26 +10,32 @@ use zentinel_agent_waf::{WafAgent, WafConfig, WebSocketConfig};
1010
1111/// Create a WebSocket-enabled agent for testing
1212fn create_websocket_agent ( ) -> WafAgent {
13- let mut config = WafConfig :: default ( ) ;
14- config. websocket = WebSocketConfig {
15- enabled : true ,
16- inspect_text_frames : true ,
17- inspect_binary_frames : false ,
18- max_frame_size : 65536 ,
19- block_mode : true ,
20- accumulate_fragments : true ,
21- max_message_size : 1048576 ,
22- block_close_code : 1008 ,
23- block_close_reason : "WAF policy violation" . to_string ( ) ,
13+ let config = WafConfig {
14+ websocket : WebSocketConfig {
15+ enabled : true ,
16+ inspect_text_frames : true ,
17+ inspect_binary_frames : false ,
18+ max_frame_size : 65536 ,
19+ block_mode : true ,
20+ accumulate_fragments : true ,
21+ max_message_size : 1048576 ,
22+ block_close_code : 1008 ,
23+ block_close_reason : "WAF policy violation" . to_string ( ) ,
24+ } ,
25+ ..Default :: default ( )
2426 } ;
2527 WafAgent :: new ( config) . expect ( "Failed to create agent" )
2628}
2729
2830/// Create a WebSocket agent with custom config
2931fn create_websocket_agent_with_config ( f : impl FnOnce ( & mut WebSocketConfig ) ) -> WafAgent {
30- let mut config = WafConfig :: default ( ) ;
31- config. websocket = WebSocketConfig :: default ( ) ;
32- config. websocket . enabled = true ;
32+ let mut config = WafConfig {
33+ websocket : WebSocketConfig {
34+ enabled : true ,
35+ ..Default :: default ( )
36+ } ,
37+ ..Default :: default ( )
38+ } ;
3339 f ( & mut config. websocket ) ;
3440 WafAgent :: new ( config) . expect ( "Failed to create agent" )
3541}
@@ -88,20 +94,18 @@ fn continuation_frame(
8894
8995/// Check if response indicates a block (close connection)
9096fn is_blocked ( response : & AgentResponse ) -> bool {
91- match & response. websocket_decision {
92- Some ( WebSocketDecision :: Close { .. } ) => true ,
93- Some ( WebSocketDecision :: Drop ) => true ,
94- _ => false ,
95- }
97+ matches ! (
98+ & response. websocket_decision,
99+ Some ( WebSocketDecision :: Close { .. } ) | Some ( WebSocketDecision :: Drop )
100+ )
96101}
97102
98103/// Check if response indicates allow
99104fn is_allowed ( response : & AgentResponse ) -> bool {
100- match & response. websocket_decision {
101- Some ( WebSocketDecision :: Allow ) => true ,
102- None => true , // Default is allow
103- _ => false ,
104- }
105+ matches ! (
106+ & response. websocket_decision,
107+ Some ( WebSocketDecision :: Allow ) | None
108+ )
105109}
106110
107111// =============================================================================
0 commit comments