(* Echo server written using state-tracked sockets. *)

#load "libsocketcap"

module EchoServer = struct
  open ASocket

  (* This is a bit different than the version in the paper, because
   * it uses exceptions. *)
  let handleClient['t] (sock: 't socket) (f: string -> string)
                       (cap: 't connected) : unit =
    let rec loop (cap: 't connected): unit =
      let (str, cap) = recv sock 1024 cap in
      let cap        = send sock (f str) cap in
        loop cap
     in try
          loop cap
        with SocketError _ -> ()

  let rec acceptLoop['t] (sock: 't socket) (f: string -> string)
                         (cap: 't listening) : unit =
    let (('s, clientsock, clientcap), cap) = accept sock cap in
      putStrLn "Opened connection";
      (Thread.fork :> (unit -o unit) -> Thread.thread)
        (fun () -> handleClient clientsock f clientcap;
                   putStrLn "Closed connection");
      acceptLoop sock f cap

  let serve (port: int) (f: string -> string) =
    let ('t, sock, cap) = socket () in
    let cap = bind sock port cap in
    let cap = listen sock cap in
      acceptLoop sock f cap
end

let serverFun (s: string) = s

let main (argv: string list) =
  match argv with
  | Cons (port, Nil) -> EchoServer.serve (int_of_string port) serverFun
  | _ -> failwith "Usage: echoServer.aff PORT\n"

in main (getArgs ())