src/libstd/sys/windows/stdio.rs

   1 // Copyright 2015 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 use prelude::v1::*;
  12 use io::prelude::*;
  13
  14 use io::{self, Cursor};
  15 use iter::repeat;
  16 use libc;
  17 use ptr;
  18 use str;
  19 use sync::Mutex;
  20 use sys::c;
  21 use sys::cvt;
  22 use sys::handle::Handle;
  23
  24 pub struct NoClose(Option<Handle>);
  25
  26 pub enum Output {
  27     Console(NoClose),
  28     Pipe(NoClose),
  29 }
  30
  31 pub struct Stdin {
  32     handle: Output,
  33     utf8: Mutex<io::Cursor<Vec<u8>>>,
  34 }
  35 pub struct Stdout(Output);
  36 pub struct Stderr(Output);
  37
  38 pub fn get(handle: libc::DWORD) -> io::Result<Output> {
  39     let handle = unsafe { c::GetStdHandle(handle) };
  40     if handle == libc::INVALID_HANDLE_VALUE {
  41         Err(io::Error::last_os_error())
  42     } else if handle.is_null() {
  43         Err(io::Error::new(io::ErrorKind::Other,
  44                            "no stdio handle available for this process"))
  45     } else {
  46         let ret = NoClose::new(handle);
  47         let mut out = 0;
  48         match unsafe { c::GetConsoleMode(handle, &mut out) } {
  49             0 => Ok(Output::Pipe(ret)),
  50             _ => Ok(Output::Console(ret)),
  51         }
  52     }
  53 }
  54
  55 fn write(out: &Output, data: &[u8]) -> io::Result<usize> {
  56     let handle = match *out {
  57         Output::Console(ref c) => c.get().raw(),
  58         Output::Pipe(ref p) => return p.get().write(data),
  59     };
  60     let utf16 = match str::from_utf8(data).ok() {
  61         Some(utf8) => utf8.utf16_units().collect::<Vec<u16>>(),
  62         None => return Err(invalid_encoding()),
  63     };
  64     let mut written = 0;
  65     try!(cvt(unsafe {
  66         c::WriteConsoleW(handle,
  67                          utf16.as_ptr() as libc::LPCVOID,
  68                          utf16.len() as u32,
  69                          &mut written,
  70                          ptr::null_mut())
  71     }));
  72
  73     // FIXME if this only partially writes the utf16 buffer then we need to
  74     //       figure out how many bytes of `data` were actually written
  75     assert_eq!(written as usize, utf16.len());
  76     Ok(data.len())
  77 }
  78
  79 impl Stdin {
  80     pub fn new() -> io::Result<Stdin> {
  81         get(c::STD_INPUT_HANDLE).map(|handle| {
  82             Stdin {
  83                 handle: handle,
  84                 utf8: Mutex::new(Cursor::new(Vec::new())),
  85             }
  86         })
  87     }
  88
  89     pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
  90         let handle = match self.handle {
  91             Output::Console(ref c) => c.get().raw(),
  92             Output::Pipe(ref p) => return p.get().read(buf),
  93         };
  94         let mut utf8 = self.utf8.lock().unwrap();
  95         // Read more if the buffer is empty
  96         if utf8.position() as usize == utf8.get_ref().len() {
  97             let mut utf16: Vec<u16> = repeat(0u16).take(0x1000).collect();
  98             let mut num = 0;
  99             try!(cvt(unsafe {
 100                 c::ReadConsoleW(handle,
 101                                 utf16.as_mut_ptr() as libc::LPVOID,
 102                                 utf16.len() as u32,
 103                                 &mut num,
 104                                 ptr::null_mut())
 105             }));
 106             utf16.truncate(num as usize);
 107             // FIXME: what to do about this data that has already been read?
 108             let data = match String::from_utf16(&utf16) {
 109                 Ok(utf8) => utf8.into_bytes(),
 110                 Err(..) => return Err(invalid_encoding()),
 111             };
 112             *utf8 = Cursor::new(data);
 113         }
 114
 115         // MemReader shouldn't error here since we just filled it
 116         utf8.read(buf)
 117     }
 118 }
 119
 120 impl Stdout {
 121     pub fn new() -> io::Result<Stdout> {
 122         get(c::STD_OUTPUT_HANDLE).map(Stdout)
 123     }
 124
 125     pub fn write(&self, data: &[u8]) -> io::Result<usize> {
 126         write(&self.0, data)
 127     }
 128 }
 129
 130 impl Stderr {
 131     pub fn new() -> io::Result<Stderr> {
 132         get(c::STD_ERROR_HANDLE).map(Stderr)
 133     }
 134
 135     pub fn write(&self, data: &[u8]) -> io::Result<usize> {
 136         write(&self.0, data)
 137     }
 138 }
 139
 140 // FIXME: right now this raw stderr handle is used in a few places because
 141 //        std::io::stderr_raw isn't exposed, but once that's exposed this impl
 142 //        should go away
 143 impl io::Write for Stderr {
 144     fn write(&mut self, data: &[u8]) -> io::Result<usize> {
 145         Stderr::write(self, data)
 146     }
 147     fn flush(&mut self) -> io::Result<()> { Ok(()) }
 148 }
 149
 150 impl NoClose {
 151     fn new(handle: libc::HANDLE) -> NoClose {
 152         NoClose(Some(Handle::new(handle)))
 153     }
 154
 155     fn get(&self) -> &Handle { self.0.as_ref().unwrap() }
 156 }
 157
 158 impl Drop for NoClose {
 159     fn drop(&mut self) {
 160         self.0.take().unwrap().into_raw();
 161     }
 162 }
 163
 164 impl Output {
 165     pub fn handle(&self) -> &Handle {
 166         let nc = match *self {
 167             Output::Console(ref c) => c,
 168             Output::Pipe(ref c) => c,
 169         };
 170         nc.0.as_ref().unwrap()
 171     }
 172 }
 173
 174 fn invalid_encoding() -> io::Error {
 175     io::Error::new(io::ErrorKind::InvalidData, "text was not valid unicode")
 176 }