@@ -11,17 +11,18 @@ use anyhow::{anyhow, Context, Result};
1111use rustix:: fs:: { flock, FlockOperation } ;
1212use uuid:: Uuid ;
1313
14- use super :: utils:: env:: find_in_path;
1514use crate :: env:: prepare_env_vars;
15+ use crate :: tty:: { run_io_host, RawTerminal } ;
1616use crate :: utils:: launch:: Launch ;
17- use rustix :: path :: Arg ;
17+ use nix :: unistd :: unlink ;
1818use std:: ops:: Range ;
19- use std:: process:: { Child , Command } ;
19+ use std:: os:: unix:: net:: UnixListener ;
20+ use std:: process:: ExitCode ;
2021
2122pub const DYNAMIC_PORT_RANGE : Range < u32 > = 50000 ..50200 ;
2223
2324pub enum LaunchResult {
24- LaunchRequested ,
25+ LaunchRequested ( ExitCode ) ,
2526 LockAcquired {
2627 cookie : Uuid ,
2728 lock_file : File ,
@@ -56,64 +57,75 @@ impl Display for LaunchError {
5657 }
5758}
5859
59- fn start_socat ( ) -> Result < ( Child , u32 ) > {
60+ fn acquire_socket_lock ( ) -> Result < ( File , u32 ) > {
6061 let run_path = env:: var ( "XDG_RUNTIME_DIR" )
6162 . map_err ( |e| anyhow ! ( "unable to get XDG_RUNTIME_DIR: {:?}" , e) ) ?;
6263 let socket_dir = Path :: new ( & run_path) . join ( "krun/socket" ) ;
63- let socat_path =
64- find_in_path ( "socat" ) ?. ok_or_else ( || anyhow ! ( "Unable to find socat in PATH" ) ) ?;
6564 for port in DYNAMIC_PORT_RANGE {
66- let path = socket_dir. join ( format ! ( "port-{}" , port) ) ;
67- if path. exists ( ) {
68- continue ;
69- }
70- let child = Command :: new ( & socat_path)
71- . arg ( format ! ( "unix-l:{}" , path. as_os_str( ) . to_string_lossy( ) ) )
72- . arg ( "-,raw,echo=0" )
73- . spawn ( ) ?;
74- return Ok ( ( child, port) ) ;
75- }
76- Err ( anyhow ! ( "Ran out of ports." ) )
77- }
78-
79- fn escape_for_socat ( s : String ) -> String {
80- let mut ret = String :: with_capacity ( s. len ( ) ) ;
81- for c in s. chars ( ) {
82- match c {
83- ':' | ',' | '!' | '"' | '\'' | '\\' | '(' | '[' | '{' => {
84- ret. push ( '\\' ) ;
65+ let path = socket_dir. join ( format ! ( "port-{port}.lock" ) ) ;
66+ return Ok ( (
67+ if !path. exists ( ) {
68+ let lock_file = File :: create ( path) . context ( "Failed to create socket lock" ) ?;
69+ flock ( & lock_file, FlockOperation :: NonBlockingLockExclusive )
70+ . context ( "Failed to acquire socket lock" ) ?;
71+ lock_file
72+ } else {
73+ let lock_file = File :: options ( )
74+ . write ( true )
75+ . read ( true )
76+ . open ( path)
77+ . context ( "Failed to open lock file" ) ?;
78+ if flock ( & lock_file, FlockOperation :: NonBlockingLockExclusive ) . is_err ( ) {
79+ continue ;
80+ }
81+ lock_file
8582 } ,
86- _ => { } ,
87- }
88- ret. push ( c) ;
83+ port,
84+ ) ) ;
8985 }
90- ret
86+ Err ( anyhow ! ( "Ran out of ports." ) )
9187}
9288
9389fn wrapped_launch (
9490 server_port : u32 ,
9591 cookie : Uuid ,
96- mut command : PathBuf ,
97- mut command_args : Vec < String > ,
92+ command : PathBuf ,
93+ command_args : Vec < String > ,
9894 env : HashMap < String , String > ,
9995 interactive : bool ,
100- ) -> Result < ( ) > {
96+ tty : bool ,
97+ ) -> Result < ExitCode > {
10198 if !interactive {
102- return request_launch ( server_port, cookie, command, command_args, env) ;
99+ request_launch ( server_port, cookie, command, command_args, env, 0 , false ) ?;
100+ return Ok ( ExitCode :: from ( 0 ) ) ;
103101 }
104- let ( mut socat, vsock_port) = start_socat ( ) ?;
105- command_args. insert ( 0 , command. to_string_lossy ( ) . into_owned ( ) ) ;
106- command_args = vec ! [
107- format!( "vsock:2:{}" , vsock_port) ,
108- format!(
109- "exec:{},pty,setsid,stderr" ,
110- escape_for_socat( command_args. join( " " ) )
111- ) ,
112- ] ;
113- command = "socat" . into ( ) ;
114- request_launch ( server_port, cookie, command, command_args, env) ?;
115- socat. wait ( ) ?;
116- Ok ( ( ) )
102+ let run_path = env:: var ( "XDG_RUNTIME_DIR" )
103+ . map_err ( |e| anyhow ! ( "unable to get XDG_RUNTIME_DIR: {:?}" , e) ) ?;
104+ let socket_dir = Path :: new ( & run_path) . join ( "krun/socket" ) ;
105+ let ( _lock, vsock_port) = acquire_socket_lock ( ) ?;
106+ let path = socket_dir. join ( format ! ( "port-{vsock_port}" ) ) ;
107+ _ = unlink ( & path) ;
108+ let listener = UnixListener :: bind ( path) . context ( "Failed to listen on vm socket" ) ?;
109+ let raw_tty = if tty {
110+ Some (
111+ RawTerminal :: set ( )
112+ . context ( "Asked to allocate a tty for the command, but stdin is not a tty" ) ?,
113+ )
114+ } else {
115+ None
116+ } ;
117+ request_launch (
118+ server_port,
119+ cookie,
120+ command,
121+ command_args,
122+ env,
123+ vsock_port,
124+ tty,
125+ ) ?;
126+ let code = run_io_host ( listener, tty) ?;
127+ drop ( raw_tty) ;
128+ Ok ( ExitCode :: from ( code) )
117129}
118130
119131pub fn launch_or_lock (
@@ -122,16 +134,17 @@ pub fn launch_or_lock(
122134 command_args : Vec < String > ,
123135 env : Vec < ( String , Option < String > ) > ,
124136 interactive : bool ,
137+ tty : bool ,
125138) -> Result < LaunchResult > {
126139 let running_server_port = env:: var ( "MUVM_SERVER_PORT" ) . ok ( ) ;
127140 if let Some ( port) = running_server_port {
128141 let port: u32 = port. parse ( ) ?;
129142 let env = prepare_env_vars ( env) ?;
130143 let cookie = read_cookie ( ) ?;
131- if let Err ( err ) = wrapped_launch ( port, cookie, command, command_args, env, interactive) {
132- return Err ( anyhow ! ( "could not request launch to server: {err}" ) ) ;
133- }
134- return Ok ( LaunchResult :: LaunchRequested ) ;
144+ return match wrapped_launch ( port, cookie, command, command_args, env, interactive, tty ) {
145+ Err ( err ) => Err ( anyhow ! ( "could not request launch to server: {err}" ) ) ,
146+ Ok ( code ) => Ok ( LaunchResult :: LaunchRequested ( code ) ) ,
147+ } ;
135148 }
136149
137150 let ( lock_file, cookie) = lock_file ( ) ?;
@@ -154,6 +167,7 @@ pub fn launch_or_lock(
154167 command_args. clone ( ) ,
155168 env. clone ( ) ,
156169 interactive,
170+ tty,
157171 ) {
158172 Err ( err) => match err. downcast_ref :: < LaunchError > ( ) {
159173 Some ( & LaunchError :: Connection ( _) ) => {
@@ -167,7 +181,7 @@ pub fn launch_or_lock(
167181 return Err ( anyhow ! ( "could not request launch to server: {err}" ) ) ;
168182 } ,
169183 } ,
170- Ok ( _ ) => return Ok ( LaunchResult :: LaunchRequested ) ,
184+ Ok ( code ) => return Ok ( LaunchResult :: LaunchRequested ( code ) ) ,
171185 }
172186 }
173187 } ,
@@ -222,6 +236,8 @@ pub fn request_launch(
222236 command : PathBuf ,
223237 command_args : Vec < String > ,
224238 env : HashMap < String , String > ,
239+ vsock_port : u32 ,
240+ tty : bool ,
225241) -> Result < ( ) > {
226242 let mut stream =
227243 TcpStream :: connect ( format ! ( "127.0.0.1:{server_port}" ) ) . map_err ( LaunchError :: Connection ) ?;
@@ -231,6 +247,8 @@ pub fn request_launch(
231247 command,
232248 command_args,
233249 env,
250+ vsock_port,
251+ tty,
234252 } ;
235253
236254 stream
0 commit comments