@@ -11,17 +11,18 @@ use anyhow::{anyhow, Context, Result};
11
11
use rustix:: fs:: { flock, FlockOperation } ;
12
12
use uuid:: Uuid ;
13
13
14
- use super :: utils:: env:: find_in_path;
15
14
use crate :: env:: prepare_env_vars;
15
+ use crate :: tty:: { run_io_host, RawTerminal } ;
16
16
use crate :: utils:: launch:: Launch ;
17
- use rustix :: path :: Arg ;
17
+ use nix :: unistd :: unlink ;
18
18
use std:: ops:: Range ;
19
- use std:: process:: { Child , Command } ;
19
+ use std:: os:: unix:: net:: UnixListener ;
20
+ use std:: process:: ExitCode ;
20
21
21
22
pub const DYNAMIC_PORT_RANGE : Range < u32 > = 50000 ..50200 ;
22
23
23
24
pub enum LaunchResult {
24
- LaunchRequested ,
25
+ LaunchRequested ( ExitCode ) ,
25
26
LockAcquired {
26
27
cookie : Uuid ,
27
28
lock_file : File ,
@@ -56,64 +57,75 @@ impl Display for LaunchError {
56
57
}
57
58
}
58
59
59
- fn start_socat ( ) -> Result < ( Child , u32 ) > {
60
+ fn acquire_socket_lock ( ) -> Result < ( File , u32 ) > {
60
61
let run_path = env:: var ( "XDG_RUNTIME_DIR" )
61
62
. map_err ( |e| anyhow ! ( "unable to get XDG_RUNTIME_DIR: {:?}" , e) ) ?;
62
63
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" ) ) ?;
65
64
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
85
82
} ,
86
- _ => { } ,
87
- }
88
- ret. push ( c) ;
83
+ port,
84
+ ) ) ;
89
85
}
90
- ret
86
+ Err ( anyhow ! ( "Ran out of ports." ) )
91
87
}
92
88
93
89
fn wrapped_launch (
94
90
server_port : u32 ,
95
91
cookie : Uuid ,
96
- mut command : PathBuf ,
97
- mut command_args : Vec < String > ,
92
+ command : PathBuf ,
93
+ command_args : Vec < String > ,
98
94
env : HashMap < String , String > ,
99
95
interactive : bool ,
100
- ) -> Result < ( ) > {
96
+ tty : bool ,
97
+ ) -> Result < ExitCode > {
101
98
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 ) ) ;
103
101
}
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) )
117
129
}
118
130
119
131
pub fn launch_or_lock (
@@ -122,16 +134,17 @@ pub fn launch_or_lock(
122
134
command_args : Vec < String > ,
123
135
env : Vec < ( String , Option < String > ) > ,
124
136
interactive : bool ,
137
+ tty : bool ,
125
138
) -> Result < LaunchResult > {
126
139
let running_server_port = env:: var ( "MUVM_SERVER_PORT" ) . ok ( ) ;
127
140
if let Some ( port) = running_server_port {
128
141
let port: u32 = port. parse ( ) ?;
129
142
let env = prepare_env_vars ( env) ?;
130
143
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
+ } ;
135
148
}
136
149
137
150
let ( lock_file, cookie) = lock_file ( ) ?;
@@ -154,6 +167,7 @@ pub fn launch_or_lock(
154
167
command_args. clone ( ) ,
155
168
env. clone ( ) ,
156
169
interactive,
170
+ tty,
157
171
) {
158
172
Err ( err) => match err. downcast_ref :: < LaunchError > ( ) {
159
173
Some ( & LaunchError :: Connection ( _) ) => {
@@ -167,7 +181,7 @@ pub fn launch_or_lock(
167
181
return Err ( anyhow ! ( "could not request launch to server: {err}" ) ) ;
168
182
} ,
169
183
} ,
170
- Ok ( _ ) => return Ok ( LaunchResult :: LaunchRequested ) ,
184
+ Ok ( code ) => return Ok ( LaunchResult :: LaunchRequested ( code ) ) ,
171
185
}
172
186
}
173
187
} ,
@@ -222,6 +236,8 @@ pub fn request_launch(
222
236
command : PathBuf ,
223
237
command_args : Vec < String > ,
224
238
env : HashMap < String , String > ,
239
+ vsock_port : u32 ,
240
+ tty : bool ,
225
241
) -> Result < ( ) > {
226
242
let mut stream =
227
243
TcpStream :: connect ( format ! ( "127.0.0.1:{server_port}" ) ) . map_err ( LaunchError :: Connection ) ?;
@@ -231,6 +247,8 @@ pub fn request_launch(
231
247
command,
232
248
command_args,
233
249
env,
250
+ vsock_port,
251
+ tty,
234
252
} ;
235
253
236
254
stream
0 commit comments