From c9b4538babbc85b971b19bbeff16bd12a4f4db54 Mon Sep 17 00:00:00 2001
From: Huon Wilson <dbau.pp+github@gmail.com>
Date: Tue, 18 Feb 2014 22:25:32 +1100
Subject: str: add a function for truncating a vector of u16 at NUL.

Many of the functions interacting with Windows APIs allocate a vector of
0's and do not retrieve a length directly from the API call, and so need
to be sure to remove the unmodified junk at the end of the vector.
---
 src/libstd/os.rs  | 10 +++++++---
 src/libstd/str.rs | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 3 deletions(-)

(limited to 'src/libstd')

diff --git a/src/libstd/os.rs b/src/libstd/os.rs
index 31e88905b30..74e2fceb6ca 100644
--- a/src/libstd/os.rs
+++ b/src/libstd/os.rs
@@ -88,7 +88,8 @@ pub fn getcwd() -> Path {
             fail!();
         }
     }
-    Path::new(str::from_utf16(buf).expect("GetCurrentDirectoryW returned invalid UTF-16"))
+    Path::new(str::from_utf16(str::truncate_utf16_at_nul(buf))
+              .expect("GetCurrentDirectoryW returned invalid UTF-16"))
 }
 
 #[cfg(windows)]
@@ -744,7 +745,8 @@ pub fn last_os_error() -> ~str {
                 fail!("[{}] FormatMessage failure", errno());
             }
 
-            str::from_utf16(buf).expect("FormatMessageW returned invalid UTF-16")
+            str::from_utf16(str::truncate_utf16_at_nul(buf))
+                .expect("FormatMessageW returned invalid UTF-16")
         }
     }
 
@@ -833,7 +835,9 @@ fn real_args() -> ~[~str] {
             while *ptr.offset(len as int) != 0 { len += 1; }
 
             // Push it onto the list.
-            let opt_s = vec::raw::buf_as_slice(ptr, len, str::from_utf16);
+            let opt_s = vec::raw::buf_as_slice(ptr, len, |buf| {
+                    str::from_utf16(str::truncate_utf16_at_nul(buf))
+                });
             args.push(opt_s.expect("CommandLineToArgvW returned invalid UTF-16"));
         }
     }
diff --git a/src/libstd/str.rs b/src/libstd/str.rs
index 20321dad600..1f94aaaa7c4 100644
--- a/src/libstd/str.rs
+++ b/src/libstd/str.rs
@@ -920,6 +920,32 @@ pub fn utf16_items<'a>(v: &'a [u16]) -> UTF16Items<'a> {
     UTF16Items { iter : v.iter() }
 }
 
+/// Return a slice of `v` ending at (and not including) the first NUL
+/// (0).
+///
+/// # Example
+///
+/// ```rust
+/// use std::str;
+///
+/// // "abcd"
+/// let mut v = ['a' as u16, 'b' as u16, 'c' as u16, 'd' as u16];
+/// // no NULs so no change
+/// assert_eq!(str::truncate_utf16_at_nul(v), v.as_slice());
+///
+/// // "ab\0d"
+/// v[2] = 0;
+/// assert_eq!(str::truncate_utf16_at_nul(v),
+///            &['a' as u16, 'b' as u16]);
+/// ```
+pub fn truncate_utf16_at_nul<'a>(v: &'a [u16]) -> &'a [u16] {
+    match v.iter().position(|c| *c == 0) {
+        // don't include the 0
+        Some(i) => v.slice_to(i),
+        None => v
+    }
+}
+
 /// Decode a UTF-16 encoded vector `v` into a string, returning `None`
 /// if `v` contains any invalid data.
 ///
@@ -3875,6 +3901,24 @@ mod tests {
         assert_eq!(from_utf16_lossy([0xD800, 0xd801, 0xdc8b, 0xD800]), ~"\uFFFD𐒋\uFFFD");
     }
 
+    #[test]
+    fn test_truncate_utf16_at_nul() {
+        let v = [];
+        assert_eq!(truncate_utf16_at_nul(v), &[]);
+
+        let v = [0, 2, 3];
+        assert_eq!(truncate_utf16_at_nul(v), &[]);
+
+        let v = [1, 0, 3];
+        assert_eq!(truncate_utf16_at_nul(v), &[1]);
+
+        let v = [1, 2, 0];
+        assert_eq!(truncate_utf16_at_nul(v), &[1, 2]);
+
+        let v = [1, 2, 3];
+        assert_eq!(truncate_utf16_at_nul(v), &[1, 2, 3]);
+    }
+
     #[test]
     fn test_char_at() {
         let s = ~"ศไทย中华Việt Nam";
-- 
cgit 1.4.1-3-g733a5