src/libstd/io/stdio.rs

   1 // Copyright 2013 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 //! Non-blocking access to stdin, stdout, and stderr.
  12 //!
  13 //! This module provides bindings to the local event loop's TTY interface, using it
  14 //! to offer synchronous but non-blocking versions of stdio. These handles can be
  15 //! inspected for information about terminal dimensions or for related information
  16 //! about the stream or terminal to which it is attached.
  17 //!
  18 //! # Example
  19 //!
  20 //! ```rust
  21 //! # #![allow(unused_must_use)]
  22 //! use std::io;
  23 //!
  24 //! let mut out = io::stdout();
  25 //! out.write(b"Hello, world!");
  26 //! ```
  27
  28 use self::StdSource::*;
  29
  30 use boxed::Box;
  31 use cell::RefCell;
  32 use clone::Clone;
  33 use failure::LOCAL_STDERR;
  34 use fmt;
  35 use io::{Reader, Writer, IoResult, IoError, OtherIoError, Buffer,
  36          standard_error, EndOfFile, LineBufferedWriter, BufferedReader};
  37 use marker::{Sync, Send};
  38 use libc;
  39 use mem;
  40 use option::Option;
  41 use option::Option::{Some, None};
  42 use ops::{Deref, DerefMut, FnOnce};
  43 use ptr;
  44 use result::Result::{Ok, Err};
  45 use rt;
  46 use slice::SliceExt;
  47 use str::StrExt;
  48 use string::String;
  49 use sys::{fs, tty};
  50 use sync::{Arc, Mutex, MutexGuard, Once, ONCE_INIT};
  51 use uint;
  52 use vec::Vec;
  53
  54 // And so begins the tale of acquiring a uv handle to a stdio stream on all
  55 // platforms in all situations. Our story begins by splitting the world into two
  56 // categories, windows and unix. Then one day the creators of unix said let
  57 // there be redirection! And henceforth there was redirection away from the
  58 // console for standard I/O streams.
  59 //
  60 // After this day, the world split into four factions:
  61 //
  62 // 1. Unix with stdout on a terminal.
  63 // 2. Unix with stdout redirected.
  64 // 3. Windows with stdout on a terminal.
  65 // 4. Windows with stdout redirected.
  66 //
  67 // Many years passed, and then one day the nation of libuv decided to unify this
  68 // world. After months of toiling, uv created three ideas: TTY, Pipe, File.
  69 // These three ideas propagated throughout the lands and the four great factions
  70 // decided to settle among them.
  71 //
  72 // The groups of 1, 2, and 3 all worked very hard towards the idea of TTY. Upon
  73 // doing so, they even enhanced themselves further then their Pipe/File
  74 // brethren, becoming the dominant powers.
  75 //
  76 // The group of 4, however, decided to work independently. They abandoned the
  77 // common TTY belief throughout, and even abandoned the fledgling Pipe belief.
  78 // The members of the 4th faction decided to only align themselves with File.
  79 //
  80 // tl;dr; TTY works on everything but when windows stdout is redirected, in that
  81 //        case pipe also doesn't work, but magically file does!
  82 enum StdSource {
  83     TTY(tty::TTY),
  84     File(fs::FileDesc),
  85 }
  86
  87 fn src<T, F>(fd: libc::c_int, _readable: bool, f: F) -> T where
  88     F: FnOnce(StdSource) -> T,
  89 {
  90     match tty::TTY::new(fd) {
  91         Ok(tty) => f(TTY(tty)),
  92         Err(_) => f(File(fs::FileDesc::new(fd, false))),
  93     }
  94 }
  95
  96 thread_local! {
  97     static LOCAL_STDOUT: RefCell<Option<Box<Writer + Send>>> = {
  98         RefCell::new(None)
  99     }
 100 }
 101
 102 struct RaceBox(BufferedReader<StdReader>);
 103
 104 unsafe impl Send for RaceBox {}
 105 unsafe impl Sync for RaceBox {}
 106
 107 /// A synchronized wrapper around a buffered reader from stdin
 108 #[derive(Clone)]
 109 pub struct StdinReader {
 110     inner: Arc<Mutex<RaceBox>>,
 111 }
 112
 113 unsafe impl Send for StdinReader {}
 114 unsafe impl Sync for StdinReader {}
 115
 116 /// A guard for exclusive access to `StdinReader`'s internal `BufferedReader`.
 117 pub struct StdinReaderGuard<'a> {
 118     inner: MutexGuard<'a, RaceBox>,
 119 }
 120
 121 impl<'a> Deref for StdinReaderGuard<'a> {
 122     type Target = BufferedReader<StdReader>;
 123
 124     fn deref(&self) -> &BufferedReader<StdReader> {
 125         &self.inner.0
 126     }
 127 }
 128
 129 impl<'a> DerefMut for StdinReaderGuard<'a> {
 130     fn deref_mut(&mut self) -> &mut BufferedReader<StdReader> {
 131         &mut self.inner.0
 132     }
 133 }
 134
 135 impl StdinReader {
 136     /// Locks the `StdinReader`, granting the calling thread exclusive access
 137     /// to the underlying `BufferedReader`.
 138     ///
 139     /// This provides access to methods like `chars` and `lines`.
 140     ///
 141     /// # Examples
 142     ///
 143     /// ```rust
 144     /// use std::io;
 145     ///
 146     /// for line in io::stdin().lock().lines() {
 147     ///     println!("{}", line.unwrap());
 148     /// }
 149     /// ```
 150     pub fn lock<'a>(&'a mut self) -> StdinReaderGuard<'a> {
 151         StdinReaderGuard {
 152             inner: self.inner.lock().unwrap()
 153         }
 154     }
 155
 156     /// Like `Buffer::read_line`.
 157     ///
 158     /// The read is performed atomically - concurrent read calls in other
 159     /// threads will not interleave with this one.
 160     pub fn read_line(&mut self) -> IoResult<String> {
 161         self.inner.lock().unwrap().0.read_line()
 162     }
 163
 164     /// Like `Buffer::read_until`.
 165     ///
 166     /// The read is performed atomically - concurrent read calls in other
 167     /// threads will not interleave with this one.
 168     pub fn read_until(&mut self, byte: u8) -> IoResult<Vec<u8>> {
 169         self.inner.lock().unwrap().0.read_until(byte)
 170     }
 171
 172     /// Like `Buffer::read_char`.
 173     ///
 174     /// The read is performed atomically - concurrent read calls in other
 175     /// threads will not interleave with this one.
 176     pub fn read_char(&mut self) -> IoResult<char> {
 177         self.inner.lock().unwrap().0.read_char()
 178     }
 179 }
 180
 181 impl Reader for StdinReader {
 182     fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
 183         self.inner.lock().unwrap().0.read(buf)
 184     }
 185
 186     // We have to manually delegate all of these because the default impls call
 187     // read more than once and we don't want those calls to interleave (or
 188     // incur the costs of repeated locking).
 189
 190     fn read_at_least(&mut self, min: uint, buf: &mut [u8]) -> IoResult<uint> {
 191         self.inner.lock().unwrap().0.read_at_least(min, buf)
 192     }
 193
 194     fn push_at_least(&mut self, min: uint, len: uint, buf: &mut Vec<u8>) -> IoResult<uint> {
 195         self.inner.lock().unwrap().0.push_at_least(min, len, buf)
 196     }
 197
 198     fn read_to_end(&mut self) -> IoResult<Vec<u8>> {
 199         self.inner.lock().unwrap().0.read_to_end()
 200     }
 201
 202     fn read_le_uint_n(&mut self, nbytes: uint) -> IoResult<u64> {
 203         self.inner.lock().unwrap().0.read_le_uint_n(nbytes)
 204     }
 205
 206     fn read_be_uint_n(&mut self, nbytes: uint) -> IoResult<u64> {
 207         self.inner.lock().unwrap().0.read_be_uint_n(nbytes)
 208     }
 209 }
 210
 211 /// Creates a new handle to the stdin of the current process.
 212 ///
 213 /// The returned handle is a wrapper around a global `BufferedReader` shared
 214 /// by all threads. If buffered access is not desired, the `stdin_raw` function
 215 /// is provided to provided unbuffered access to stdin.
 216 ///
 217 /// See `stdout()` for more notes about this function.
 218 pub fn stdin() -> StdinReader {
 219     // We're following the same strategy as kimundi's lazy_static library
 220     static mut STDIN: *const StdinReader = 0 as *const StdinReader;
 221     static ONCE: Once = ONCE_INIT;
 222
 223     unsafe {
 224         ONCE.call_once(|| {
 225             // The default buffer capacity is 64k, but apparently windows doesn't like
 226             // 64k reads on stdin. See #13304 for details, but the idea is that on
 227             // windows we use a slightly smaller buffer that's been seen to be
 228             // acceptable.
 229             let stdin = if cfg!(windows) {
 230                 BufferedReader::with_capacity(8 * 1024, stdin_raw())
 231             } else {
 232                 BufferedReader::new(stdin_raw())
 233             };
 234             let stdin = StdinReader {
 235                 inner: Arc::new(Mutex::new(RaceBox(stdin)))
 236             };
 237             STDIN = mem::transmute(box stdin);
 238
 239             // Make sure to free it at exit
 240             rt::at_exit(|| {
 241                 mem::transmute::<_, Box<StdinReader>>(STDIN);
 242                 STDIN = ptr::null();
 243             });
 244         });
 245
 246         (*STDIN).clone()
 247     }
 248 }
 249
 250 /// Creates a new non-blocking handle to the stdin of the current process.
 251 ///
 252 /// Unlike `stdin()`, the returned reader is *not* a buffered reader.
 253 ///
 254 /// See `stdout()` for more notes about this function.
 255 pub fn stdin_raw() -> StdReader {
 256     src(libc::STDIN_FILENO, true, |src| StdReader { inner: src })
 257 }
 258
 259 /// Creates a line-buffered handle to the stdout of the current process.
 260 ///
 261 /// Note that this is a fairly expensive operation in that at least one memory
 262 /// allocation is performed. Additionally, this must be called from a runtime
 263 /// task context because the stream returned will be a non-blocking object using
 264 /// the local scheduler to perform the I/O.
 265 ///
 266 /// Care should be taken when creating multiple handles to an output stream for
 267 /// a single process. While usage is still safe, the output may be surprising if
 268 /// no synchronization is performed to ensure a sane output.
 269 pub fn stdout() -> LineBufferedWriter<StdWriter> {
 270     LineBufferedWriter::new(stdout_raw())
 271 }
 272
 273 /// Creates an unbuffered handle to the stdout of the current process
 274 ///
 275 /// See notes in `stdout()` for more information.
 276 pub fn stdout_raw() -> StdWriter {
 277     src(libc::STDOUT_FILENO, false, |src| StdWriter { inner: src })
 278 }
 279
 280 /// Creates a line-buffered handle to the stderr of the current process.
 281 ///
 282 /// See `stdout()` for notes about this function.
 283 pub fn stderr() -> LineBufferedWriter<StdWriter> {
 284     LineBufferedWriter::new(stderr_raw())
 285 }
 286
 287 /// Creates an unbuffered handle to the stderr of the current process
 288 ///
 289 /// See notes in `stdout()` for more information.
 290 pub fn stderr_raw() -> StdWriter {
 291     src(libc::STDERR_FILENO, false, |src| StdWriter { inner: src })
 292 }
 293
 294 /// Resets the task-local stdout handle to the specified writer
 295 ///
 296 /// This will replace the current task's stdout handle, returning the old
 297 /// handle. All future calls to `print` and friends will emit their output to
 298 /// this specified handle.
 299 ///
 300 /// Note that this does not need to be called for all new tasks; the default
 301 /// output handle is to the process's stdout stream.
 302 pub fn set_stdout(stdout: Box<Writer + Send>) -> Option<Box<Writer + Send>> {
 303     let mut new = Some(stdout);
 304     LOCAL_STDOUT.with(|slot| {
 305         mem::replace(&mut *slot.borrow_mut(), new.take())
 306     }).and_then(|mut s| {
 307         let _ = s.flush();
 308         Some(s)
 309     })
 310 }
 311
 312 /// Resets the task-local stderr handle to the specified writer
 313 ///
 314 /// This will replace the current task's stderr handle, returning the old
 315 /// handle. Currently, the stderr handle is used for printing panic messages
 316 /// during task panic.
 317 ///
 318 /// Note that this does not need to be called for all new tasks; the default
 319 /// output handle is to the process's stderr stream.
 320 pub fn set_stderr(stderr: Box<Writer + Send>) -> Option<Box<Writer + Send>> {
 321     let mut new = Some(stderr);
 322     LOCAL_STDERR.with(|slot| {
 323         mem::replace(&mut *slot.borrow_mut(), new.take())
 324     }).and_then(|mut s| {
 325         let _ = s.flush();
 326         Some(s)
 327     })
 328 }
 329
 330 // Helper to access the local task's stdout handle
 331 //
 332 // Note that this is not a safe function to expose because you can create an
 333 // aliased pointer very easily:
 334 //
 335 //  with_task_stdout(|io1| {
 336 //      with_task_stdout(|io2| {
 337 //          // io1 aliases io2
 338 //      })
 339 //  })
 340 fn with_task_stdout<F>(f: F) where F: FnOnce(&mut Writer) -> IoResult<()> {
 341     let mut my_stdout = LOCAL_STDOUT.with(|slot| {
 342         slot.borrow_mut().take()
 343     }).unwrap_or_else(|| {
 344         box stdout() as Box<Writer + Send>
 345     });
 346     let result = f(&mut *my_stdout);
 347     let mut var = Some(my_stdout);
 348     LOCAL_STDOUT.with(|slot| {
 349         *slot.borrow_mut() = var.take();
 350     });
 351     match result {
 352         Ok(()) => {}
 353         Err(e) => panic!("failed printing to stdout: {:?}", e),
 354     }
 355 }
 356
 357 /// Flushes the local task's stdout handle.
 358 ///
 359 /// By default, this stream is a line-buffering stream, so flushing may be
 360 /// necessary to ensure that all output is printed to the screen (if there are
 361 /// no newlines printed).
 362 ///
 363 /// Note that logging macros do not use this stream. Using the logging macros
 364 /// will emit output to stderr, and while they are line buffered the log
 365 /// messages are always terminated in a newline (no need to flush).
 366 pub fn flush() {
 367     with_task_stdout(|io| io.flush())
 368 }
 369
 370 /// Prints a string to the stdout of the current process. No newline is emitted
 371 /// after the string is printed.
 372 pub fn print(s: &str) {
 373     with_task_stdout(|io| io.write(s.as_bytes()))
 374 }
 375
 376 /// Prints a string to the stdout of the current process. A literal
 377 /// `\n` character is printed to the console after the string.
 378 pub fn println(s: &str) {
 379     with_task_stdout(|io| {
 380         io.write(s.as_bytes()).and_then(|()| io.write(&[b'\n']))
 381     })
 382 }
 383
 384 /// Similar to `print`, but takes a `fmt::Arguments` structure to be compatible
 385 /// with the `format_args!` macro.
 386 pub fn print_args(fmt: fmt::Arguments) {
 387     with_task_stdout(|io| write!(io, "{}", fmt))
 388 }
 389
 390 /// Similar to `println`, but takes a `fmt::Arguments` structure to be
 391 /// compatible with the `format_args!` macro.
 392 pub fn println_args(fmt: fmt::Arguments) {
 393     with_task_stdout(|io| writeln!(io, "{}", fmt))
 394 }
 395
 396 /// Representation of a reader of a standard input stream
 397 pub struct StdReader {
 398     inner: StdSource
 399 }
 400
 401 impl StdReader {
 402     /// Returns whether this stream is attached to a TTY instance or not.
 403     pub fn isatty(&self) -> bool {
 404         match self.inner {
 405             TTY(..) => true,
 406             File(..) => false,
 407         }
 408     }
 409 }
 410
 411 impl Reader for StdReader {
 412     fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
 413         let ret = match self.inner {
 414             TTY(ref mut tty) => {
 415                 // Flush the task-local stdout so that weird issues like a
 416                 // print!'d prompt not being shown until after the user hits
 417                 // enter.
 418                 flush();
 419                 tty.read(buf).map(|i| i as uint)
 420             },
 421             File(ref mut file) => file.read(buf).map(|i| i as uint),
 422         };
 423         match ret {
 424             // When reading a piped stdin, libuv will return 0-length reads when
 425             // stdin reaches EOF. For pretty much all other streams it will
 426             // return an actual EOF error, but apparently for stdin it's a
 427             // little different. Hence, here we convert a 0 length read to an
 428             // end-of-file indicator so the caller knows to stop reading.
 429             Ok(0) => { Err(standard_error(EndOfFile)) }
 430             ret @ Ok(..) | ret @ Err(..) => ret,
 431         }
 432     }
 433 }
 434
 435 /// Representation of a writer to a standard output stream
 436 pub struct StdWriter {
 437     inner: StdSource
 438 }
 439
 440 unsafe impl Send for StdWriter {}
 441 unsafe impl Sync for StdWriter {}
 442
 443 impl StdWriter {
 444     /// Gets the size of this output window, if possible. This is typically used
 445     /// when the writer is attached to something like a terminal, this is used
 446     /// to fetch the dimensions of the terminal.
 447     ///
 448     /// If successful, returns `Ok((width, height))`.
 449     ///
 450     /// # Error
 451     ///
 452     /// This function will return an error if the output stream is not actually
 453     /// connected to a TTY instance, or if querying the TTY instance fails.
 454     pub fn winsize(&mut self) -> IoResult<(int, int)> {
 455         match self.inner {
 456             TTY(ref mut tty) => {
 457                 tty.get_winsize()
 458             }
 459             File(..) => {
 460                 Err(IoError {
 461                     kind: OtherIoError,
 462                     desc: "stream is not a tty",
 463                     detail: None,
 464                 })
 465             }
 466         }
 467     }
 468
 469     /// Controls whether this output stream is a "raw stream" or simply a normal
 470     /// stream.
 471     ///
 472     /// # Error
 473     ///
 474     /// This function will return an error if the output stream is not actually
 475     /// connected to a TTY instance, or if querying the TTY instance fails.
 476     pub fn set_raw(&mut self, raw: bool) -> IoResult<()> {
 477         match self.inner {
 478             TTY(ref mut tty) => {
 479                 tty.set_raw(raw)
 480             }
 481             File(..) => {
 482                 Err(IoError {
 483                     kind: OtherIoError,
 484                     desc: "stream is not a tty",
 485                     detail: None,
 486                 })
 487             }
 488         }
 489     }
 490
 491     /// Returns whether this stream is attached to a TTY instance or not.
 492     pub fn isatty(&self) -> bool {
 493         match self.inner {
 494             TTY(..) => true,
 495             File(..) => false,
 496         }
 497     }
 498 }
 499
 500 impl Writer for StdWriter {
 501     fn write(&mut self, buf: &[u8]) -> IoResult<()> {
 502         // As with stdin on windows, stdout often can't handle writes of large
 503         // sizes. For an example, see #14940. For this reason, chunk the output
 504         // buffer on windows, but on unix we can just write the whole buffer all
 505         // at once.
 506         //
 507         // For some other references, it appears that this problem has been
 508         // encountered by others [1] [2]. We choose the number 8KB just because
 509         // libuv does the same.
 510         //
 511         // [1]: https://tahoe-lafs.org/trac/tahoe-lafs/ticket/1232
 512         // [2]: http://www.mail-archive.com/log4net-dev@logging.apache.org/msg00661.html
 513         let max_size = if cfg!(windows) {8192} else {uint::MAX};
 514         for chunk in buf.chunks(max_size) {
 515             try!(match self.inner {
 516                 TTY(ref mut tty) => tty.write(chunk),
 517                 File(ref mut file) => file.write(chunk),
 518             })
 519         }
 520         Ok(())
 521     }
 522 }
 523
 524 #[cfg(test)]
 525 mod tests {
 526     use prelude::v1::*;
 527
 528     use super::*;
 529     use sync::mpsc::channel;
 530     use thread::Thread;
 531
 532     #[test]
 533     fn smoke() {
 534         // Just make sure we can acquire handles
 535         stdin();
 536         stdout();
 537         stderr();
 538     }
 539
 540     #[test]
 541     fn capture_stdout() {
 542         use io::{ChanReader, ChanWriter};
 543
 544         let (tx, rx) = channel();
 545         let (mut r, w) = (ChanReader::new(rx), ChanWriter::new(tx));
 546         let _t = Thread::spawn(move|| {
 547             set_stdout(box w);
 548             println!("hello!");
 549         });
 550         assert_eq!(r.read_to_string().unwrap(), "hello!\n");
 551     }
 552
 553     #[test]
 554     fn capture_stderr() {
 555         use io::{ChanReader, ChanWriter, Reader};
 556
 557         let (tx, rx) = channel();
 558         let (mut r, w) = (ChanReader::new(rx), ChanWriter::new(tx));
 559         let _t = Thread::spawn(move || -> () {
 560             set_stderr(box w);
 561             panic!("my special message");
 562         });
 563         let s = r.read_to_string().unwrap();
 564         assert!(s.contains("my special message"));
 565     }
 566 }