Skip to content

Commit 2f26db6

Browse files
RuaRua
committed
transpile: Do not leak main argument strings
1 parent 4f00c7f commit 2f26db6

File tree

1 file changed

+89
-38
lines changed

1 file changed

+89
-38
lines changed

‎c2rust-transpile/src/translator/main_function.rs‎

Lines changed: 89 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
//! Rust.
55
66
use super::*;
7+
use c2rust_ast_builder::CaptureBy;
78
use failure::format_err;
89
use proc_macro2::{TokenStream, TokenTree};
910

@@ -53,45 +54,96 @@ impl<'c> Translation<'c> {
5354
if n >= 2 {
5455
// `argv` and `argc`
5556

56-
stmts.push(mk().local_stmt(Box::new(mk().local(
57-
mk().mutbl().ident_pat("args"),
58-
Some(mk().path_ty(vec![mk().path_segment_with_args(
57+
stmts.push(mk().local_stmt(Box::new({
58+
// Vec<Vec<u8>>
59+
let ty = mk().path_ty(vec![mk().path_segment_with_args(
5960
"Vec",
60-
mk().angle_bracketed_args(vec![
61-
mk().mutbl().ptr_ty(mk().abs_path_ty(vec!["core", "ffi", "c_char"])),
62-
]),
63-
)])),
64-
Some(mk().call_expr(mk().path_expr(vec!["Vec", "new"]), vec![])),
65-
))));
66-
stmts.push(mk().semi_stmt(mk().for_expr(
67-
mk().ident_pat("arg"),
68-
mk().call_expr(args_fn, vec![]),
69-
mk().block(vec![mk().semi_stmt(mk().method_call_expr(
70-
mk().path_expr(vec!["args"]),
71-
"push",
72-
vec![mk().method_call_expr(
73-
mk().method_call_expr(
74-
mk().call_expr(
75-
// TODO(kkysen) change `"std"` to `"alloc"` after `#![feature(alloc_c_string)]` is stabilized in `1.63.0`
76-
mk().abs_path_expr(vec!["std", "ffi", "CString", "new"]),
77-
vec![mk().path_expr(vec!["arg"])],
78-
),
79-
"expect",
80-
vec![mk().lit_expr("Failed to convert argument into CString.")],
81-
),
82-
"into_raw",
61+
mk().angle_bracketed_args(vec![mk().path_ty(vec![mk()
62+
.path_segment_with_args(
63+
"Vec",
64+
mk().angle_bracketed_args(vec![mk().ident_ty("u8")]),
65+
)])]),
66+
)]);
67+
// |arg| {
68+
// (::std::ffi::CString::new(arg))
69+
// .expect("Failed to convert argument into CString.")
70+
// .into_bytes_with_nul()
71+
// }
72+
let cstring_call = mk().call_expr(
73+
mk().abs_path_expr(vec!["std", "ffi", "CString", "new"]),
74+
vec![mk().path_expr(vec!["arg"])],
75+
);
76+
let expect_arg = mk().lit_expr("Failed to convert argument into CString.");
77+
let map_arg = mk().closure_expr(
78+
CaptureBy::Ref,
79+
Movability::Movable,
80+
vec![mk().ident_pat("arg")],
81+
ReturnType::Default,
82+
mk().method_chain_expr(
83+
cstring_call,
84+
vec![
85+
(mk().path_segment("expect"), vec![expect_arg]),
86+
(mk().path_segment("into_bytes_with_nul"), vec![]),
87+
],
88+
),
89+
);
90+
// (args_fn)
91+
// .map(map_arg)
92+
// .collect();
93+
let init = mk().method_chain_expr(
94+
mk().call_expr(args_fn, vec![]),
95+
vec![
96+
(mk().path_segment("map"), vec![map_arg]),
97+
(mk().path_segment("collect"), vec![]),
98+
],
99+
);
100+
mk().local(mk().mutbl().ident_pat("args_strings"), Some(ty), Some(init))
101+
})));
102+
103+
stmts.push(mk().local_stmt(Box::new({
104+
// Vec<*mut ::core::ffi::c_char>
105+
let ty = mk().path_ty(vec![mk().path_segment_with_args(
106+
"Vec",
107+
mk().angle_bracketed_args(vec![mk()
108+
.mutbl()
109+
.ptr_ty(mk().abs_path_ty(vec!["core", "ffi", "c_char"]))]),
110+
)]);
111+
// |arg| arg.as_mut_ptr() as *mut ::core::ffi::c_char
112+
let map_arg = mk().closure_expr(
113+
CaptureBy::Ref,
114+
Movability::Movable,
115+
vec![mk().ident_pat("arg")],
116+
ReturnType::Default,
117+
mk().cast_expr(
118+
mk().method_call_expr(mk().ident_expr("arg"), "as_mut_ptr", vec![]),
119+
mk().mutbl()
120+
.ptr_ty(mk().abs_path_ty(vec!["core", "ffi", "c_char"])),
121+
),
122+
);
123+
// ::core::iter::once(::core::ptr::null_mut())
124+
let chain_arg = mk().call_expr(
125+
mk().abs_path_expr(vec!["core", "iter", "once"]),
126+
vec![mk().call_expr(
127+
mk().abs_path_expr(vec!["core", "ptr", "null_mut"]),
83128
vec![],
84129
)],
85-
))]),
86-
None::<Ident>,
87-
)));
88-
stmts.push(mk().semi_stmt(mk().method_call_expr(
89-
mk().path_expr(vec!["args"]),
90-
"push",
91-
vec![
92-
mk().call_expr(mk().abs_path_expr(vec!["core", "ptr", "null_mut"]), vec![]),
93-
],
94-
)));
130+
);
131+
// args_strings
132+
// .iter_mut()
133+
// .map(map_arg)
134+
// .chain(chain_arg)
135+
// .collect()
136+
let init = mk().method_chain_expr(
137+
mk().ident_expr("args_strings"),
138+
vec![
139+
(mk().path_segment("iter_mut"), vec![]),
140+
(mk().path_segment("map"), vec![map_arg]),
141+
(mk().path_segment("chain"), vec![chain_arg]),
142+
(mk().path_segment("collect"), vec![]),
143+
],
144+
);
145+
mk().local(mk().mutbl().ident_pat("args_ptrs"), Some(ty), Some(init))
146+
})));
95147

96148
let argc_ty: Box<Type> = match self.ast_context.index(parameters[0]).kind {
97149
CDeclKind::Variable { ref typ, .. } => self.convert_type(typ.ctype),
@@ -105,8 +157,7 @@ impl<'c> Translation<'c> {
105157
"Cannot find type of 'argv' argument in main function",
106158
)),
107159
}?;
108-
109-
let args = mk().ident_expr("args");
160+
let args = mk().ident_expr("args_ptrs");
110161
let argc = mk().binary_expr(
111162
BinOp::Sub(Default::default()),
112163
mk().method_call_expr(args.clone(), "len", no_args.clone()),

0 commit comments

Comments
 (0)