Http upgrades

async/httpaf_async.ml

@@ -90,14 +90,28 @@ let read fd buffer =
 open Httpaf
 
 module Server = struct
-  let create_connection_handler ?(config=Config.default) ~request_handler ~error_handler =
+  let create_connection_handler
+        ?(config=Config.default) ~request_handler ~error_handler ~upgrade_handler =
     fun client_addr socket ->
       let fd     = Socket.fd socket in
       let writev = Faraday_async.writev_of_fd fd in
       let request_handler = request_handler client_addr in
       let error_handler   = error_handler client_addr in
       let conn = Server_connection.create ~config ~error_handler request_handler in
       let read_complete = Ivar.create () in
+      let write_complete = Ivar.create () in
+      let upgrade_read, upgrade_write = Ivar.create (), Ivar.create () in
+      upon
+        (Deferred.both (Ivar.read upgrade_read) (Ivar.read upgrade_write))
+        (fun ((), ()) ->
+          match upgrade_handler with
+          | None -> failwith "HTTP upgrades not supported"
+          | Some upgrade_handler ->
+            upgrade_handler client_addr socket
+            >>> fun () ->
+            if not (Fd.is_closed fd) then Socket.shutdown socket `Both;
+            Ivar.fill read_complete ();
+            Ivar.fill write_complete ());
       let buffer = Buffer.create config.read_buffer_size in
       let rec reader_thread () =
         match Server_connection.next_read_operation conn with
@@ -119,13 +133,13 @@ module Server = struct
         | `Yield  ->
           (* Log.Global.printf "read_yield(%d)%!" (Fd.to_int_exn fd); *)
           Server_connection.yield_reader conn reader_thread
+        | `Upgrade -> Ivar.fill upgrade_read ()
         | `Close ->
           (* Log.Global.printf "read_close(%d)%!" (Fd.to_int_exn fd); *)
           Ivar.fill read_complete ();
           if not (Fd.is_closed fd)
           then Socket.shutdown socket `Receive
       in
-      let write_complete = Ivar.create () in
       let rec writer_thread () =
         match Server_connection.next_write_operation conn with
         | `Write iovecs ->
@@ -136,6 +150,7 @@ module Server = struct
         | `Yield ->
           (* Log.Global.printf "write_yield(%d)%!" (Fd.to_int_exn fd); *)
           Server_connection.yield_writer conn writer_thread;
+        | `Upgrade -> Ivar.fill upgrade_write ()
         | `Close _ ->
           (* Log.Global.printf "write_close(%d)%!" (Fd.to_int_exn fd); *)
           Ivar.fill write_complete ();

async/httpaf_async.mli

@@ -8,6 +8,7 @@ module Server : sig
     :  ?config         : Config.t
     -> request_handler : ('a -> Server_connection.request_handler)
     -> error_handler   : ('a -> Server_connection.error_handler)
+    -> upgrade_handler : ('a -> ([`Active], 'a) Socket.t -> unit Deferred.t) option
     -> ([< Socket.Address.t] as 'a)
     -> ([`Active], 'a) Socket.t
     -> unit Deferred.t

examples/async/async_echo_post.ml

@@ -10,7 +10,10 @@ let main port max_accepts_per_batch () =
   let where_to_listen = Tcp.Where_to_listen.of_port port in
   Tcp.(Server.create_sock ~on_handler_error:`Raise
       ~backlog:10_000 ~max_connections:10_000 ~max_accepts_per_batch where_to_listen)
-    (Server.create_connection_handler ~request_handler ~error_handler)
+    (Server.create_connection_handler
+       ~request_handler
+       ~error_handler
+       ~upgrade_handler:None)
   >>= fun _server ->
     Stdio.printf "Listening on port %i and echoing POST requests.\n" port;
     Stdio.printf "To send a POST request, try one of the following\n\n";

examples/async/async_echo_upgrade.ml

@@ -0,0 +1,47 @@
+open Core
+open Async
+
+open Httpaf_async
+
+let request_handler (_ : Socket.Address.Inet.t) = Httpaf_examples.Server.upgrade
+let error_handler (_ : Socket.Address.Inet.t) = Httpaf_examples.Server.error_handler
+
+let upgrade_handler (_ : Socket.Address.Inet.t) socket =
+  let fd = Socket.fd socket in
+  let reader = Reader.create fd in
+  let writer = Writer.create fd in
+  Reader.read_one_chunk_at_a_time reader ~handle_chunk:(fun bigstring ~pos ~len ->
+    Writer.write_bigstring writer bigstring ~pos ~len;
+    return `Continue)
+  >>| function
+  | `Eof | `Stopped _ | `Eof_with_unconsumed_data _ -> ()
+;;
+
+let main port max_accepts_per_batch () =
+  let where_to_listen = Tcp.Where_to_listen.of_port port in
+  Tcp.(Server.create_sock ~on_handler_error:`Raise
+      ~backlog:10_000 ~max_connections:10_000 ~max_accepts_per_batch where_to_listen)
+    (Server.create_connection_handler
+       ~request_handler
+       ~error_handler
+       ~upgrade_handler:(Some upgrade_handler))
+  >>= fun _server ->
+    Stdio.printf "Listening on port %i, upgrading, and echoing data.\n" port;
+    Stdio.printf "To send an interactive upgrade request, try\n\n";
+    Stdio.printf "  examples/script/upgrade-connect\n%!";
+    Deferred.never ()
+;;
+
+let () =
+  Command.async
+    ~summary:"Echo POST requests"
+    Command.Param.(
+      map (both
+          (flag "-p" (optional_with_default 8080 int)
+            ~doc:"int Source port to listen on")
+          (flag "-a" (optional_with_default 1 int)
+            ~doc:"int Maximum accepts per batch"))
+        ~f:(fun (port, accepts) ->
+              (fun () -> main port accepts ())))
+  |> Command.run
+;;

examples/async/dune

@@ -1,6 +1,6 @@
 (executables
  (libraries httpaf httpaf-async httpaf_examples async core)
- (names     async_echo_post async_get async_post))
+ (names     async_echo_post async_echo_upgrade async_get async_post))
 
 (alias
  (name examples)

examples/lib/httpaf_examples.ml

@@ -86,4 +86,14 @@ module Server = struct
     end;
     Body.Writer.close response_body
   ;;
+
+  let upgrade reqd =
+    if Request.is_upgrade (Reqd.request reqd) then (
+      let headers = Headers.of_list [ "connection", "upgrade" ] in
+      Reqd.respond_with_upgrade reqd headers;
+    ) else (
+      let headers = Headers.of_list [ "connection", "close" ] in
+      Reqd.respond_with_string reqd (Response.create ~headers `Not_found) ""
+    )
+  ;;
 end

examples/lwt/dune

@@ -1,6 +1,6 @@
 (executables
  (libraries httpaf httpaf-lwt-unix httpaf_examples base stdio lwt lwt.unix)
- (names lwt_get lwt_post lwt_echo_post))
+ (names lwt_get lwt_post lwt_echo_post lwt_echo_upgrade))
 
 (alias
  (name examples)

examples/lwt/lwt_echo_post.ml

@@ -12,7 +12,10 @@ let main port =
   Lwt.async (fun () ->
     Lwt_io.establish_server_with_client_socket
       listen_address
-      (Server.create_connection_handler ~request_handler ~error_handler)
+      (Server.create_connection_handler
+         ~request_handler
+         ~error_handler
+         ~upgrade_handler:None)
     >|= fun _server ->
       Stdio.printf "Listening on port %i and echoing POST requests.\n" port;
       Stdio.printf "To send a POST request, try one of the following\n\n";

examples/lwt/lwt_echo_upgrade.ml

@@ -0,0 +1,47 @@
+open Base
+open Lwt.Infix
+module Arg = Caml.Arg
+
+open Httpaf_lwt_unix
+
+let request_handler (_ : Unix.sockaddr) = Httpaf_examples.Server.upgrade
+let error_handler (_ : Unix.sockaddr) = Httpaf_examples.Server.error_handler
+
+let upgrade_handler (_ : Unix.sockaddr) (fd : Lwt_unix.file_descr) =
+  let input = Lwt_io.of_fd fd ~mode:Input in
+  let output = Lwt_io.of_fd fd ~mode:Output in
+  let rec loop () =
+    Lwt_io.read input ~count:4096
+    >>= function
+    | "" -> Lwt.return_unit
+    | data -> Lwt_io.write output data >>= loop
+  in
+  loop ()
+;;
+
+let main port =
+  let listen_address = Unix.(ADDR_INET (inet_addr_loopback, port)) in
+  Lwt.async (fun () ->
+    Lwt_io.establish_server_with_client_socket
+      listen_address
+      (Server.create_connection_handler
+         ~request_handler
+         ~error_handler
+         ~upgrade_handler:(Some upgrade_handler))
+    >|= fun _server ->
+      Stdio.printf "Listening on port %i, upgrading, and echoing data.\n" port;
+      Stdio.printf "To send an interactive upgrade request, try\n\n";
+      Stdio.printf "  examples/script/upgrade-connect\n%!");
+  let forever, _ = Lwt.wait () in
+  Lwt_main.run forever
+;;
+
+let () =
+  let port = ref 8080 in
+  Arg.parse
+    ["-p", Arg.Set_int port, " Listening port number (8080 by default)"]
+    ignore
+    "Echoes POST requests. Runs forever.";
+  main !port
+;;
+

examples/script/upgrade-connect

@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+function headers {
+  printf "\
+GET / HTTP/1.1\r
+Host: localhost\r
+Connection: upgrade\r
+\r
+"
+}
+
+( headers; echo hello; cat; echo bye ) | nc localhost 8080 --close

lib/httpaf.mli

@@ -564,6 +564,10 @@ module Request : sig
       more details. *)
 
   val pp_hum : Format.formatter -> t -> unit [@@ocaml.toplevel_printer]
+
+  val is_upgrade : t -> bool
+  (** [is_upgrade t] returns true if the request has the "Connection: upgrade"
+      header. *)
 end
 
 
@@ -659,6 +663,16 @@ module Reqd : sig
   val respond_with_bigstring : t -> Response.t -> Bigstringaf.t -> unit
   val respond_with_streaming : ?flush_headers_immediately:bool -> t -> Response.t -> Body.Writer.t
 
+  val respond_with_upgrade : ?reason:string -> t -> Headers.t -> unit
+  (** Initiate an HTTP upgrade. [Server_connection.next_write_request] and
+      [next_read_request] will begin returning [`Upgrade] once the response
+      headers have been written, which indicates that the runtime should take
+      over direct control of the socket rather than shuttling bytes through
+      httpaf.
+
+      The headers must indicate a valid upgrade message, e.g. must include
+      "Connection: upgrade". See [Request.is_upgrade]. *)
+
   (** {3 Exception Handling} *)
 
   val report_exn : t -> exn -> unit
@@ -700,7 +714,7 @@ module Server_connection : sig
   (** [create ?config ?error_handler ~request_handler] creates a connection
       handler that will service individual requests with [request_handler]. *)
 
-  val next_read_operation : t -> [ `Read | `Yield | `Close ]
+  val next_read_operation : t -> [ `Read | `Yield | `Close | `Upgrade ]
   (** [next_read_operation t] returns a value describing the next operation
       that the caller should conduct on behalf of the connection. *)
 
@@ -727,6 +741,7 @@ module Server_connection : sig
   val next_write_operation : t -> [
     | `Write of Bigstringaf.t IOVec.t list
     | `Yield
+    | `Upgrade
     | `Close of int ]
   (** [next_write_operation t] returns a value describing the next operation
       that the caller should conduct on behalf of the connection. *)

lib/parse.ml

@@ -245,6 +245,8 @@ module Reader = struct
       | `Fixed 0L  ->
         handler request Body.Reader.empty;
         ok
+      | `Fixed _ | `Chunked when Request.is_upgrade request ->
+        return (Error (`Bad_request request))
       | `Fixed _ | `Chunked as encoding ->
         let request_body = Body.Reader.create Bigstringaf.empty in
         handler request request_body;