@@ -4,6 +4,7 @@ use std::sync::{Arc, Mutex as StdMutex};
44
55use contextforge_gateway_rs_cpex:: CpexRuntimeRegistry ;
66use cpex:: cpex_core:: cmf:: Role ;
7+ use cpex:: cpex_core:: config:: CpexConfig ;
78use cpex:: cpex_core:: hooks:: types:: cmf_hook_names;
89use rmcp:: {
910 ClientHandler ,
@@ -13,7 +14,7 @@ use rmcp::{
1314 } ,
1415 service:: { NotificationContext , PeerRequestOptions , RequestHandle , RoleClient , RunningService } ,
1516} ;
16- use serde_json:: Value ;
17+ use serde_json:: { Map , Value , json } ;
1718
1819use support:: {
1920 POST_DENY_ERROR_CODE , PRE_DENY_ERROR_CODE , REWRITTEN_SUM_A , REWRITTEN_SUM_B , RunningGateway , TEST_USER_ID ,
@@ -124,6 +125,43 @@ fn raw_tool_call(tool_name: &str, request_id: i64, progress_token: &str) -> Valu
124125 } )
125126}
126127
128+ fn fake_aws_access_key ( suffix : & str ) -> String {
129+ [ "AKIA" , suffix] . concat ( )
130+ }
131+
132+ fn sum_request_with_secret ( secret_field : & str , secret : String ) -> CallToolRequestParams {
133+ let mut request = sum_request ( "sum" , 1 , 2 ) ;
134+ request
135+ . arguments
136+ . as_mut ( )
137+ . expect ( "sum request has arguments" )
138+ . insert ( secret_field. to_owned ( ) , Value :: String ( secret) ) ;
139+ request
140+ }
141+
142+ fn reflect_text_request ( text : String ) -> CallToolRequestParams {
143+ CallToolRequestParams :: new ( "reflect_text" )
144+ . with_arguments ( Map :: from_iter ( [ ( "text" . to_owned ( ) , Value :: String ( text) ) ] ) )
145+ }
146+
147+ async fn runtime_with_secrets_detection ( hooks : Vec < & ' static str > , plugin_config : Value ) -> Arc < CpexRuntimeRegistry > {
148+ let mut runtime = CpexRuntimeRegistry :: default ( ) ;
149+ runtime
150+ . register_factory ( secrets_detection_rust:: KIND , Box :: new ( secrets_detection_rust:: SecretsDetectionFactory ) )
151+ . expect ( "secrets detection factory registers" ) ;
152+ let config: CpexConfig = serde_json:: from_value ( json ! ( {
153+ "plugins" : [ {
154+ "name" : "secrets-detection" ,
155+ "kind" : secrets_detection_rust:: KIND ,
156+ "hooks" : hooks,
157+ "config" : plugin_config,
158+ } ]
159+ } ) )
160+ . expect ( "secrets detection CPEX config parses" ) ;
161+ runtime. apply_config ( Some ( config) ) . await . expect ( "secrets detection runtime applies" ) ;
162+ Arc :: new ( runtime)
163+ }
164+
127165fn sse_data_values ( body : & str ) -> Vec < Value > {
128166 let values = body
129167 . lines ( )
@@ -333,6 +371,139 @@ async fn disabled_runtime_does_not_invoke_registered_plugin() {
333371 assert_eq ! ( 0 , post_observations. lock( ) . expect( "observations lock poisoned" ) . post_calls) ;
334372}
335373
374+ #[ tokio:: test( flavor = "multi_thread" , worker_threads = 1 ) ]
375+ async fn secrets_detection_pre_hook_redacts_tool_arguments_before_backend_call ( ) {
376+ let runtime = runtime_with_secrets_detection (
377+ vec ! [ cmf_hook_names:: TOOL_PRE_INVOKE ] ,
378+ json ! ( {
379+ "redact" : true ,
380+ "redaction_text" : "[redacted]" ,
381+ "block_on_detection" : false ,
382+ } ) ,
383+ )
384+ . await ;
385+ let gateway = start_gateway ( "admin@example.com" , true , runtime) . await ;
386+ let service = gateway. connect ( "admin@example.com" ) . await ;
387+
388+ let result = service
389+ . call_tool ( sum_request_with_secret ( "credential" , fake_aws_access_key ( "1111111111111111" ) ) )
390+ . await
391+ . expect ( "secret argument is redacted and call succeeds" ) ;
392+
393+ assert_eq ! ( "3" , text( & result) ) ;
394+ let backend_calls = gateway. backend_state . calls . lock ( ) . expect ( "backend calls lock poisoned" ) ;
395+ assert_eq ! ( 1 , backend_calls. len( ) ) ;
396+ assert_eq ! (
397+ Some ( & Value :: from( "[redacted]" ) ) ,
398+ backend_calls[ 0 ] . args. as_ref( ) . and_then( |args| args. get( "credential" ) )
399+ ) ;
400+ }
401+
402+ #[ tokio:: test( flavor = "multi_thread" , worker_threads = 1 ) ]
403+ async fn secrets_detection_clean_tool_payload_passes_through_unchanged ( ) {
404+ let runtime = runtime_with_secrets_detection (
405+ vec ! [ cmf_hook_names:: TOOL_PRE_INVOKE , cmf_hook_names:: TOOL_POST_INVOKE ] ,
406+ json ! ( {
407+ "redact" : true ,
408+ "redaction_text" : "[redacted]" ,
409+ "block_on_detection" : true ,
410+ } ) ,
411+ )
412+ . await ;
413+ let gateway = start_gateway ( "admin@example.com" , true , runtime) . await ;
414+ let service = gateway. connect ( "admin@example.com" ) . await ;
415+
416+ let result =
417+ service. call_tool ( sum_request ( "sum" , 1 , 2 ) ) . await . expect ( "clean argument payload passes through unchanged" ) ;
418+
419+ let result_text = text ( & result) ;
420+ assert_eq ! ( "3" , result_text. as_str( ) ) ;
421+ assert ! ( !result_text. contains( "[redacted]" ) ) ;
422+ let backend_calls = gateway. backend_state . calls . lock ( ) . expect ( "backend calls lock poisoned" ) ;
423+ assert_eq ! ( 1 , backend_calls. len( ) ) ;
424+ assert_eq ! ( "sum" , backend_calls[ 0 ] . tool_name) ;
425+ let args = backend_calls[ 0 ] . args . as_ref ( ) . expect ( "backend call has args" ) ;
426+ assert_eq ! ( 2 , args. len( ) ) ;
427+ assert_eq ! ( Some ( & Value :: from( 1 ) ) , args. get( "a" ) ) ;
428+ assert_eq ! ( Some ( & Value :: from( 2 ) ) , args. get( "b" ) ) ;
429+ }
430+
431+ #[ tokio:: test( flavor = "multi_thread" , worker_threads = 1 ) ]
432+ async fn secrets_detection_pre_hook_blocks_tool_arguments_before_backend_call ( ) {
433+ let runtime = runtime_with_secrets_detection (
434+ vec ! [ cmf_hook_names:: TOOL_PRE_INVOKE ] ,
435+ json ! ( {
436+ "redact" : false ,
437+ "block_on_detection" : true ,
438+ } ) ,
439+ )
440+ . await ;
441+ let gateway = start_gateway ( "admin@example.com" , true , runtime) . await ;
442+ let service = gateway. connect ( "admin@example.com" ) . await ;
443+
444+ let error = service
445+ . call_tool ( sum_request_with_secret ( "credential" , fake_aws_access_key ( "2222222222222222" ) ) )
446+ . await
447+ . expect_err ( "secret argument blocks the call" ) ;
448+
449+ assert_eq ! ( ErrorCode :: INVALID_REQUEST , error_code( error) ) ;
450+ assert ! ( gateway. backend_state. calls. lock( ) . expect( "backend calls lock poisoned" ) . is_empty( ) ) ;
451+ }
452+
453+ #[ tokio:: test( flavor = "multi_thread" , worker_threads = 1 ) ]
454+ async fn secrets_detection_post_hook_redacts_tool_result_before_client_response ( ) {
455+ let runtime = runtime_with_secrets_detection (
456+ vec ! [ cmf_hook_names:: TOOL_POST_INVOKE ] ,
457+ json ! ( {
458+ "redact" : true ,
459+ "redaction_text" : "[redacted]" ,
460+ "block_on_detection" : false ,
461+ } ) ,
462+ )
463+ . await ;
464+ let gateway = start_gateway ( "admin@example.com" , true , runtime) . await ;
465+ let service = gateway. connect ( "admin@example.com" ) . await ;
466+
467+ let result = service
468+ . call_tool ( reflect_text_request ( fake_aws_access_key ( "3333333333333333" ) ) )
469+ . await
470+ . expect ( "secret result is redacted and call succeeds" ) ;
471+
472+ assert_eq ! ( "[redacted]" , text( & result) ) ;
473+ assert_eq ! ( 1 , gateway. backend_state. calls. lock( ) . expect( "backend calls lock poisoned" ) . len( ) ) ;
474+ }
475+
476+ #[ tokio:: test( flavor = "multi_thread" , worker_threads = 1 ) ]
477+ async fn secrets_detection_pre_hook_respects_field_allowlist ( ) {
478+ let runtime = runtime_with_secrets_detection (
479+ vec ! [ cmf_hook_names:: TOOL_PRE_INVOKE ] ,
480+ json ! ( {
481+ "redact" : true ,
482+ "redaction_text" : "[redacted]" ,
483+ "block_on_detection" : false ,
484+ "field_allowlist" : [ "credential" ] ,
485+ } ) ,
486+ )
487+ . await ;
488+ let gateway = start_gateway ( "admin@example.com" , true , runtime) . await ;
489+ let service = gateway. connect ( "admin@example.com" ) . await ;
490+ let ignored_secret = fake_aws_access_key ( "4444444444444444" ) ;
491+ let mut request = sum_request_with_secret ( "credential" , fake_aws_access_key ( "5555555555555555" ) ) ;
492+ request
493+ . arguments
494+ . as_mut ( )
495+ . expect ( "sum request has arguments" )
496+ . insert ( "ignored" . to_owned ( ) , Value :: String ( ignored_secret. clone ( ) ) ) ;
497+
498+ let result = service. call_tool ( request) . await . expect ( "allowed field is redacted" ) ;
499+
500+ assert_eq ! ( "3" , text( & result) ) ;
501+ let backend_calls = gateway. backend_state . calls . lock ( ) . expect ( "backend calls lock poisoned" ) ;
502+ let args = backend_calls[ 0 ] . args . as_ref ( ) . expect ( "backend call has args" ) ;
503+ assert_eq ! ( Some ( & Value :: from( "[redacted]" ) ) , args. get( "credential" ) ) ;
504+ assert_eq ! ( Some ( & Value :: from( ignored_secret) ) , args. get( "ignored" ) ) ;
505+ }
506+
336507#[ tokio:: test( flavor = "multi_thread" , worker_threads = 1 ) ]
337508async fn pre_hook_modifies_backend_arguments_without_rerouting_tool ( ) {
338509 let plugin = Arc :: new ( TestPlugin :: new ( "pre" , vec ! [ cmf_hook_names:: TOOL_PRE_INVOKE ] ) . with_pre_rewrite ( ) ) ;
0 commit comments