--binary-prefix --text-prefix --base64 --base64-text

Partially addresses #78, #47.

Closes #20.

Author: vi

Cargo.lock

@@ -53,6 +53,7 @@ tokio-timer = "0.2.0"
 tempfile = "3.0.8"
 net2 = "0.2.33"
 anymap = "0.12.1"
+base64 = "0.10"
 #anymap = { path = "/mnt/src/git/anymap"}
 
 # Rust 1.30.1 compat:

Cargo.toml

@@ -227,7 +227,7 @@ Pre-built binaries for Linux (usual and musl), Windows, OS X and Android are ava
 Limitations
 ---
 
-* Websocat can't mix sending text and binary WebSocket messages. This also affects `mirror:` specifier, making it a bit different from ws://echo.websocket.org.
+* It is not convenient when text and binary WebSocket messages are mixed. This affects `mirror:` specifier, making it a bit different from ws://echo.websocket.org. There are `--binary-prefix`, `--text-prefix` and `--base64` options to handle mixture of binary and text.
 * Current version of Websocat don't receive notification about closed sockets. This makes serving without `-E` or `-u` options or in backpressure scenarios prone to socket leak.
 
 Building from source code

README.md

@@ -451,6 +451,30 @@ struct Opt {
     /// [A] Delay before reconnect attempt for `autoreconnect:` overlay.
     #[structopt(long = "--autoreconnect-delay-millis", default_value="20")]
     autoreconnect_delay_millis: u64,
+
+
+    /// [A] Prepend specified text to each received WebSocket text message.
+    /// Also strip this prefix from outgoing messages, explicitly marking
+    /// them as text even if `--binary` is specified
+    #[structopt(long = "--text-prefix")]
+    pub ws_text_prefix: Option<String>,
+    
+    /// [A] Prepend specified text to each received WebSocket binary message.
+    /// Also strip this prefix from outgoing messages, explicitly marking
+    /// them as binary even if `--text` is specified
+    #[structopt(long = "--binary-prefix")]
+    pub ws_binary_prefix: Option<String>,
+
+    /// Encode incoming binary WebSocket messages in one-line Base64
+    /// If `--binary-prefix` (see `--help=full`) is set, outgoing WebSocket messages
+    /// that start with the prefix are decoded from base64 prior to sending.
+    #[structopt(long = "--base64")]
+    pub ws_binary_base64: bool,
+
+    /// [A] Encode incoming text WebSocket messages in one-line Base64.
+    /// I don't know whether it can be ever useful, but it's for symmetry with `--base64`.
+    #[structopt(long = "--base64-text")]
+    pub ws_text_base64: bool,
 }
 
 // TODO: make it byte-oriented/OsStr?
@@ -693,6 +717,10 @@ fn run() -> Result<()> {
             max_messages
             max_messages_rev
             autoreconnect_delay_millis
+            ws_text_prefix
+            ws_binary_prefix
+            ws_binary_base64
+            ws_text_base64
         );
         #[cfg(feature = "ssl")]
         {

src/main.rs

@@ -83,4 +83,9 @@ pub struct Options {
     pub request_headers: Vec<(http::header::HeaderName, http::header::HeaderValue)>,
 
     pub autoreconnect_delay_millis: u64,
+
+    pub ws_text_prefix: Option<String>,
+    pub ws_binary_prefix: Option<String>,
+    pub ws_binary_base64: bool,
+    pub ws_text_base64: bool,
 }

src/options.rs

@@ -1,5 +1,6 @@
 extern crate tokio_codec;
 extern crate websocket;
+extern crate base64;
 
 use self::websocket::stream::r#async::Stream as WsStream;
 use self::websocket::OwnedMessage;
@@ -37,6 +38,11 @@ pub struct WsReadWrapper<T: WsStream + 'static> {
     pub debt: ReadDebt,
     pub pong_timeout: Option<(::tokio_timer::Delay, ::std::time::Duration)>,
     pub ping_aborter: Option<::futures::unsync::oneshot::Sender<()>>,
+
+    pub text_prefix: Option<String>,
+    pub binary_prefix: Option<String>,
+    pub binary_base64: bool,
+    pub text_base64: bool,
 }
 
 impl<T: WsStream + 'static> AsyncRead for WsReadWrapper<T> {}
@@ -54,6 +60,35 @@ impl<T: WsStream + 'static> Read for WsReadWrapper<T> {
                 brokenpipe()
             }};
         }
+        fn process_prefixes_and_base64<'a>(qbuf :&'a mut Vec<u8>, q: &mut &'a [u8], prefix: &Option<String>, base64: bool) {
+            match (prefix, base64) {
+                (None, false) => (),
+                (Some(pr), false) => {
+                    debug!("prepending prefix");
+                    qbuf.reserve_exact(pr.len() + q.len());
+                    qbuf.extend_from_slice(pr.as_bytes());
+                    qbuf.extend_from_slice(q);
+                    *q = &mut qbuf[..];
+                }
+                (None, true) => {
+                    debug!("encoding to base64");
+                    qbuf.resize(q.len() * 3 / 2 + 3, 0);
+                    let r = base64::encode_config_slice(q, base64::STANDARD, &mut qbuf[..]);
+                    qbuf.resize(r, 0);
+                    qbuf.push(b'\n');
+                    *q = &mut qbuf[..];
+                },
+                (Some(pr), true) => {
+                    debug!("prepending prefix and encoding to base64");
+                    qbuf.extend_from_slice(pr.as_bytes());
+                    qbuf.resize(pr.len() + q.len() * 3 / 2 + 3, 0);
+                    let r = base64::encode_config_slice(q, base64::STANDARD, &mut qbuf[pr.len()..]);
+                    qbuf.resize(pr.len()+r, 0);
+                    qbuf.push(b'\n');
+                    *q = &mut qbuf[..];
+                },
+            }
+        }
         loop {
             return match self.s.poll().map_err(io_other_error)? {
                 Ready(Some(OwnedMessage::Close(_))) => {
@@ -97,14 +132,20 @@ impl<T: WsStream + 'static> Read for WsReadWrapper<T> {
                 }
                 Ready(Some(OwnedMessage::Text(x))) => {
                     debug!("incoming text");
-                    match self.debt.process_message(buf, x.as_str().as_bytes()) {
+                    let mut qbuf : Vec<u8> = vec![];
+                    let mut q : &[u8] = x.as_str().as_bytes();
+                    process_prefixes_and_base64(&mut qbuf, &mut q, &self.text_prefix, self.text_base64);
+                    match self.debt.process_message(buf, q) {
                         ProcessMessageResult::Return(x) => x,
                         ProcessMessageResult::Recurse => continue,
                     }
                 }
                 Ready(Some(OwnedMessage::Binary(x))) => {
                     debug!("incoming binary");
-                    match self.debt.process_message(buf, x.as_slice()) {
+                    let mut qbuf : Vec<u8> = vec![];
+                    let mut q : &[u8] = x.as_slice();
+                    process_prefixes_and_base64(&mut qbuf, &mut q, &self.binary_prefix, self.binary_base64);
+                    match self.debt.process_message(buf, q) {
                         ProcessMessageResult::Return(x) => x,
                         ProcessMessageResult::Recurse => continue,
                     }
@@ -135,14 +176,23 @@ pub enum Mode1 {
     Binary,
 }
 
-pub struct WsWriteWrapper<T: WsStream + 'static>(pub MultiProducerWsSink<T>, pub Mode1, pub bool);
+pub struct WsWriteWrapper<T: WsStream + 'static> {
+    pub sink: MultiProducerWsSink<T>,
+    pub mode: Mode1,
+    pub close_on_shutdown: bool,
+
+    pub text_prefix: Option<String>,
+    pub binary_prefix: Option<String>,
+    pub binary_base64: bool,
+    pub text_base64: bool,
+}
 
 impl<T: WsStream + 'static> AsyncWrite for WsWriteWrapper<T> {
     fn shutdown(&mut self) -> futures::Poll<(), std::io::Error> {
-        if !self.2 {
+        if !self.close_on_shutdown {
             return Ok(Ready(()));
         }
-        let mut sink = self.0.borrow_mut();
+        let mut sink = self.sink.borrow_mut();
         match sink
             .start_send(OwnedMessage::Close(None))
             .map_err(io_other_error)?
@@ -160,16 +210,56 @@ impl<T: WsStream + 'static> AsyncWrite for WsWriteWrapper<T> {
 }
 
 impl<T: WsStream + 'static> Write for WsWriteWrapper<T> {
-    fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
-        let om = match self.1 {
+    fn write(&mut self, buf_: &[u8]) -> IoResult<usize> {
+        let bufv;
+        let mut effective_mode = self.mode;
+
+        let mut buf : &[u8] = buf_;
+
+        let origlen = buf.len();
+
+        if let Some(pr) = &self.text_prefix {
+            if buf.starts_with(pr.as_bytes()) {
+                effective_mode = Mode1::Text;
+                buf = &buf[pr.len()..];
+            }
+        }
+        if let Some(pr) = &self.binary_prefix {
+            if buf.starts_with(pr.as_bytes()) {
+                effective_mode = Mode1::Binary;
+                buf = &buf[pr.len()..];
+            }
+        }
+
+        let decode_base64 = match effective_mode {
+            Mode1::Binary => self.binary_base64,
+            Mode1::Text => self.text_base64,
+        };
+
+        if decode_base64 {
+            if buf.last() == Some(&b'\n') {
+                buf = &buf[..(buf.len()-1)];
+            }
+            if buf.last() == Some(&b'\r') {
+                buf = &buf[..(buf.len()-1)];
+            }
+            if let Ok(v) = base64::decode(buf) {
+                bufv = v;
+                buf = &bufv[..];
+            } else {
+                error!("Failed to decode user-supplised base64 buffer. Sending message as is.");
+            }
+        }
+
+        let om = match effective_mode {
             Mode1::Binary => OwnedMessage::Binary(buf.to_vec()),
             Mode1::Text => {
                 let text_tmp;
                 let text = match ::std::str::from_utf8(buf) {
                     Ok(x) => x,
                     Err(_) => {
                         error!(
-                            "Invalid UTF-8 in --text mode. Sending lossy data. May be \
+                            "Invalid UTF-8 in a text WebSocket message. Sending lossy data. May be \
                              caused by unlucky buffer splits."
                         );
                         text_tmp = String::from_utf8_lossy(buf);
@@ -179,14 +269,14 @@ impl<T: WsStream + 'static> Write for WsWriteWrapper<T> {
                 OwnedMessage::Text(text.to_string())
             }
         };
-        match self.0.borrow_mut().start_send(om).map_err(io_other_error)? {
+        match self.sink.borrow_mut().start_send(om).map_err(io_other_error)? {
             futures::AsyncSink::NotReady(_) => wouldblock(),
-            futures::AsyncSink::Ready => Ok(buf.len()),
+            futures::AsyncSink::Ready => Ok(origlen),
         }
     }
     fn flush(&mut self) -> IoResult<()> {
         match self
-            .0
+            .sink
             .borrow_mut()
             .poll_complete()
             .map_err(io_other_error)?
@@ -362,8 +452,21 @@ pub fn finish_building_ws_peer<S>(opts: &super::Options, duplex: Duplex<S>, clos
         debt: super::readdebt::ReadDebt(Default::default(), opts.read_debt_handling, zmsgh),
         pong_timeout,
         ping_aborter,
+        text_prefix: opts.ws_text_prefix.clone(),
+        binary_prefix: opts.ws_binary_prefix.clone(),
+        binary_base64: opts.ws_binary_base64,
+        text_base64: opts.ws_text_base64,
+    };
+    let ws_sin = WsWriteWrapper{
+        sink: mpsink,
+        mode: mode1,
+        close_on_shutdown,
+
+        text_prefix: opts.ws_text_prefix.clone(),
+        binary_prefix: opts.ws_binary_prefix.clone(),
+        binary_base64: opts.ws_binary_base64,
+        text_base64: opts.ws_text_base64,
     };
-    let ws_sin = WsWriteWrapper(mpsink, mode1, close_on_shutdown);
 
     Peer::new(ws_str, ws_sin, hup)
 }