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