<commit_before>open Core.Std
open Lwt
open Cohttp_lwt_unix
open Har_j
open Settings

type protocol = HTTP | HTTPS
exception Too_many_requests
exception Cant_resolve_ip of string * string

let make_server config =
	let nb_current = ref 0 in
	let sock = Network.get_ZMQ_sock config.zmq_host config.zmq_port in
	let global_archive = Option.map ~f:(fun k -> (module Archive.Make (struct let key = k end) : Archive.Sig_make)) config.key in
	let empty_archive = Some (module Archive.Make (struct let key = "" end) : Archive.Sig_make) in

	let send_har archive environment req req_uri res t_client_body t_provider_body client_ip server_ip (t0, har_send, har_wait) =
		let send () = try_lwt (
				t_client_body
			>>= fun (req_length, req_b64) ->
				t_provider_body
			>>= fun (res_length, res_b64) ->
				let module KeyArchive = (val archive : Archive.Sig_make) in

				(* Bypass for XHR *)
				if KeyArchive.key = "" then return () else

				let open Archive in
				let archive_input = {
					environment;
					req;
					req_uri;
					res;
					req_length;
					res_length;
					req_b64;
					res_b64;
					client_ip;
					server_ip;
					timings = (har_send, har_wait, ((get_timestamp_ms ()) - t0 - har_wait));
				} in

				let har_string = KeyArchive.get_alf archive_input |> string_of_alf |> fun x -> "alf_1.0.0 " ^ x in
				Lwt.pick [
					(Lwt_zmq.Socket.send sock har_string
					>>= fun () ->
						if config.debug then
							Lwt_io.printlf "SENT\n%s\n" har_string
						else
							return ()
					);
					(Lwt_unix.sleep Settings.zmq_flush_timeout
					>>= fun () ->
						Lwt_io.printlf "ERROR Could not flush this datapoint to the ZMQ driver within %f seconds:\n%s\n" Settings.zmq_flush_timeout har_string
					);
				]
		) with ex ->
			match ex with
			| e ->
				Lwt_io.printlf "ERROR\n%s\n%s" (Exn.to_string e) (Exn.backtrace ())
		in
		match config.filter_ua with
		| None -> send ()
		| Some filter when Cohttp.Header.get (Request.headers req) "User-Agent" |> Option.value ~default:"" |> Regex.matches filter |> not -> send ()
		| Some _ -> return ()
	in
	let callback (ch, _) req client_body protcol =
		(* Initiate counters and other bookeeping *)
		nb_current := (!nb_current + 1);
		let t0 = Archive.get_timestamp_ms () in
		let client_ip = Network.get_addr_from_ch ch in
		let client_headers = Request.headers req in
		let environment = Option.first_some (Cohttp.Header.get client_headers "Mashape-Environment") config.environment in

		(* This is necessary for now due to https://github.com/mirage/ocaml-cohttp/issues/248 *)
		let uri = Request.uri req
		|> fun uri ->
			match Cohttp.Header.get client_headers "Mashape-Host-Override" with
			| None -> uri
			| Some x -> Uri.with_host uri (Some x)
		in

		(* Prepare the target *)
		let target = (match config.reverse with
		| None -> uri
		| Some (reverse_host, reverse_port) ->
			uri
			|> fun uri -> Uri.with_host uri reverse_host
			|> fun uri ->
				match reverse_port with
				| Some p -> Uri.with_port uri p
				| None -> uri
		) |> fun uri ->
			let protocol_header = Cohttp.Header.get client_headers "Mashape-Upstream-Protocol" |> Option.map ~f:String.lowercase in
			match (protocol_header, protcol) with
			| (Some "https", _) | (None, HTTPS) -> Uri.with_scheme uri (Some "https")
			| (Some "http", _) | (None, HTTP) | (Some _, _) -> Uri.with_scheme uri (Some "http")
		in

		let uri_fixed = Uri.with_scheme uri (Uri.scheme target) in

		(* Debug output *)
		ignore_result (if config.debug then
			Lwt_io.printlf "RECEIVED %s\n> protocol: %s\n> host: %s\n> port: %s\n> path: %s\n"
				(Uri.to_string target)
				(Uri.scheme target |> Option.value ~default:"<<no protocol>>")
				(Uri.host target |> Option.value ~default:"<<no host>>")
				(Uri.port target |> Option.map ~f:Int.to_string |> Option.value ~default:"<<no port>>")
				(Uri.path target)
			else return ());

		(* Start fetching the target IP in advance *)
		let target_to_resolve = target |> Uri.host |> Option.value ~default:"" in
		let t_dns = Network.dns_lookup target_to_resolve in

		(* More bookeeping *)
		let t_client_body = Http_utils.process_body client_body config.replays in
		let local_archive = Option.map (Cohttp.Header.get client_headers "Mashape-Service-Token") ~f:(fun k ->
			(module Archive.Make (struct let key = k end) : Archive.Sig_make)) in
		let har_send = (Archive.get_timestamp_ms ()) - t0 in

		(* Main block. Throws lwt exceptions for any invalid request or error *)
		let response = try_lwt (
			(* Throw some exceptions if needed, then set some headers *)
			if !nb_current > config.concurrent then Lwt.fail Too_many_requests else

			let archive = match Request.meth req with
			| `OPTIONS -> Option.first_some local_archive empty_archive
			| _ -> Option.first_some local_archive global_archive
			in

			match archive with
			| None ->
				Lwt.fail (Failure "Mashape-Service-Token header missing")
			| Some archive ->
				(* So we're doing an upstream request *)
				let client_headers_ready = client_headers
				|> (fun h -> match Cohttp.Header.get h "Mashape-Host-Override" with
					| None -> h
					| Some x -> Cohttp.Header.replace h "Host" x)
				|> fun h -> Cohttp.Header.remove h "Mashape-Service-Token"
				|> fun h -> Cohttp.Header.remove h "Mashape-Host-Override"
				|> fun h -> Cohttp.Header.remove h "Mashape-Environment"
				|> fun h -> Cohttp.Header.remove h "Mashape-Upstream-Protocol"
				|> fun h -> Http_utils.set_x_forwarded_for h client_ip
				in
<commit_msg>Play nicer with other proxies: X-Forwarded-Proto<commit_after>open Core.Std
open Lwt
open Cohttp_lwt_unix
open Har_j
open Settings

type protocol = HTTP | HTTPS
exception Too_many_requests
exception Cant_resolve_ip of string * string

let make_server config =
	let nb_current = ref 0 in
	let sock = Network.get_ZMQ_sock config.zmq_host config.zmq_port in
	let global_archive = Option.map ~f:(fun k -> (module Archive.Make (struct let key = k end) : Archive.Sig_make)) config.key in
	let empty_archive = Some (module Archive.Make (struct let key = "" end) : Archive.Sig_make) in

	let send_har archive environment req req_uri res t_client_body t_provider_body client_ip server_ip (t0, har_send, har_wait) =
		let send () = try_lwt (
				t_client_body
			>>= fun (req_length, req_b64) ->
				t_provider_body
			>>= fun (res_length, res_b64) ->
				let module KeyArchive = (val archive : Archive.Sig_make) in

				(* Bypass for XHR *)
				if KeyArchive.key = "" then return () else

				let open Archive in
				let archive_input = {
					environment;
					req;
					req_uri;
					res;
					req_length;
					res_length;
					req_b64;
					res_b64;
					client_ip;
					server_ip;
					timings = (har_send, har_wait, ((get_timestamp_ms ()) - t0 - har_wait));
				} in

				let har_string = KeyArchive.get_alf archive_input |> string_of_alf |> fun x -> "alf_1.0.0 " ^ x in
				Lwt.pick [
					(Lwt_zmq.Socket.send sock har_string
					>>= fun () ->
						if config.debug then
							Lwt_io.printlf "SENT\n%s\n" har_string
						else
							return ()
					);
					(Lwt_unix.sleep Settings.zmq_flush_timeout
					>>= fun () ->
						Lwt_io.printlf "ERROR Could not flush this datapoint to the ZMQ driver within %f seconds:\n%s\n" Settings.zmq_flush_timeout har_string
					);
				]
		) with ex ->
			match ex with
			| e ->
				Lwt_io.printlf "ERROR\n%s\n%s" (Exn.to_string e) (Exn.backtrace ())
		in
		match config.filter_ua with
		| None -> send ()
		| Some filter when Cohttp.Header.get (Request.headers req) "User-Agent" |> Option.value ~default:"" |> Regex.matches filter |> not -> send ()
		| Some _ -> return ()
	in
	let callback (ch, _) req client_body protocol =
		(* Initiate counters and other bookeeping *)
		nb_current := (!nb_current + 1);
		let t0 = Archive.get_timestamp_ms () in
		let client_ip = Network.get_addr_from_ch ch in
		let client_headers = Request.headers req in
		let environment = Option.first_some (Cohttp.Header.get client_headers "Mashape-Environment") config.environment in

		let uri = Request.uri req
		|> fun uri ->
			match Cohttp.Header.get client_headers "Mashape-Host-Override" with
			| None -> uri
			| Some x -> Uri.with_host uri (Some x)
		in

		(* Prepare the target *)
		let target = (match config.reverse with
		| None -> uri
		| Some (reverse_host, reverse_port) ->
			uri
			|> fun uri -> Uri.with_host uri reverse_host
			|> fun uri ->
				match reverse_port with
				| Some p -> Uri.with_port uri p
				| None -> uri
		) |> fun uri ->
			let protocol_header = (Option.first_some
				(Cohttp.Header.get client_headers "Mashape-Upstream-Protocol")
				(Cohttp.Header.get client_headers "X-Forwarded-Proto"))
			|> Option.map ~f:String.lowercase
			in
			match (protocol_header, protocol) with
			| (Some "https", _) | (None, HTTPS) -> Uri.with_scheme uri (Some "https")
			| (Some "http", _) | (None, HTTP) | (Some _, _) -> Uri.with_scheme uri (Some "http")
		in

		(* This is the URI used to generate the ALF. Same as target, except that it doesn't contain the Host/Target/IP manipulations *)
		let uri_fixed = Uri.with_scheme uri (Uri.scheme target) in

		(* Debug output *)
		ignore_result (if config.debug then
			Lwt_io.printlf "RECEIVED %s\n> protocol: %s\n> host: %s\n> port: %s\n> path: %s\n"
				(Uri.to_string target)
				(Uri.scheme target |> Option.value ~default:"<<no protocol>>")
				(Uri.host target |> Option.value ~default:"<<no host>>")
				(Uri.port target |> Option.map ~f:Int.to_string |> Option.value ~default:"<<no port>>")
				(Uri.path target)
			else return ());

		(* Start fetching the target IP in advance *)
		let target_to_resolve = target |> Uri.host |> Option.value ~default:"" in
		let t_dns = Network.dns_lookup target_to_resolve in

		(* More bookeeping *)
		let t_client_body = Http_utils.process_body client_body config.replays in
		let local_archive = Option.map (Cohttp.Header.get client_headers "Mashape-Service-Token") ~f:(fun k ->
			(module Archive.Make (struct let key = k end) : Archive.Sig_make)) in
		let har_send = (Archive.get_timestamp_ms ()) - t0 in

		(* Main block. Throws lwt exceptions for any invalid request or error *)
		let response = try_lwt (
			(* Throw some exceptions if needed, then set some headers *)
			if !nb_current > config.concurrent then Lwt.fail Too_many_requests else

			let archive = match Request.meth req with
			| `OPTIONS -> Option.first_some local_archive empty_archive
			| _ -> Option.first_some local_archive global_archive
			in

			match archive with
			| None ->
				Lwt.fail (Failure "Mashape-Service-Token header missing")
			| Some archive ->
				(* So we're doing an upstream request *)
				let client_headers_ready = client_headers
				|> (fun h -> match Cohttp.Header.get h "Mashape-Host-Override" with
					| None -> h
					| Some x -> Cohttp.Header.replace h "Host" x)
				|> fun h -> Cohttp.Header.remove h "Mashape-Service-Token"
				|> fun h -> Cohttp.Header.remove h "Mashape-Host-Override"
				|> fun h -> Cohttp.Header.remove h "Mashape-Environment"
				|> fun h -> Cohttp.Header.remove h "Mashape-Upstream-Protocol"
				|> fun h -> Cohttp.Header.remove h "X-Forwarded-Proto"
				|> fun h -> Http_utils.set_x_forwarded_for h client_ip
				in

{'fraction_non_alphanumeric': 0.11641255605381166, 'fraction_numerical': 0.0033183856502242153, 'mean_word_length': 3.1984186746987953, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 159, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '449700', 'n_tokens_mistral': 4440, 'n_tokens_neox': 3734, 'n_words': 1080}