1 use crate::cli::{IsTerminal, StdoutStream}; 2 use crate::p2; 3 use bytes::Bytes; 4 use std::io::{self, Write}; 5 use std::pin::Pin; 6 use std::task::{Context, Poll}; 7 use tokio::io::AsyncWrite; 8 use wasmtime_wasi_io::streams::OutputStream; 9 10 // Implementation for tokio::io::Stdout 11 impl IsTerminal for tokio::io::Stdout { is_terminal(&self) -> bool12 fn is_terminal(&self) -> bool { 13 std::io::stdout().is_terminal() 14 } 15 } 16 impl StdoutStream for tokio::io::Stdout { p2_stream(&self) -> Box<dyn OutputStream>17 fn p2_stream(&self) -> Box<dyn OutputStream> { 18 Box::new(StdioOutputStream::Stdout) 19 } async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync>20 fn async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync> { 21 Box::new(StdioOutputStream::Stdout) 22 } 23 } 24 25 // Implementation for std::io::Stdout 26 impl IsTerminal for std::io::Stdout { is_terminal(&self) -> bool27 fn is_terminal(&self) -> bool { 28 std::io::IsTerminal::is_terminal(self) 29 } 30 } 31 impl StdoutStream for std::io::Stdout { p2_stream(&self) -> Box<dyn OutputStream>32 fn p2_stream(&self) -> Box<dyn OutputStream> { 33 Box::new(StdioOutputStream::Stdout) 34 } async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync>35 fn async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync> { 36 Box::new(StdioOutputStream::Stdout) 37 } 38 } 39 40 // Implementation for tokio::io::Stderr 41 impl IsTerminal for tokio::io::Stderr { is_terminal(&self) -> bool42 fn is_terminal(&self) -> bool { 43 std::io::stderr().is_terminal() 44 } 45 } 46 impl StdoutStream for tokio::io::Stderr { p2_stream(&self) -> Box<dyn OutputStream>47 fn p2_stream(&self) -> Box<dyn OutputStream> { 48 Box::new(StdioOutputStream::Stderr) 49 } async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync>50 fn async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync> { 51 Box::new(StdioOutputStream::Stderr) 52 } 53 } 54 55 // Implementation for std::io::Stderr 56 impl IsTerminal for std::io::Stderr { is_terminal(&self) -> bool57 fn is_terminal(&self) -> bool { 58 std::io::IsTerminal::is_terminal(self) 59 } 60 } 61 impl StdoutStream for std::io::Stderr { p2_stream(&self) -> Box<dyn OutputStream>62 fn p2_stream(&self) -> Box<dyn OutputStream> { 63 Box::new(StdioOutputStream::Stderr) 64 } async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync>65 fn async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync> { 66 Box::new(StdioOutputStream::Stderr) 67 } 68 } 69 70 enum StdioOutputStream { 71 Stdout, 72 Stderr, 73 } 74 75 impl OutputStream for StdioOutputStream { write(&mut self, bytes: Bytes) -> p2::StreamResult<()>76 fn write(&mut self, bytes: Bytes) -> p2::StreamResult<()> { 77 match self { 78 StdioOutputStream::Stdout => std::io::stdout().write_all(&bytes), 79 StdioOutputStream::Stderr => std::io::stderr().write_all(&bytes), 80 } 81 .map_err(|e| p2::StreamError::LastOperationFailed(wasmtime::format_err!(e))) 82 } 83 flush(&mut self) -> p2::StreamResult<()>84 fn flush(&mut self) -> p2::StreamResult<()> { 85 match self { 86 StdioOutputStream::Stdout => std::io::stdout().flush(), 87 StdioOutputStream::Stderr => std::io::stderr().flush(), 88 } 89 .map_err(|e| p2::StreamError::LastOperationFailed(wasmtime::format_err!(e))) 90 } 91 check_write(&mut self) -> p2::StreamResult<usize>92 fn check_write(&mut self) -> p2::StreamResult<usize> { 93 Ok(1024 * 1024) 94 } 95 } 96 97 impl AsyncWrite for StdioOutputStream { poll_write( self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8], ) -> Poll<io::Result<usize>>98 fn poll_write( 99 self: Pin<&mut Self>, 100 _cx: &mut Context<'_>, 101 buf: &[u8], 102 ) -> Poll<io::Result<usize>> { 103 Poll::Ready(match *self { 104 StdioOutputStream::Stdout => std::io::stdout().write(buf), 105 StdioOutputStream::Stderr => std::io::stderr().write(buf), 106 }) 107 } poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>>108 fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { 109 Poll::Ready(match *self { 110 StdioOutputStream::Stdout => std::io::stdout().flush(), 111 StdioOutputStream::Stderr => std::io::stderr().flush(), 112 }) 113 } poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>>114 fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { 115 Poll::Ready(Ok(())) 116 } 117 } 118 119 #[async_trait::async_trait] 120 impl p2::Pollable for StdioOutputStream { ready(&mut self)121 async fn ready(&mut self) {} 122 } 123