Implement buffered transport on binary protocol

README.md

@@ -7,7 +7,7 @@
 
 This package contains an implementation of [Thrift](https://thrift.apache.org/)
 for Elixir. It includes a Thrift IDL parser, an Elixir code generator, and
-binary framed client and server implementations.
+binary protocol clients and servers with framed and buffered transports.
 
 The generated serialization code is highly optimized and has been measured at
 **10 and 25 times faster**<sub>[why?](#why-is-it-faster-than-the-apache-implementation)</sub>
@@ -50,6 +50,23 @@ defmodule MyProject.Mixfile do
 end
 ```
 
+## Buffered Transport
+
+Generated services also provide `Binary.Buffered.Client` and
+`Binary.Buffered.Server`, with the same RPC methods and connection options
+as their framed counterparts:
+
+```elixir
+{:ok, server} = MyService.Binary.Buffered.Server.start_link(MyHandler, 2345)
+{:ok, client} = MyService.Binary.Buffered.Client.start_link("localhost", 2345)
+```
+
+Both endpoints must use the same transport. Buffered transport sends binary
+protocol messages without frame-length prefixes, writing each RPC in one send.
+It accepts versioned and legacy headers and handles messages split across TCP
+reads. Incoming messages are limited to 16 MiB and 64 levels of nesting.
+Regenerate existing service code to obtain the buffered modules.
+
 ## RPC Service Support
 
 We provide full client and server support for Thrift RPC services. The examples

lib/thrift/binary/buffered/client.ex

@@ -0,0 +1,20 @@
+defmodule Thrift.Binary.Buffered.Client do
+  @moduledoc """
+  Binary protocol client using buffered (unframed) transport.
+
+  Writes each RPC in one send and reads a complete binary protocol message.
+  Incoming messages are limited to 16 MiB and 64 levels of nesting.
+
+  Options are the same as `Thrift.Binary.Framed.Client.start_link/3`.
+  Generated services expose this as `YourService.Binary.Buffered.Client`.
+  """
+  alias Thrift.Binary.Framed.Client
+
+  def start_link(host, port, opts \\ []),
+    do: Client.start_link(host, port, Keyword.put(opts, :packet, 0))
+
+  defdelegate close(conn), to: Client
+  defdelegate connect(conn, opts), to: Client
+  defdelegate call(conn, name, args, module, opts), to: Client
+  defdelegate oneway(conn, name, args, opts), to: Client
+end

lib/thrift/binary/buffered/message.ex

@@ -0,0 +1,102 @@
+defmodule Thrift.Binary.Buffered.Message do
+  @moduledoc false
+
+  # Exact reads leave subsequent messages in the socket buffer. Bound both
+  # message size and nesting, and use one deadline for the entire message.
+  @max_size 16 * 1024 * 1024
+  @max_depth 64
+
+  def recv(transport, socket, timeout) do
+    deadline = if timeout == :infinity, do: :infinity, else: now() + timeout
+    state = {transport, socket, deadline, [], 0}
+
+    try do
+      {header, state} = read(state, 4)
+
+      state =
+        case header do
+          <<0x80, 1, 0, type>> when type in 1..4 ->
+            state = value(11, state, 0)
+            elem(read(state, 4), 1)
+
+          <<size::signed-32>> when size >= 0 ->
+            {_, state} = read(state, size)
+            {<<type>>, state} = read(state, 1)
+            if type not in 1..4, do: throw({:error, :invalid_message})
+            elem(read(state, 4), 1)
+
+          _ ->
+            throw({:error, :invalid_message})
+        end
+
+      {_, _, _, chunks, _} = value(12, state, 0)
+      {:ok, chunks |> Enum.reverse() |> IO.iodata_to_binary()}
+    catch
+      {:error, _} = error -> error
+    end
+  end
+
+  defp value(_, _, depth) when depth > @max_depth, do: throw({:error, :invalid_message})
+  defp value(type, state, _) when type in [2, 3], do: elem(read(state, 1), 1)
+  defp value(6, state, _), do: elem(read(state, 2), 1)
+  defp value(8, state, _), do: elem(read(state, 4), 1)
+  defp value(type, state, _) when type in [4, 10], do: elem(read(state, 8), 1)
+
+  defp value(11, state, _) do
+    {size, state} = count(state)
+    elem(read(state, size), 1)
+  end
+
+  defp value(12, state, depth) do
+    case read(state, 1) do
+      {<<0>>, state} ->
+        state
+
+      {<<type>>, state} ->
+        {_, state} = read(state, 2)
+        state = value(type, state, depth + 1)
+        value(12, state, depth)
+    end
+  end
+
+  defp value(13, state, depth) do
+    {<<key, val>>, state} = read(state, 2)
+    {size, state} = count(state)
+    repeat(size, [key, val], state, depth + 1)
+  end
+
+  defp value(type, state, depth) when type in [14, 15] do
+    {<<element>>, state} = read(state, 1)
+    {size, state} = count(state)
+    repeat(size, [element], state, depth + 1)
+  end
+
+  defp value(_, _, _), do: throw({:error, :invalid_message})
+
+  defp repeat(0, _, state, _), do: state
+
+  defp repeat(count, types, state, depth) do
+    state = Enum.reduce(types, state, &value(&1, &2, depth))
+    repeat(count - 1, types, state, depth)
+  end
+
+  defp count(state) do
+    {<<size::signed-32>>, state} = read(state, 4)
+    if size < 0 or size > @max_size, do: throw({:error, :invalid_message})
+    {size, state}
+  end
+
+  defp read(state, 0), do: {<<>>, state}
+
+  defp read({transport, socket, deadline, chunks, total}, size) do
+    if total + size > @max_size, do: throw({:error, :message_too_large})
+    timeout = if deadline == :infinity, do: :infinity, else: max(deadline - now(), 0)
+
+    case transport.recv(socket, size, timeout) do
+      {:ok, bytes} -> {bytes, {transport, socket, deadline, [bytes | chunks], total + size}}
+      {:error, _} = error -> throw(error)
+    end
+  end
+
+  defp now, do: System.monotonic_time(:millisecond)
+end

lib/thrift/binary/buffered/server.ex

@@ -0,0 +1,16 @@
+defmodule Thrift.Binary.Buffered.Server do
+  @moduledoc """
+  Binary protocol server using buffered (unframed) transport.
+
+  Incoming messages are limited to 16 MiB and 64 levels of nesting.
+
+  Options are the same as `Thrift.Binary.Framed.Server.start_link/4`.
+  Generated services expose this as `YourService.Binary.Buffered.Server`.
+  """
+  alias Thrift.Binary.Framed.Server
+
+  def start_link(server_module, port, handler_module, opts \\ []),
+    do: Server.start_link(server_module, port, handler_module, Keyword.put(opts, :packet, 0))
+
+  defdelegate stop(pid), to: Server
+end

lib/thrift/binary/framed/client.ex

@@ -55,7 +55,8 @@ defmodule Thrift.Binary.Framed.Client do
             ssl_opts: [SSL.option()],
             timeout: integer,
             sock: {:gen_tcp, :gen_tcp.socket()} | {:ssl, :ssl.sslsocket()},
-            seq_id: integer
+            seq_id: integer,
+            packet: 0 | 4
           }
     defstruct host: nil,
               port: nil,
@@ -64,7 +65,8 @@ defmodule Thrift.Binary.Framed.Client do
               ssl_opts: nil,
               timeout: 5000,
               sock: nil,
-              seq_id: 0
+              seq_id: 0,
+              packet: 4
   end
 
   require Logger
@@ -82,7 +84,8 @@ defmodule Thrift.Binary.Framed.Client do
       port: port,
       tcp_opts: tcp_opts,
       ssl_opts: ssl_opts,
-      timeout: timeout
+      timeout: timeout,
+      packet: Keyword.get(opts, :packet, 4)
     }
 
     {:connect, :init, s}
@@ -141,6 +144,7 @@ defmodule Thrift.Binary.Framed.Client do
     opts =
       opts
       |> Keyword.merge(@immutable_tcp_opts)
+      |> Keyword.put(:packet, s.packet)
       |> Keyword.put_new(:send_timeout, 1000)
 
     case :gen_tcp.connect(host, port, opts, timeout) do
@@ -253,7 +257,7 @@ defmodule Thrift.Binary.Framed.Client do
     timeout = Keyword.get(tcp_opts, :timeout, default_timeout)
 
     with :ok <- transport.send(sock, [message | serialized_args]),
-         {:ok, message} <- transport.recv(sock, 0, timeout) do
+         {:ok, message} <- receive_message(transport, sock, timeout, s.packet) do
       reply = deserialize_message_reply(message, rpc_name, seq_id)
       {:reply, reply, s}
     else
@@ -291,6 +295,12 @@ defmodule Thrift.Binary.Framed.Client do
     end
   end
 
+  defp receive_message(transport, socket, timeout, 0),
+    do: Thrift.Binary.Buffered.Message.recv(transport, socket, timeout)
+
+  defp receive_message(transport, socket, timeout, 4),
+    do: transport.recv(socket, 0, timeout)
+
   def deserialize_message_reply(message, rpc_name, seq_id) do
     handle_message(Binary.deserialize(:message_begin, message), seq_id, rpc_name)
   end
@@ -333,9 +343,7 @@ defmodule Thrift.Binary.Framed.Client do
       TApplicationException.exception(
         type: :bad_sequence_id,
         message:
-          "Invalid sequence id. The client sent #{seq_id}, but the server replied with #{
-            mismatched_seq_id
-          }"
+          "Invalid sequence id. The client sent #{seq_id}, but the server replied with #{mismatched_seq_id}"
       )
 
     {:error, {:exception, exception}}

lib/thrift/binary/framed/protocol_handler.ex

@@ -25,7 +25,8 @@ defmodule Thrift.Binary.Framed.ProtocolHandler do
     :server_module,
     :handler_module,
     :on_connect,
-    :recv_timeout
+    :recv_timeout,
+    :packet
   ]
 
   @spec start_link(reference, port, transport, {module, module, transport_opts, [SSL.option()]}) ::
@@ -60,15 +61,17 @@ defmodule Thrift.Binary.Framed.ProtocolHandler do
     :ok = :ranch.accept_ack(ref)
 
     {recv_timeout, tcp_opts} = Keyword.pop(tcp_opts, :recv_timeout, @default_timeout)
-    :ok = :inet.setopts(socket, Keyword.put(tcp_opts, :packet, 4))
+    packet = Keyword.get(tcp_opts, :packet, 4)
+    :ok = :inet.setopts(socket, tcp_opts)
 
     state = %__MODULE__{
       transport: :gen_tcp,
       socket: socket,
       server_module: server_module,
       handler_module: handler_module,
       on_connect: on_connect,
-      recv_timeout: recv_timeout
+      recv_timeout: recv_timeout,
+      packet: packet
     }
 
     case SSL.configuration(ssl_opts) do
@@ -97,7 +100,7 @@ defmodule Thrift.Binary.Framed.ProtocolHandler do
     with :ok <- :inet.setopts(socket, packet: :raw),
          {:ok, first_bytes} <- :gen_tcp.recv(socket, 4, recv_timeout),
          :ok <- :gen_tcp.unrecv(socket, first_bytes),
-         :ok <- :inet.setopts(socket, packet: 4) do
+         :ok <- :inet.setopts(socket, packet: state.packet) do
       <<first_byte::8-unsigned, _::binary>> = first_bytes
       {:ok, first_byte}
     end
@@ -114,15 +117,15 @@ defmodule Thrift.Binary.Framed.ProtocolHandler do
         finish_span(span, result: "ssl_rejected")
         close(state)
 
-      {:ok, @tcp_header_byte} when optional == :required ->
+      {:ok, byte} when byte in [@tcp_header_byte, 0x80] and optional == :required ->
         Logger.error(fn ->
           format_log("peek_first_byte", "SSL required", state)
         end)
 
         finish_span(span, result: "tcp_rejected")
         close(state)
 
-      {:ok, @tcp_header_byte} when optional in [:disabled, :optional] ->
+      {:ok, byte} when byte in [@tcp_header_byte, 0x80] and optional in [:disabled, :optional] ->
         finish_span(span, result: "tcp")
         on_connect(state)
 
@@ -200,7 +203,7 @@ defmodule Thrift.Binary.Framed.ProtocolHandler do
        ) do
     span = start_span(:receive_message, state)
 
-    case transport.recv(socket, 0, recv_timeout) do
+    case receive_packet(transport, socket, recv_timeout, state.packet) do
       {:ok, serialized_message} ->
         finish_span(span, result: "success")
         deserialize_message(serialized_message, state)
@@ -219,6 +222,12 @@ defmodule Thrift.Binary.Framed.ProtocolHandler do
     end
   end
 
+  defp receive_packet(transport, socket, timeout, 0),
+    do: Thrift.Binary.Buffered.Message.recv(transport, socket, timeout)
+
+  defp receive_packet(transport, socket, timeout, 4),
+    do: transport.recv(socket, 0, timeout)
+
   defp deserialize_message(serialized_message, state) do
     case Binary.deserialize(:message_begin, serialized_message) do
       {:ok, message} ->
@@ -293,7 +302,11 @@ defmodule Thrift.Binary.Framed.ProtocolHandler do
       server_module.handle_thrift(method_name, args_binary, handler_module)
     end)
 
-    send_reply(<<0>>, :continue, state)
+    if state.packet == 0 do
+      receive_message(state)
+    else
+      send_reply(<<0>>, :continue, state)
+    end
   end
 
   defp send_reply(data, after_reply, %__MODULE__{transport: transport, socket: socket} = state) do

lib/thrift/binary/framed/server.ex

@@ -47,7 +47,10 @@ defmodule Thrift.Binary.Framed.Server do
   def start_link(server_module, port, handler_module, opts) do
     name = Keyword.get(opts, :name, handler_module)
     worker_count = Keyword.get(opts, :worker_count, 1)
-    tcp_opts = Keyword.get(opts, :tcp_opts, [])
+
+    tcp_opts =
+      Keyword.put(Keyword.get(opts, :tcp_opts, []), :packet, Keyword.get(opts, :packet, 4))
+
     ssl_opts = Keyword.get(opts, :ssl_opts, [])
     on_connect = Keyword.get(opts, :on_connect)
 

lib/thrift/generator/binary/framed/client.ex

@@ -4,18 +4,21 @@ defmodule Thrift.Generator.Binary.Framed.Client do
   alias Thrift.AST.Function
   alias Thrift.Generator.{Service, Utils}
 
-  def generate(service) do
+  def generate(service, transport \\ :framed) do
     functions =
       service.functions
       |> Map.values()
       |> Enum.map(&generate_handler_function(&1))
       |> Utils.merge_blocks()
 
+    module = Module.concat([Binary, Macro.camelize(Atom.to_string(transport)), Client])
+    implementation = Module.concat([Thrift, module])
+
     quote do
-      defmodule Binary.Framed.Client do
+      defmodule unquote(module) do
         @moduledoc false
 
-        alias Thrift.Binary.Framed.Client, as: ClientImpl
+        alias unquote(implementation), as: ClientImpl
 
         defdelegate close(conn), to: ClientImpl
         defdelegate connect(conn, opts), to: ClientImpl

lib/thrift/generator/binary/framed/server.ex

@@ -9,18 +9,21 @@ defmodule Thrift.Generator.Binary.Framed.Server do
 
   alias Thrift.Parser.FileGroup
 
-  def generate(service_module, service, file_group) do
+  def generate(service_module, service, file_group, transport \\ :framed) do
     functions =
       service.functions
       |> Map.values()
       |> Enum.map(&generate_handler_function(file_group, service_module, &1))
 
+    module = Module.concat([Binary, Macro.camelize(Atom.to_string(transport)), Server])
+    implementation = Module.concat([Thrift, module])
+
     quote do
-      defmodule Binary.Framed.Server do
+      defmodule unquote(module) do
         @moduledoc false
         require Logger
 
-        alias Thrift.Binary.Framed.Server, as: ServerImpl
+        alias unquote(implementation), as: ServerImpl
         defdelegate stop(name), to: ServerImpl
 
         def start_link(handler_module, port, opts \\ []) do