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 ptr;
  16 use str;
  17 use sync::Mutex;
  18 use sys::c;
  19 use sys::cvt;
  20 use sys::handle::Handle;
  21
  22 pub struct NoClose(Option<Handle>);
  23
  24 pub enum Output {
  25     Console(NoClose),
  26     Pipe(NoClose),
  27 }
  28
  29 pub struct Stdin {
  30     handle: Output,
  31     utf8: Mutex<io::Cursor<Vec<u8>>>,
  32 }
  33 pub struct Stdout(Output);
  34 pub struct Stderr(Output);
  35
  36 pub fn get(handle: c::DWORD) -> io::Result<Output> {
  37     let handle = unsafe { c::GetStdHandle(handle) };
  38     if handle == c::INVALID_HANDLE_VALUE {
  39         Err(io::Error::last_os_error())
  40     } else if handle.is_null() {
  41         Err(io::Error::new(io::ErrorKind::Other,
  42                            "no stdio handle available for this process"))
  43     } else {
  44         let ret = NoClose::new(handle);
  45         let mut out = 0;
  46         match unsafe { c::GetConsoleMode(handle, &mut out) } {
  47             0 => Ok(Output::Pipe(ret)),
  48             _ => Ok(Output::Console(ret)),
  49         }
  50     }
  51 }
  52
  53 fn write(out: &Output, data: &[u8]) -> io::Result<usize> {
  54     let handle = match *out {
  55         Output::Console(ref c) => c.get().raw(),
  56         Output::Pipe(ref p) => return p.get().write(data),
  57     };
  58     let utf16 = match str::from_utf8(data).ok() {
  59         Some(utf8) => utf8.encode_utf16().collect::<Vec<u16>>(),
  60         None => return Err(invalid_encoding()),
  61     };
  62     let mut written = 0;
  63     try!(cvt(unsafe {
  64         c::WriteConsoleW(handle,
  65                          utf16.as_ptr() as c::LPCVOID,
  66                          utf16.len() as u32,
  67                          &mut written,
  68                          ptr::null_mut())
  69     }));
  70
  71     // FIXME if this only partially writes the utf16 buffer then we need to
  72     //       figure out how many bytes of `data` were actually written
  73     assert_eq!(written as usize, utf16.len());
  74     Ok(data.len())
  75 }
  76
  77 impl Stdin {
  78     pub fn new() -> io::Result<Stdin> {
  79         get(c::STD_INPUT_HANDLE).map(|handle| {
  80             Stdin {
  81                 handle: handle,
  82                 utf8: Mutex::new(Cursor::new(Vec::new())),
  83             }
  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![0u16; 0x1000];
  96             let mut num = 0;
  97             try!(cvt(unsafe {
  98                 c::ReadConsoleW(handle,
  99                                 utf16.as_mut_ptr() as c::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() -> io::Result<Stdout> {
 120         get(c::STD_OUTPUT_HANDLE).map(Stdout)
 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() -> io::Result<Stderr> {
 130         get(c::STD_ERROR_HANDLE).map(Stderr)
 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: c::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::InvalidData, "text was not valid unicode")
 174 }