44use super :: * ;
55use crate :: grpc:: mutation_replay:: tests:: { reason, state_for} ;
66use crate :: grpc:: mutation_replay:: {
7- Admission , Completion , OBJECT_TYPE , SUCCESS_TTL_MS , fingerprint, run,
7+ Admission , AuthorizationBarrier , Completion , OBJECT_TYPE , OriginalMutation , SUCCESS_TTL_MS ,
8+ fingerprint, run,
89} ;
910use crate :: grpc:: test_support:: authed_request;
1011use crate :: persistence:: { ObjectType , WriteCondition , current_time_ms} ;
@@ -13,9 +14,83 @@ use openshell_core::proto::{
1314 GatewayMessage , Sandbox , SandboxPhase , SandboxStatus , gateway_message,
1415} ;
1516use openshell_core:: { ObjectId , ObjectName } ;
17+ use prost:: Message ;
1618use tokio:: sync:: { mpsc, oneshot} ;
1719use tokio_stream:: StreamExt ;
1820
21+ type ExecClient =
22+ openshell_core:: proto:: open_shell_client:: OpenShellClient < tonic:: transport:: Channel > ;
23+
24+ async fn wire_client (
25+ state : Arc < ServerState > ,
26+ barrier : Option < Arc < AuthorizationBarrier > > ,
27+ original : Option < OriginalMutation > ,
28+ ) -> ( ExecClient , tokio:: task:: JoinHandle < ( ) > ) {
29+ use crate :: grpc:: OpenShellService ;
30+ use openshell_core:: proto:: open_shell_server:: OpenShellServer ;
31+ use tokio_stream:: wrappers:: TcpListenerStream ;
32+
33+ let listener = tokio:: net:: TcpListener :: bind ( "127.0.0.1:0" ) . await . unwrap ( ) ;
34+ let address = listener. local_addr ( ) . unwrap ( ) ;
35+ let principal = authed_request ( ( ) )
36+ . extensions ( )
37+ . get :: < Principal > ( )
38+ . unwrap ( )
39+ . clone ( ) ;
40+ let service = OpenShellServer :: with_interceptor (
41+ OpenShellService :: new ( state) ,
42+ move |mut request : Request < ( ) > | {
43+ request. extensions_mut ( ) . insert ( principal. clone ( ) ) ;
44+ if let Some ( barrier) = & barrier {
45+ request. extensions_mut ( ) . insert ( barrier. clone ( ) ) ;
46+ }
47+ if let Some ( original) = & original {
48+ request. extensions_mut ( ) . insert ( original. clone ( ) ) ;
49+ }
50+ Ok ( request)
51+ } ,
52+ ) ;
53+ let server = tokio:: spawn ( async move {
54+ tonic:: transport:: Server :: builder ( )
55+ . add_service ( service)
56+ . serve_with_incoming ( TcpListenerStream :: new ( listener) )
57+ . await
58+ . unwrap ( ) ;
59+ } ) ;
60+ let client = ExecClient :: connect ( format ! ( "http://{address}" ) )
61+ . await
62+ . unwrap ( ) ;
63+ ( client, server)
64+ }
65+
66+ fn interactive_start ( req : ExecSandboxRequest ) -> ExecSandboxInput {
67+ ExecSandboxInput {
68+ payload : Some ( exec_sandbox_input:: Payload :: Start ( req) ) ,
69+ }
70+ }
71+
72+ fn original_exec ( req : ExecSandboxRequest , interactive : bool ) -> OriginalMutation {
73+ OriginalMutation ( if interactive {
74+ interactive_start ( req) . encode_to_vec ( )
75+ } else {
76+ req. encode_to_vec ( )
77+ } )
78+ }
79+
80+ async fn exec_rpc (
81+ client : & mut ExecClient ,
82+ req : ExecSandboxRequest ,
83+ interactive : bool ,
84+ ) -> Result < Response < tonic:: Streaming < ExecSandboxEvent > > , Status > {
85+ if interactive {
86+ client
87+ . exec_sandbox_interactive ( tokio_stream:: iter ( [ interactive_start ( req) ] ) )
88+ . await
89+ } else {
90+ client. exec_sandbox ( req) . await
91+ }
92+ }
93+
1994async fn setup ( url : & str , directory : & tempfile:: TempDir ) -> Arc < ServerState > {
2095 let mut state = state_for ( Store :: connect ( url) . await . unwrap ( ) ) . await ;
2196 let key = directory. path ( ) . join ( "key" ) ;
@@ -377,6 +452,125 @@ async fn exec_restart_preserves_pending_and_terminal_fences_and_reauthorizes() {
377452 ) ;
378453}
379454
455+ #[ tokio:: test]
456+ async fn noninteractive_rejects_replacement_between_authorizations ( ) {
457+ replacement_between_authorizations ( false ) . await ;
458+ }
459+
460+ #[ tokio:: test]
461+ async fn interactive_rejects_replacement_between_authorizations ( ) {
462+ replacement_between_authorizations ( true ) . await ;
463+ }
464+
465+ async fn replacement_between_authorizations ( interactive : bool ) {
466+ // A payload-only transformation must not be mistaken for redirection.
467+ for change_command in [ false , true ] {
468+ let directory = tempfile:: tempdir ( ) . unwrap ( ) ;
469+ let state = setup ( "sqlite::memory:" , & directory) . await ;
470+ let mut req = request ( & state) . await ;
471+ let original = change_command. then ( || original_exec ( req. clone ( ) , interactive) ) ;
472+ if change_command {
473+ req. command . push ( "transformed-argument" . into ( ) ) ;
474+ }
475+ let id = sandbox_id ( & state, & req) . await ;
476+ let mut old_messages = register ( & state, & id) ;
477+ let barrier = Arc :: new ( AuthorizationBarrier :: default ( ) ) ;
478+ let ( mut client, server) =
479+ wire_client ( state. clone ( ) , Some ( barrier. clone ( ) ) , original) . await ;
480+ let call = tokio:: spawn ( async move { exec_rpc ( & mut client, req, interactive) . await } ) ;
481+ tokio:: time:: timeout (
482+ std:: time:: Duration :: from_secs ( 5 ) ,
483+ barrier. resolved . notified ( ) ,
484+ )
485+ . await
486+ . unwrap ( ) ;
487+
488+ let mut replacement: Sandbox = state. store . get_message ( & id) . await . unwrap ( ) . unwrap ( ) ;
489+ state
490+ . store
491+ . delete ( Sandbox :: object_type ( ) , & id)
492+ . await
493+ . unwrap ( ) ;
494+ replacement. metadata . as_mut ( ) . unwrap ( ) . id = uuid:: Uuid :: new_v4 ( ) . to_string ( ) ;
495+ state. store . put_message ( & replacement) . await . unwrap ( ) ;
496+ let mut replacement_messages = register ( & state, replacement. object_id ( ) ) ;
497+ barrier. resume . notify_one ( ) ;
498+
499+ let status = tokio:: time:: timeout ( std:: time:: Duration :: from_secs ( 5 ) , call)
500+ . await
501+ . unwrap ( )
502+ . unwrap ( )
503+ . unwrap_err ( ) ;
504+ assert_eq ! ( reason( & status) , "REQUEST_REPLAY_UNAVAILABLE" ) ;
505+ assert ! ( old_messages. try_recv( ) . is_err( ) ) ;
506+ assert ! ( replacement_messages. try_recv( ) . is_err( ) ) ;
507+ assert ! (
508+ state
509+ . store
510+ . list_by_type_after( OBJECT_TYPE , None , 100 )
511+ . await
512+ . unwrap( )
513+ . is_empty( )
514+ ) ;
515+ server. abort ( ) ;
516+ }
517+ }
518+
519+ #[ tokio:: test]
520+ async fn exec_transformations_preserve_unchanged_targets_and_explicit_redirection ( ) {
521+ use openshell_core:: proto:: Workspace ;
522+
523+ for interactive in [ false , true ] {
524+ for redirect in [ "command-only" , "sandbox" , "workspace" ] {
525+ let directory = tempfile:: tempdir ( ) . unwrap ( ) ;
526+ let state = setup ( "sqlite::memory:" , & directory) . await ;
527+ let mut req = request ( & state) . await ;
528+ let original = original_exec ( req. clone ( ) , interactive) ;
529+ let original_id = sandbox_id ( & state, & req) . await ;
530+ let mut target: Sandbox = state
531+ . store
532+ . get_message ( & original_id)
533+ . await
534+ . unwrap ( )
535+ . unwrap ( ) ;
536+ req. command . push ( "transformed-argument" . into ( ) ) ;
537+ if redirect != "command-only" {
538+ target. metadata . as_mut ( ) . unwrap ( ) . id = uuid:: Uuid :: new_v4 ( ) . to_string ( ) ;
539+ if redirect == "sandbox" {
540+ req. sandbox = "redirected" . into ( ) ;
541+ target. metadata . as_mut ( ) . unwrap ( ) . name = req. sandbox . clone ( ) ;
542+ } else {
543+ let workspace = Workspace {
544+ metadata : Some ( ObjectMeta {
545+ id : uuid:: Uuid :: new_v4 ( ) . to_string ( ) ,
546+ name : "redirected" . into ( ) ,
547+ ..Default :: default ( )
548+ } ) ,
549+ ..Default :: default ( )
550+ } ;
551+ state. store . put_message ( & workspace) . await . unwrap ( ) ;
552+ target. metadata . as_mut ( ) . unwrap ( ) . workspace = "redirected" . into ( ) ;
553+ req. workspace_scope =
554+ Some ( openshell_core:: proto:: workspace_selector ( "redirected" ) ) ;
555+ }
556+ state. store . put_message ( & target) . await . unwrap ( ) ;
557+ }
558+ let mut messages = register ( & state, target. object_id ( ) ) ;
559+ let ( mut client, server) = wire_client ( state. clone ( ) , None , Some ( original) ) . await ;
560+ drop ( exec_rpc ( & mut client, req, interactive) . await . unwrap ( ) ) ;
561+ receive_relay ( & state, & mut messages) . await ;
562+ let owner = completion ( & state) . await ;
563+ let admission: Admission = serde_json:: from_slice ( & owner. admission ) . unwrap ( ) ;
564+ assert_eq ! ( admission. target_id. as_deref( ) , Some ( original_id. as_str( ) ) ) ;
565+ owner. ensure_target ( target. object_id ( ) ) . unwrap ( ) ;
566+ if redirect != "command-only" {
567+ assert ! ( owner. ensure_target( & original_id) . is_err( ) ) ;
568+ }
569+ server. abort ( ) ;
570+ }
571+ }
572+ }
573+
380574#[ tokio:: test]
381575async fn replacement_cannot_reuse_an_exec_claim_or_receive_its_admitted_launch ( ) {
382576 for terminal in [ false , true ] {
0 commit comments