Added binary messages

src/main/scala/com/twitter/finagle/HttpWebSocket.scala

@@ -11,10 +11,16 @@ import java.net.{SocketAddress, URI}
 
 trait WebSocketRichClient {
   def open(out: Offer[String], uri: String): Future[WebSocket] =
-    open(out, new URI(uri))
+    open(out, Offer.never, new URI(uri))
 
-  def open(out: Offer[String], uri: URI): Future[WebSocket] = {
-    val socket = WebSocket(messages = out, uri = uri)
+  def open(out: Offer[String], uri: URI): Future[WebSocket] =
+    open(out, Offer.never, uri)
+
+  def open(out: Offer[String], binaryOut: Offer[Array[Byte]], uri: String): Future[WebSocket] =
+    open(out, binaryOut, new URI(uri))
+
+  def open(out: Offer[String], binaryOut: Offer[Array[Byte]], uri: URI): Future[WebSocket] = {
+    val socket = WebSocket(messages = out, binaryMessages = binaryOut, uri = uri)
     val addr = uri.getHost + ":" + uri.getPort
     HttpWebSocket.newClient(addr).toService(socket)
   }

src/main/scala/com/twitter/finagle/websocket/WebSocket.scala

@@ -8,6 +8,7 @@ import java.net.SocketAddress
 
 case class WebSocket(
   messages: Offer[String],
+  binaryMessages: Offer[Array[Byte]],
   uri: URI,
   headers: Map[String, String] = Map.empty[String, String],
   remoteAddress: SocketAddress = new SocketAddress {},

src/main/scala/com/twitter/finagle/websocket/WebSocketHandler.scala

@@ -6,7 +6,7 @@ import com.twitter.finagle.netty3.Conversions._
 import com.twitter.finagle.netty3.{Cancelled, Ok, Error}
 import com.twitter.util.{Promise, Return, Throw, Try}
 import java.net.{URI, InetSocketAddress}
-import org.jboss.netty.buffer.ChannelBuffer
+import org.jboss.netty.buffer.ChannelBuffers
 import org.jboss.netty.channel._
 import org.jboss.netty.handler.codec.http.websocketx._
 import org.jboss.netty.handler.codec.http.{
@@ -15,6 +15,7 @@ import scala.collection.JavaConversions._
 
 class WebSocketHandler extends SimpleChannelHandler {
   protected[this] val messagesBroker = new Broker[String]
+  protected[this] val binaryMessagesBroker = new Broker[Array[Byte]]
   protected[this] val closer = new Promise[Unit]
 
   protected[this] def write(
@@ -37,12 +38,22 @@ class WebSocketHandler extends SimpleChannelHandler {
         }
 
       case None =>
-        sock.messages { message =>
-          val frame = new TextWebSocketFrame(message)
-          val writeFuture = Channels.future(ctx.getChannel)
-          Channels.write(ctx, writeFuture, frame)
-          write(ctx, sock, Some(writeFuture.toTwitterFuture.toOffer))
-        }
+        Offer.choose(
+          sock.messages {
+            message =>
+              val frame = new TextWebSocketFrame(message)
+              val writeFuture = Channels.future(ctx.getChannel)
+              Channels.write(ctx, writeFuture, frame)
+              write(ctx, sock, Some(writeFuture.toTwitterFuture.toOffer))
+          },
+          sock.binaryMessages {
+            binary =>
+              val frame = new BinaryWebSocketFrame(ChannelBuffers.wrappedBuffer(binary))
+              val writeFuture = Channels.future(ctx.getChannel)
+              Channels.write(ctx, writeFuture, frame)
+              write(ctx, sock, Some(writeFuture.toTwitterFuture.toOffer))
+          }
+        )
     }
     awaitAck.sync()
   }
@@ -71,6 +82,7 @@ class WebSocketServerHandler extends WebSocketHandler {
 
             val webSocket = WebSocket(
               messages = messagesBroker.recv,
+              binaryMessages = binaryMessagesBroker.recv,
               uri = new URI(req.getUri),
               headers = req.getHeaderNames().map(name => name -> req.getHeader(name)).toMap,
               remoteAddress = ctx.getChannel.getRemoteAddress,
@@ -91,6 +103,11 @@ class WebSocketServerHandler extends WebSocketHandler {
         ch.setReadable(false)
         (messagesBroker ! frame.getText) ensure { ch.setReadable(true) }
 
+      case frame: BinaryWebSocketFrame =>
+        val ch = ctx.getChannel
+        ch.setReadable(false)
+        (binaryMessagesBroker ! frame.getBinaryData.array) ensure { ch.setReadable(true) }
+
       case invalid =>
         Channels.fireExceptionCaught(ctx,
           new IllegalArgumentException("invalid message \"%s\"".format(invalid)))
@@ -136,6 +153,11 @@ class WebSocketClientHandler extends WebSocketHandler {
         ch.setReadable(false)
         (messagesBroker ! frame.getText) ensure { ch.setReadable(true) }
 
+      case frame: BinaryWebSocketFrame =>
+        val ch = ctx.getChannel
+        ch.setReadable(false)
+        (binaryMessagesBroker ! frame.getBinaryData.array) ensure { ch.setReadable(true) }
+
       case invalid =>
         Channels.fireExceptionCaught(ctx,
           new IllegalArgumentException("invalid message \"%s\"".format(invalid)))
@@ -152,6 +174,7 @@ class WebSocketClientHandler extends WebSocketHandler {
 
         val webSocket = sock.copy(
           messages = messagesBroker.recv,
+          binaryMessages = binaryMessagesBroker.recv,
           onClose = closer,
           close = close)
 

src/test/scala/com/twitter/finagle/websocket/EndToEndTest.scala

@@ -8,22 +8,29 @@ import com.twitter.concurrent.Broker
 import com.twitter.finagle.{HttpWebSocket, Service}
 import com.twitter.util._
 import java.net.InetSocketAddress
+import scala.collection.mutable.ArrayBuffer
 
 @RunWith(classOf[JUnitRunner])
 class EndToEndTest extends FunSuite {
   test("multi client") {
     var result = ""
+    val binaryResult = ArrayBuffer.empty[Byte]
     val addr = RandomSocket()
-    val latch = new CountDownLatch(5)
+    val latch = new CountDownLatch(10)
 
     val server = HttpWebSocket.serve(addr, new Service[WebSocket, WebSocket] {
       def apply(req: WebSocket): Future[WebSocket] = {
         val outgoing = new Broker[String]
-        val socket = req.copy(messages = outgoing.recv)
+        val binaryOutgoing = new Broker[Array[Byte]]
+        val socket = req.copy(messages = outgoing.recv, binaryMessages = binaryOutgoing.recv)
         req.messages foreach { msg =>
           synchronized { result += msg }
           latch.countDown()
         }
+        req.binaryMessages foreach { binary =>
+          synchronized { binaryResult ++= binary }
+          latch.countDown()
+        }
         Future.value(socket)
       }
     })
@@ -32,15 +39,18 @@ class EndToEndTest extends FunSuite {
 
     val brokers = (0 until 5) map { _ =>
       val out = new Broker[String]
-      Await.ready(HttpWebSocket.open(out.recv, target))
-      out
+      val binaryOut = new Broker[Array[Byte]]
+      Await.ready(HttpWebSocket.open(out.recv, binaryOut.recv, target))
+      (out, binaryOut)
     }
 
-    brokers foreach { out =>
-      FuturePool.unboundedPool { out !! "1" }
+    brokers foreach { pair =>
+      FuturePool.unboundedPool { pair._1 !! "1" }
+      FuturePool.unboundedPool { pair._2 !! Array[Byte](0x01) }
     }
 
     latch.within(1.second)
     assert(result === "11111")
+    assert(binaryResult === ArrayBuffer(0x01, 0x01, 0x01, 0x01, 0x01))
   }
 }