#!/usr/bin/env ocaml

(*
 * Copyright (c) 2014-2015 Anil Madhavapeddy <anil@recoil.org>
 * Copyright (c) 2014 Romain Calascibetta <romain.calascibetta@gmail.com>
 * Copyright (c) 2014 David Sheets <sheets@alum.mit.edu>
 * Copyright (c) 2015-2017 Thomas Gazagnaire <thomas@gazagnaire.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 *)

let () =
  try Topdirs.dir_directory (Sys.getenv "OCAML_TOPLEVEL_PATH")
  with Not_found -> ()
;;

#use "topfind";;
#require "cmdliner";;
#require "cohttp.lwt";;

open Printf

open Lwt
open Cohttp
open Cohttp_lwt_unix

let hum f x = Sexplib.Sexp.to_string_hum (f x)

let handler ~verbose (ch,conn) req body =
  printf "New connection: %s\n%!" (hum Conduit_lwt_unix.sexp_of_flow ch);

  let uri = Cohttp.Request.uri req in
  let uri = Uri.with_scheme uri (Some "https") in
  let uri = Uri.with_host uri (Some "github.com") in
  let uri = Uri.with_port uri None in

  (* Strip out hop-by-hop connection headers *)
  let headers =
    Request.headers req |> fun h ->
    Header.remove h "host" |> fun h ->
    Header.add    h "host" "github.com" |> fun h ->
    Header.remove h "content-length" |> fun h ->
    Header.remove h "transfer-encoding" |> fun h ->
    Header.remove h "connection" |> fun h ->
    Header.add h "accept-encoding" "identity"
  in

  (* Log the request to the console *)
  Cohttp_lwt_body.to_string body >>= fun body ->
  begin if verbose then (
      printf "> %s\n%!" (hum Request.sexp_of_t req);
      printf "> %s\n%!" body;
      Lwt.return_unit
    ) else Lwt.return_unit
  end >>= fun () ->
  let body = Cohttp_lwt_body.of_string body in

  (* Fetch the remote URI *)
  let meth = Request.meth req in
  Client.call ~headers ~body meth uri >>= fun (resp, body) ->
  Cohttp_lwt_body.to_string body >>= fun body ->
  begin if verbose then (
      printf "< %s\n%!" (hum Response.sexp_of_t resp);
      printf "< %s\n%!" body;
      Lwt.return_unit
    ) else Lwt.return_unit
  end >>= fun () ->
  let body = Cohttp_lwt_body.of_string body in

  let status = Response.status resp in
  let headers =
    Response.headers resp |> fun h ->
    Header.remove h "transfer-encoding" |> fun h ->
    Header.remove h "content-length" |> fun h ->
    Header.remove h "connection"
  in
  Server.respond ~headers ~status ~body ()

let start_proxy port host verbose =
  printf "Listening for HTTP request on: %s %d\n%!" host port;
  let conn_closed (ch, conn) =
    printf "Connection %s closed\n%!" (hum Conduit_lwt_unix.sexp_of_flow ch)
  in
  let callback = handler ~verbose in
  let config = Server.make ~callback ~conn_closed () in
  let mode = `TCP (`Port port) in
  Server.create ~mode config

let lwt_start_proxy port host verbose =
  Lwt_main.run (start_proxy port host verbose)

open Cmdliner

let host =
  let doc = "IP address to listen on." in
  Arg.(value & opt string "0.0.0.0" & info ["s"] ~docv:"HOST" ~doc)

let port =
  let doc = "TCP port to listen on." in
  Arg.(value & opt int 8080 & info ["p"] ~docv:"PORT" ~doc)

let verb =
  let doc = "Logging output to console." in
  Arg.(value & flag & info ["v"; "verbose"] ~doc)

let cmd =
  let doc = "a simple http proxy to github.com" in
  let man = [
    `S "DESCRIPTION";
    `P "$(tname) sets up a simple http proxy with lwt as backend";
    `S "BUGS";
    `P "Report them to via e-mail to <mirageos-devel@lists.xenproject.org>, or
        on the issue tracker at <https://github.com/mirage/ocaml-cohttp/issues>";
  ] in
  Term.(pure lwt_start_proxy $ port $ host $ verb),
  Term.info "cohttp-proxy" ~version:"1.0.0" ~doc ~man

let () =
  match Term.eval cmd with
  | `Error _ -> exit 1
  | _ -> exit 0
