77mod  tests; 
88
99use  super :: os:: current_exe; 
10- use  crate :: ffi:: OsString ; 
10+ use  crate :: ffi:: { OsStr ,   OsString } ; 
1111use  crate :: fmt; 
1212use  crate :: io; 
1313use  crate :: num:: NonZeroU16 ; 
@@ -17,6 +17,7 @@ use crate::sys::path::get_long_path;
1717use  crate :: sys:: process:: ensure_no_nuls; 
1818use  crate :: sys:: { c,  to_u16s} ; 
1919use  crate :: sys_common:: wstr:: WStrUnits ; 
20+ use  crate :: sys_common:: AsInner ; 
2021use  crate :: vec; 
2122
2223use  crate :: iter; 
@@ -262,16 +263,92 @@ pub(crate) fn append_arg(cmd: &mut Vec<u16>, arg: &Arg, force_quotes: bool) -> i
262263 Ok ( ( ) ) 
263264} 
264265
266+ fn  append_bat_arg ( cmd :  & mut  Vec < u16 > ,  arg :  & OsStr ,  mut  quote :  bool )  -> io:: Result < ( ) >  { 
267+  ensure_no_nuls ( arg) ?; 
268+  // If an argument has 0 characters then we need to quote it to ensure 
269+  // that it actually gets passed through on the command line or otherwise 
270+  // it will be dropped entirely when parsed on the other end. 
271+  // 
272+  // We also need to quote the argument if it ends with `\` to guard against 
273+  // bat usage such as `"%~2"` (i.e. force quote arguments) otherwise a 
274+  // trailing slash will escape the closing quote. 
275+  if  arg. is_empty ( )  || arg. as_encoded_bytes ( ) . last ( )  == Some ( & b'\\' )  { 
276+  quote = true ; 
277+  } 
278+  for  cp in  arg. as_inner ( ) . inner . code_points ( )  { 
279+  if  let  Some ( cp)  = cp. to_char ( )  { 
280+  // Rather than trying to find every ascii symbol that must be quoted, 
281+  // we assume that all ascii symbols must be quoted unless they're known to be good. 
282+  // We also quote Unicode control blocks for good measure. 
283+  // Note an unquoted `\` is fine so long as the argument isn't otherwise quoted. 
284+  static  UNQUOTED :  & str  = r"#$*+-./:?@\_" ; 
285+  let  ascii_needs_quotes =
286+  cp. is_ascii ( )  && !( cp. is_ascii_alphanumeric ( )  || UNQUOTED . contains ( cp) ) ; 
287+  if  ascii_needs_quotes || cp. is_control ( )  { 
288+  quote = true ; 
289+  } 
290+  } 
291+  } 
292+ 
293+  if  quote { 
294+  cmd. push ( '"'  as  u16 ) ; 
295+  } 
296+  // Loop through the string, escaping `\` only if followed by `"`. 
297+  // And escaping `"` by doubling them. 
298+  let  mut  backslashes:  usize  = 0 ; 
299+  for  x in  arg. encode_wide ( )  { 
300+  if  x == '\\'  as  u16  { 
301+  backslashes += 1 ; 
302+  }  else  { 
303+  if  x == '"'  as  u16  { 
304+  // Add n backslashes to total 2n before internal `"`. 
305+  cmd. extend ( ( 0 ..backslashes) . map ( |_| '\\'  as  u16 ) ) ; 
306+  // Appending an additional double-quote acts as an escape. 
307+  cmd. push ( b'"'  as  u16 ) 
308+  }  else  if  x == '%'  as  u16  || x == '\r'  as  u16  { 
309+  // yt-dlp hack: replaces `%` with `%%cd:~,%` to stop %VAR% being expanded as an environment variable. 
310+  // 
311+  // # Explanation 
312+  // 
313+  // cmd supports extracting a substring from a variable using the following syntax: 
314+  // %variable:~start_index,end_index% 
315+  // 
316+  // In the above command `cd` is used as the variable and the start_index and end_index are left blank. 
317+  // `cd` is a built-in variable that dynamically expands to the current directory so it's always available. 
318+  // Explicitly omitting both the start and end index creates a zero-length substring. 
319+  // 
320+  // Therefore it all resolves to nothing. However, by doing this no-op we distract cmd.exe 
321+  // from potentially expanding %variables% in the argument. 
322+  cmd. extend_from_slice ( & [ 
323+  '%'  as  u16 ,  '%'  as  u16 ,  'c'  as  u16 ,  'd'  as  u16 ,  ':'  as  u16 ,  '~'  as  u16 , 
324+  ','  as  u16 , 
325+  ] ) ; 
326+  } 
327+  backslashes = 0 ; 
328+  } 
329+  cmd. push ( x) ; 
330+  } 
331+  if  quote { 
332+  // Add n backslashes to total 2n before ending `"`. 
333+  cmd. extend ( ( 0 ..backslashes) . map ( |_| '\\'  as  u16 ) ) ; 
334+  cmd. push ( '"'  as  u16 ) ; 
335+  } 
336+  Ok ( ( ) ) 
337+ } 
338+ 
265339pub ( crate )  fn  make_bat_command_line ( 
266340 script :  & [ u16 ] , 
267341 args :  & [ Arg ] , 
268342 force_quotes :  bool , 
269343)  -> io:: Result < Vec < u16 > >  { 
344+  const  INVALID_ARGUMENT_ERROR :  io:: Error  =
345+  io:: const_io_error!( io:: ErrorKind :: InvalidInput ,  r#"batch file arguments are invalid"# ) ; 
270346 // Set the start of the command line to `cmd.exe /c "` 
271347 // It is necessary to surround the command in an extra pair of quotes, 
272348 // hence the trailing quote here. It will be closed after all arguments 
273349 // have been added. 
274-  let  mut  cmd:  Vec < u16 >  = "cmd.exe /d /c \" " . encode_utf16 ( ) . collect ( ) ; 
350+  // Using /e:ON enables "command extensions" which is essential for the `%` hack to work. 
351+  let  mut  cmd:  Vec < u16 >  = "cmd.exe /e:ON /v:OFF /d /c \" " . encode_utf16 ( ) . collect ( ) ; 
275352
276353 // Push the script name surrounded by its quote pair. 
277354 cmd. push ( b'"'  as  u16 ) ; 
@@ -291,18 +368,22 @@ pub(crate) fn make_bat_command_line(
291368 // reconstructed by the batch script by default. 
292369 for  arg in  args { 
293370 cmd. push ( ' '  as  u16 ) ; 
294-  // Make sure to always quote special command prompt characters, including: 
295-  // * Characters `cmd /?` says require quotes. 
296-  // * `%` for environment variables, as in `%TMP%`. 
297-  // * `|<>` pipe/redirect characters. 
298-  const  SPECIAL :  & [ u8 ]  = b"\t  &()[]{}^=;!'+,`~%|<>" ; 
299-  let  force_quotes = match  arg { 
300-  Arg :: Regular ( arg)  if  !force_quotes => { 
301-  arg. as_encoded_bytes ( ) . iter ( ) . any ( |c| SPECIAL . contains ( c) ) 
371+  match  arg { 
372+  Arg :: Regular ( arg_os)  => { 
373+  let  arg_bytes = arg_os. as_encoded_bytes ( ) ; 
374+  // Disallow \r and \n as they may truncate the arguments. 
375+  const  DISALLOWED :  & [ u8 ]  = b"\r \n " ; 
376+  if  arg_bytes. iter ( ) . any ( |c| DISALLOWED . contains ( c) )  { 
377+  return  Err ( INVALID_ARGUMENT_ERROR ) ; 
378+  } 
379+  append_bat_arg ( & mut  cmd,  arg_os,  force_quotes) ?; 
380+  } 
381+  _ => { 
382+  // Raw arguments are passed on as-is. 
383+  // It's the user's responsibility to properly handle arguments in this case. 
384+  append_arg ( & mut  cmd,  arg,  force_quotes) ?; 
302385 } 
303-  _ => force_quotes, 
304386 } ; 
305-  append_arg ( & mut  cmd,  arg,  force_quotes) ?; 
306387 } 
307388
308389 // Close the quote we left opened earlier. 
0 commit comments