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() -> Stdin {
  81         Stdin {
  82             handle: get(c::STD_INPUT_HANDLE).unwrap(),
  83             utf8: Mutex::new(Cursor::new(Vec::new())),
  84         }
  85     }
  86
  87     pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
  88         let handle = match self.handle {
  89             Output::Console(ref c) => c.get().raw(),
  90             Output::Pipe(ref p) => return p.get().read(buf),
  91         };
  92         let mut utf8 = self.utf8.lock().unwrap();
  93         // Read more if the buffer is empty
  94         if utf8.position() as usize == utf8.get_ref().len() {
  95             let mut utf16: Vec<u16> = repeat(0u16).take(0x1000).collect();
  96             let mut num = 0;
  97             try!(cvt(unsafe {
  98                 c::ReadConsoleW(handle,
  99                                 utf16.as_mut_ptr() as libc::LPVOID,
 100                                 utf16.len() as u32,
 101                                 &mut num,
 102                                 ptr::null_mut())
 103             }));
 104             utf16.truncate(num as usize);
 105             // FIXME: what to do about this data that has already been read?
 106             let data = match String::from_utf16(&utf16) {
 107                 Ok(utf8) => utf8.into_bytes(),
 108                 Err(..) => return Err(invalid_encoding()),
 109             };
 110             *utf8 = Cursor::new(data);
 111         }
 112
 113         // MemReader shouldn't error here since we just filled it
 114         utf8.read(buf)
 115     }
 116 }
 117
 118 impl Stdout {
 119     pub fn new() -> Stdout {
 120         Stdout(get(c::STD_OUTPUT_HANDLE).unwrap())
 121     }
 122
 123     pub fn write(&self, data: &[u8]) -> io::Result<usize> {
 124         write(&self.0, data)
 125     }
 126 }
 127
 128 impl Stderr {
 129     pub fn new() -> Stderr {
 130         Stderr(get(c::STD_ERROR_HANDLE).unwrap())
 131     }
 132
 133     pub fn write(&self, data: &[u8]) -> io::Result<usize> {
 134         write(&self.0, data)
 135     }
 136 }
 137
 138 // FIXME: right now this raw stderr handle is used in a few places because
 139 //        std::io::stderr_raw isn't exposed, but once that's exposed this impl
 140 //        should go away
 141 impl io::Write for Stderr {
 142     fn write(&mut self, data: &[u8]) -> io::Result<usize> {
 143         Stderr::write(self, data)
 144     }
 145     fn flush(&mut self) -> io::Result<()> { Ok(()) }
 146 }
 147
 148 impl NoClose {
 149     fn new(handle: libc::HANDLE) -> NoClose {
 150         NoClose(Some(Handle::new(handle)))
 151     }
 152
 153     fn get(&self) -> &Handle { self.0.as_ref().unwrap() }
 154 }
 155
 156 impl Drop for NoClose {
 157     fn drop(&mut self) {
 158         self.0.take().unwrap().into_raw();
 159     }
 160 }
 161
 162 impl Output {
 163     pub fn handle(&self) -> &Handle {
 164         let nc = match *self {
 165             Output::Console(ref c) => c,
 166             Output::Pipe(ref c) => c,
 167         };
 168         nc.0.as_ref().unwrap()
 169     }
 170 }
 171
 172 fn invalid_encoding() -> io::Error {
 173     io::Error::new(io::ErrorKind::InvalidInput, "text was not valid unicode")
 174 }