open Cohttp
module String_io = Cohttp.Private.String_io
module StringRequest = Request.Private.Make (String_io.M)

let user_agent = Cohttp.Header.user_agent
let uri_userinfo = Uri.of_string "http://foo:bar%2525@ocaml.org"

let header_auth =
  let h = Header.init () in
  let h = Header.add_authorization h (`Basic ("qux", "qwerty")) in
  h

let is_some = function None -> false | Some _ -> true

let header_has_auth _ =
  Alcotest.check Alcotest.bool "Test header has auth"
    (header_auth |> Header.get_authorization |> is_some)
    true

let uri_has_userinfo _ =
  Alcotest.check Alcotest.bool "Uri has user info"
    (uri_userinfo |> Uri.userinfo |> is_some)
    true

let t_credentials =
  Alcotest.testable
    (fun fmt c ->
      let sexp = Cohttp.Auth.sexp_of_credential c in
      Sexplib0.Sexp.pp_hum fmt sexp)
    ( = )

let auth_uri_no_override _ =
  let r = Request.make ~headers:header_auth uri_userinfo in
  Alcotest.check
    (Alcotest.option t_credentials)
    "auth uri no override"
    (r |> Request.headers |> Header.get_authorization)
    (Header.get_authorization header_auth)

let auth_uri _ =
  let r = Request.make uri_userinfo in
  Alcotest.check
    (Alcotest.option t_credentials)
    "auth_uri"
    (r |> Request.headers |> Header.get_authorization)
    (Some (`Basic ("foo", "bar%25")))

let t_encoding =
  Alcotest.testable
    (fun fmt e ->
      let sexp = Cohttp.Transfer.sexp_of_encoding e in
      Sexplib0.Sexp.pp fmt sexp)
    ( = )

let encoding_content_length_header () =
  let r =
    Request.make
      ~headers:(Cohttp.Header.of_list [ ("content-length", "100") ])
      (Uri.of_string "http://someuri.com")
  in
  Alcotest.check t_encoding "body encoding determined by content-length header"
    (r |> Request.encoding) (Fixed 100L)

let encoding_transfer_encoding_header () =
  let r =
    Request.make
      ~headers:(Cohttp.Header.of_list [ ("transfer-encoding", "chunked") ])
      (Uri.of_string "http://someuri.com")
  in
  Alcotest.check t_encoding
    "body encoding determined by transfer-encoding header"
    (r |> Request.encoding) Chunked

let encoding_both_headers () =
  let r =
    Request.make
      ~headers:
        (Cohttp.Header.of_list
           [ ("transfer-encoding", "chunked"); ("content-length", "100") ])
      (Uri.of_string "http://someuri.com")
  in
  Alcotest.check t_encoding
    "body encoding with content-length and transfer-encoding headers."
    (r |> Request.encoding) Chunked

let encoding_header_opt_argument () =
  let r =
    Request.make ~encoding:Chunked
      ~headers:(Cohttp.Header.of_list [ ("content-length", "100") ])
      (Uri.of_string "http://someuri.com")
  in
  Alcotest.check t_encoding
    "body encoding with content-length and transfer-encoding headers."
    (r |> Request.encoding) (Fixed 100L)

module Parse_result = struct
  type 'a t = [ `Ok of 'a | `Invalid of string | `Eof ]

  let map t ~f =
    match t with `Ok x -> `Ok (f x) | (`Invalid _ | `Eof) as e -> e
end

let uri_testable : Uri.t Alcotest.testable =
  Alcotest.testable Uri.pp_hum Uri.equal

let t_parse_result_uri : Uri.t Parse_result.t Alcotest.testable =
  Alcotest.testable
    (fun fmt -> function
      | `Invalid s -> Format.fprintf fmt "`Invalid %s" s
      | `Eof -> Format.fprintf fmt "`Eof"
      | `Ok u -> Uri.pp_hum fmt u)
    (fun x y ->
      match (x, y) with `Ok x, `Ok y -> Uri.equal x y | x, y -> x = y)

let parse_request_uri_ r (expected : Uri.t Parse_result.t) name =
  String_io.M.(
    StringRequest.read (String_io.open_in r)
    >>= fun (result : Http.Request.t Parse_result.t) ->
    let uri = Parse_result.map result ~f:Request.uri in
    return @@ Alcotest.check t_parse_result_uri name uri expected)

let bad_request = `Invalid "bad request URI"

let parse_request_uri _ =
  let r = "GET / HTTP/1.1\r\n\r\n" in
  let uri = `Ok (Uri.of_string "/") in
  parse_request_uri_ r uri "parse_request_uri"

let parse_request_uri_host _ =
  let r = "GET / HTTP/1.1\r\nHost: example.com\r\n\r\n" in
  let uri = `Ok (Uri.of_string "//example.com/") in
  parse_request_uri_ r uri "parse_request_uri_host"

let parse_request_uri_host_port _ =
  let r = "GET / HTTP/1.1\r\nHost: example.com:8080\r\n\r\n" in
  let uri = `Ok (Uri.of_string "//example.com:8080/") in
  parse_request_uri_ r uri "parse_request_uri_host_port"

let parse_request_uri_double_slash _ =
  let r = "GET // HTTP/1.1\r\n\r\n" in
  let uri = `Ok (Uri.with_path (Uri.of_string "") "//") in
  parse_request_uri_ r uri "parse_request_uri_double_slash"

let parse_request_uri_host_double_slash _ =
  let r = "GET // HTTP/1.1\r\nHost: example.com\r\n\r\n" in
  let uri = `Ok (Uri.of_string "//example.com//") in
  parse_request_uri_ r uri "parse_request_uri_host_double_slash"

let parse_request_uri_triple_slash _ =
  let r = "GET /// HTTP/1.1\r\n\r\n" in
  let uri = `Ok (Uri.with_path (Uri.of_string "") "///") in
  parse_request_uri_ r uri "parse_request_uri_triple_slash"

let parse_request_uri_host_triple_slash _ =
  let r = "GET /// HTTP/1.1\r\nHost: example.com\r\n\r\n" in
  let uri = `Ok (Uri.of_string "//example.com///") in
  parse_request_uri_ r uri "parse_request_uri_host_triple_slash"

let parse_request_uri_no_slash _ =
  let r = "GET foo HTTP/1.1\r\n\r\n" in
  parse_request_uri_ r bad_request "parse_request_uri_no_slash"

let parse_request_uri_host_no_slash _ =
  let r = "GET foo HTTP/1.1\r\nHost: example.com\r\n\r\n" in
  parse_request_uri_ r bad_request "parse_request_uri_host_no_slash"

let parse_request_uri_empty _ =
  let r = "GET  HTTP/1.1\r\n\r\n" in
  let uri = `Ok (Uri.of_string "/") in
  parse_request_uri_ r uri "parse_request_uri_empty"

let parse_request_uri_host_empty _ =
  let r = "GET  HTTP/1.1\r\nHost: example.com\r\n\r\n" in
  let uri = `Ok (Uri.of_string "//example.com/") in
  parse_request_uri_ r uri "parse_request_uri_host_empty"

let parse_request_uri_path_like_scheme _ =
  let r = "GET http://example.net HTTP/1.1\r\n\r\n" in
  let uri = `Ok (Uri.of_string "http://example.net/") in
  parse_request_uri_ r uri "parse_request_uri_path_like_scheme"

let parse_request_uri_host_path_like_scheme _ =
  let r = "GET http://example.net HTTP/1.1\r\nHost: example.com\r\n\r\n" in
  let uri = `Ok (Uri.of_string "http://example.net/") in
  parse_request_uri_ r uri "parse_request_uri_host_path_like_scheme"

let parse_request_uri_path_like_host_port _ =
  let path = "//example.net:8080" in
  let r = "GET " ^ path ^ " HTTP/1.1\r\n\r\n" in
  let uri = `Ok (Uri.with_path (Uri.of_string "") path) in
  parse_request_uri_ r uri "parse_request_uri_path_like_host_port"

let parse_request_uri_host_path_like_host_port _ =
  let path = "//example.net:8080" in
  let r = "GET " ^ path ^ " HTTP/1.1\r\nHost: example.com\r\n\r\n" in
  let uri = `Ok (Uri.with_path (Uri.of_string "//example.com") path) in
  parse_request_uri_ r uri "parse_request_uri_host_path_like_host_port"

let parse_request_uri_query _ =
  let pqs = "/?foo" in
  let r = "GET " ^ pqs ^ " HTTP/1.1\r\n\r\n" in
  let uri = `Ok (Uri.of_string pqs) in
  parse_request_uri_ r uri "parse_request_uri_query"

let parse_request_uri_host_query _ =
  let pqs = "/?foo" in
  let r = "GET " ^ pqs ^ " HTTP/1.1\r\nHost: example.com\r\n\r\n" in
  let uri = `Ok (Uri.of_string ("//example.com" ^ pqs)) in
  parse_request_uri_ r uri "parse_request_uri_host_query"

let parse_request_uri_query_no_slash _ =
  let r = "GET ?foo HTTP/1.1\r\n\r\n" in
  parse_request_uri_ r bad_request "parse_request_uri_query_no_slash"

let parse_request_uri_host_query_no_slash _ =
  let r = "GET ?foo HTTP/1.1\r\nHost: example.com\r\n\r\n" in
  parse_request_uri_ r bad_request "parse_request_uri_host_query_no_slash"

let parse_request_connect _ =
  let r = "CONNECT vpn.example.net:443 HTTP/1.1\r\n\r\n" in
  let uri = `Ok (Uri.of_string "//vpn.example.net:443") in
  parse_request_uri_ r uri "parse_request_connect"

let parse_request_connect_host _ =
  let r =
    "CONNECT vpn.example.net:443 HTTP/1.1\r\nHost: vpn.example.com:443\r\n\r\n"
  in
  let uri = `Ok (Uri.of_string "//vpn.example.net:443") in
  parse_request_uri_ r uri "parse_request_connect_host"

let parse_request_options _ =
  let r = "OPTIONS * HTTP/1.1\r\n\r\n" in
  let uri = `Ok (Uri.of_string "") in
  parse_request_uri_ r uri "parse_request_options"

let parse_request_options_host _ =
  let r = "OPTIONS * HTTP/1.1\r\nHost: example.com:443\r\n\r\n" in
  let uri = `Ok (Uri.of_string "//example.com:443") in
  parse_request_uri_ r uri "parse_request_options_host"

let parse_request_uri_traversal _ =
  let r = "GET /../../../../etc/shadow HTTP/1.1\r\n\r\n" in
  let uri = `Ok (Uri.of_string "/etc/shadow") in
  parse_request_uri_ r uri "parse_request_uri_traversal"

let parse_request_uri_host_traversal _ =
  let r = "GET /../../../../etc/shadow HTTP/1.1\r\nHost: example.com\r\n\r\n" in
  let uri = `Ok (Uri.of_string "//example.com/etc/shadow") in
  parse_request_uri_ r uri "parse_request_uri_host_traversal"

let uri_round_trip _ =
  let expected_uri =
    let uri = Uri.of_string "https://www.example.com/test" in
    Uri.with_userinfo uri (Some "foo")
  in
  let actual_uri =
    let uri = Request.make expected_uri |> Request.uri in
    (* These are the fields that aren't preserved: *)
    let uri = Uri.with_scheme uri (Uri.scheme expected_uri) in
    Uri.with_userinfo uri (Uri.userinfo expected_uri)
  in
  Alcotest.check uri_testable "Request.make uri round-trip" actual_uri
    expected_uri

let () = Printexc.record_backtrace true

module Buffer = struct
  include Http_bytebuffer.Bytebuffer

  include
    Http_bytebuffer.Bytebuffer.Make
      (struct
        type 'a t = 'a

        let ( >>= ) v f = f v
        let ( >>| ) v f = f v
        let return v = v
      end)
      (struct
        type src = string

        let refill s buffer ~pos ~len =
          if String.equal s "" then `Eof
          else
            let len = min len (String.length s) in
            let () = Bytes.blit_string s 0 buffer pos len in
            `Ok len
      end)
end

module Test_io = struct
  type 'a t = 'a

  let ( >>= ) v f = f v
  let return v = v

  type ic = Buffer.t
  type oc = Buffer.t
  type conn = unit

  let refill _ = `Eof

  let with_input_buffer b ~f =
    let contents = Buffer.to_string b in
    let res, read = f contents ~pos:0 ~len:(String.length contents) in
    let () = Buffer.drop b read in
    res

  let read_line buffer = Buffer.read_line buffer ""
  let read buffer = Buffer.read buffer ""
  let write buffer string = Buffer.refill buffer string |> ignore
  let flush _ = ()
end

module Request = Request.Private.Make (Test_io)

let null_content_length_header () =
  let output = Buffer.create 1024 in
  let () =
    (* The user-agent in releases contentsontains the version, we need to strip
       it for the test *)
    let r =
      Cohttp.Request.make_for_client ~chunked:false ~body_length:0L `PUT
        (Uri.of_string "http://someuri.com")
    in
    Request.write_header r output
  in
  let expected =
    "PUT / HTTP/1.1\r\nhost: someuri.com\r\nuser-agent: "
    ^ user_agent
    ^ "\r\ncontent-length: 0\r\n\r\n"
  in
  Alcotest.(check string)
    "null content-length header are sent" expected (Buffer.to_string output)

let useless_null_content_length_header () =
  let output = Buffer.create 1024 in
  let () =
    let r =
      Cohttp.Request.make_for_client `GET (Uri.of_string "http://someuri.com")
    in
    Request.write_header r output
  in
  let expected =
    "GET / HTTP/1.1\r\nhost: someuri.com\r\nuser-agent: "
    ^ user_agent
    ^ "\r\n\r\n"
  in
  Alcotest.(check string)
    "null content-length header are not sent for bodyless methods" expected
    (Buffer.to_string output)

let () =
  Alcotest.run "test_request"
    [
      ( "Auth",
        [
          ("header has auth", `Quick, header_has_auth);
          ("URI has user info", `Quick, uri_has_userinfo);
          ("from URI - do not override", `Quick, auth_uri_no_override);
          ("from URI", `Quick, auth_uri);
        ] );
      ( "Encoding",
        [
          ("from content-length header", `Quick, encoding_content_length_header);
          ( "from transfer-encoding header",
            `Quick,
            encoding_transfer_encoding_header );
          ("with both headers", `Quick, encoding_both_headers);
          ( "from both optional argument and headers",
            `Quick,
            encoding_header_opt_argument );
          ("null content-length", `Quick, null_content_length_header);
          ( "useless null content-length",
            `Quick,
            useless_null_content_length_header );
        ] );
      ( "Parse URI",
        [
          ("simple", `Quick, parse_request_uri);
          ("with host", `Quick, parse_request_uri_host);
          ("with host and port", `Quick, parse_request_uri_host_port);
          ("double slash", `Quick, parse_request_uri_double_slash);
          ("double slash with host", `Quick, parse_request_uri_host_double_slash);
          ("triple slash", `Quick, parse_request_uri_triple_slash);
          ("triple slash with host", `Quick, parse_request_uri_host_triple_slash);
          ("no slash", `Quick, parse_request_uri_no_slash);
          ("no slash with host", `Quick, parse_request_uri_host_no_slash);
          ("empty", `Quick, parse_request_uri_empty);
          ("empty with host", `Quick, parse_request_uri_host_empty);
          ("path like scheme", `Quick, parse_request_uri_path_like_scheme);
          ( "path like scheme with host",
            `Quick,
            parse_request_uri_host_path_like_scheme );
          ("path like host:port", `Quick, parse_request_uri_path_like_host_port);
          ( "path like host:port with host",
            `Quick,
            parse_request_uri_host_path_like_host_port );
          ("with query string", `Quick, parse_request_uri_query);
          ("with query with host", `Quick, parse_request_uri_host_query);
          ( "no slash with query string",
            `Quick,
            parse_request_uri_query_no_slash );
          ( "no slash with query with host",
            `Quick,
            parse_request_uri_host_query_no_slash );
          ("CONNECT", `Quick, parse_request_connect);
          ("CONNECT with host", `Quick, parse_request_connect_host);
          ("OPTIONS", `Quick, parse_request_options);
          ("OPTIONS with host", `Quick, parse_request_options_host);
          ("parent traversal", `Quick, parse_request_uri_traversal);
          ( "parent traversal with host",
            `Quick,
            parse_request_uri_host_traversal );
          ("uri round-trip", `Quick, uri_round_trip);
        ] );
    ]