From 563c9ab34eae7b88c4d1796101e17bbc8d438021 Mon Sep 17 00:00:00 2001
From: dokaptur <dokaptur@gmail.com>
Date: Sun, 14 Dec 2014 01:20:49 +0100
Subject: [PATCH] preloading improvement

---
 src/stdbuf/libstdbuf.rs    | 13 ++++++++++---
 src/stdbuf/prepare_libs.sh |  6 +++---
 src/stdbuf/stdbuf.rs       | 35 +++++++++++++++++++++++++++++------
 3 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/src/stdbuf/libstdbuf.rs b/src/stdbuf/libstdbuf.rs
index 3187ef55c..f10775bb1 100644
--- a/src/stdbuf/libstdbuf.rs
+++ b/src/stdbuf/libstdbuf.rs
@@ -1,16 +1,23 @@
-#![crate_type = "dylib"]
+#![crate_name = "libstdbuf"]
+#![crate_type = "staticlib"]
+#![feature(macro_rules)]
 
 extern crate libc;
 use libc::{c_int, size_t, c_char, FILE, _IOFBF, _IONBF, _IOLBF, setvbuf};
 use std::ptr;
 use std::os;
 
+#[path = "../common/util.rs"]
+mod util;
+
 extern {
     static stdin: *mut FILE;
     static stdout: *mut FILE;
     static stderr: *mut FILE;
 }
 
+static NAME: &'static str = "libstdbuf";
+
 fn set_buffer(stream: *mut FILE, value: &str) {
     let (mode, size): (c_int, size_t) = match value {
         "0" => (_IONBF, 0 as size_t),
@@ -18,7 +25,7 @@ fn set_buffer(stream: *mut FILE, value: &str) {
         input => {
             let buff_size: uint = match from_str(input) {
                 Some(num) => num,
-                None => panic!("incorrect size of buffer!")
+                None => crash!(1, "incorrect size of buffer!")
             };
             (_IOFBF, buff_size as size_t)
         }
@@ -30,7 +37,7 @@ fn set_buffer(stream: *mut FILE, value: &str) {
         res = libc::setvbuf(stream, buffer, mode, size);
     }
     if res != 0 {
-        panic!("error while calling setvbuf!");
+        crash!(res, "error while calling setvbuf!");
     }
 }
 
diff --git a/src/stdbuf/prepare_libs.sh b/src/stdbuf/prepare_libs.sh
index bc0869123..6be6d920d 100755
--- a/src/stdbuf/prepare_libs.sh
+++ b/src/stdbuf/prepare_libs.sh
@@ -2,8 +2,8 @@
 
 cd $(dirname $0)
 rustc libstdbuf.rs
-gcc -c -Wall -Werror -fpic libstdbuf.c -L. -llibstdbuf.so
-gcc -shared -o libstdbuf.so libstdbuf.o
+gcc -c -Wall -Werror -fpic libstdbuf.c -L. -llibstdbuf.a
+gcc -shared -o libstdbuf.so -Wl,--whole-archive liblibstdbuf.a -Wl,--no-whole-archive libstdbuf.o -lpthread
 mv *.so ../../build/
 rm *.o
-export LD_LIBRARY_PATH="$LD_LIBRARY_PATH":"$PWD"/../../build
+rm *.a
diff --git a/src/stdbuf/stdbuf.rs b/src/stdbuf/stdbuf.rs
index 548e23351..e01a8c60f 100644
--- a/src/stdbuf/stdbuf.rs
+++ b/src/stdbuf/stdbuf.rs
@@ -16,22 +16,21 @@ use getopts::{optopt, optflag, getopts, usage, Matches, OptGroup};
 use std::io::process::{Command, StdioContainer};
 use std::iter::range_inclusive;
 use std::num::Int;
+use std::os;
 
 #[path = "../common/util.rs"]
 mod util;
 
 static NAME: &'static str = "stdbuf";
 static VERSION: &'static str = "1.0.0";
-static LIBSTDBUF_PATH: &'static str = "liblibstdbuf.so libstdbuf.so"; 
+static LIBSTDBUF: &'static str = "libstdbuf.so"; 
 
-#[deriving(Show)]
 enum BufferType {
     Default,
     Line,
     Size(u64)
 }
 
-#[deriving(Show)]
 struct ProgramOptions {
     stdin: BufferType,
     stdout: BufferType,
@@ -49,6 +48,22 @@ enum OkMsg {
     Version
 }
 
+#[cfg(target_os = "linux")]
+fn preload_env() -> &'static str { 
+    "LD_PRELOAD"
+}
+
+#[cfg(target_os = "macos")]
+fn preload_env() -> &'static str { 
+    "DYLD_INSERT_LIBRARIES"
+}
+
+#[cfg(not(any(target_os = "linux", target_os = "macos")))]
+fn preload_env() -> &'static str { 
+    crash!(1, "Command not supported for this operating system!")
+}
+
+
 fn print_version() {
     println!("{} version {}", NAME, VERSION);
 }
@@ -181,7 +196,7 @@ pub fn uumain(args: Vec<String>) -> int {
     for i in range_inclusive(1, args.len()) {
         match parse_options(args.slice(1, i), &mut options, &optgrps) {
             Ok(OkMsg::Buffering) => {
-                command_idx = i-1;
+                command_idx = i - 1;
                 break;
             },
             Ok(OkMsg::Help) => {
@@ -201,7 +216,16 @@ pub fn uumain(args: Vec<String>) -> int {
     }
     let ref command_name = args[command_idx];
     let mut command = Command::new(command_name);
-    command.args(args.slice_from(command_idx+1)).env("LD_PRELOAD", LIBSTDBUF_PATH);
+    let mut path = match os::self_exe_path() {
+        Some(exe_path) => exe_path,
+        None => crash!(1, "Impossible to fetch the path of this executable.")
+    };
+    path.push(LIBSTDBUF);
+    let libstdbuf = match path.as_str() {
+        Some(s) => s,
+        None => crash!(1, "Error while converting path.")
+    };
+    command.args(args.slice_from(command_idx+1)).env(preload_env(), libstdbuf);
     command.stdin(StdioContainer::InheritFd(0)).stdout(StdioContainer::InheritFd(1)).stderr(StdioContainer::InheritFd(2));
     set_command_env(&mut command, "_STDBUF_I", options.stdin);
     set_command_env(&mut command, "_STDBUF_O", options.stdout);
@@ -212,4 +236,3 @@ pub fn uumain(args: Vec<String>) -> int {
     };
     0
 }
-