Source file response.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
open Common_

type body =
  [ `String of string | `Stream of IO.Input.t | `Writer of IO.Writer.t | `Void ]

type t = { code: Response_code.t; headers: Headers.t; body: body }

let set_body body self = { self with body }
let set_headers headers self = { self with headers }
let update_headers f self = { self with headers = f self.headers }
let set_header k v self = { self with headers = Headers.set k v self.headers }
let remove_header k self = { self with headers = Headers.remove k self.headers }
let set_code code self = { self with code }

let make_raw ?(headers = []) ~code body : t =
  (* add content length to response *)
  let headers =
    Headers.set "Content-Length" (string_of_int (String.length body)) headers
  in
  { code; headers; body = `String body }

let make_raw_stream ?(headers = []) ~code body : t =
  let headers = Headers.set "Transfer-Encoding" "chunked" headers in
  { code; headers; body = `Stream body }

let make_raw_writer ?(headers = []) ~code body : t =
  let headers = Headers.set "Transfer-Encoding" "chunked" headers in
  { code; headers; body = `Writer body }

let make_void_force_ ?(headers = []) ~code () : t =
  { code; headers; body = `Void }

let make_void ?(headers = []) ~code () : t =
  let is_ok = code < 200 || code = 204 || code = 304 in
  if is_ok then
    make_void_force_ ~headers ~code ()
  else
    make_raw ~headers ~code "" (* invalid to not have a body *)

let make_string ?headers ?(code = 200) r =
  match r with
  | Ok body -> make_raw ?headers ~code body
  | Error (code, msg) -> make_raw ?headers ~code msg

let make_stream ?headers ?(code = 200) r =
  match r with
  | Ok body -> make_raw_stream ?headers ~code body
  | Error (code, msg) -> make_raw ?headers ~code msg

let make_writer ?headers ?(code = 200) r : t =
  match r with
  | Ok body -> make_raw_writer ?headers ~code body
  | Error (code, msg) -> make_raw ?headers ~code msg

let make ?headers ?(code = 200) r : t =
  match r with
  | Ok (`String body) -> make_raw ?headers ~code body
  | Ok (`Stream body) -> make_raw_stream ?headers ~code body
  | Ok `Void -> make_void ?headers ~code ()
  | Ok (`Writer f) -> make_raw_writer ?headers ~code f
  | Error (code, msg) -> make_raw ?headers ~code msg

let fail ?headers ~code fmt =
  Printf.ksprintf (fun msg -> make_raw ?headers ~code msg) fmt

exception Bad_req = Bad_req

let fail_raise ~code fmt =
  Printf.ksprintf (fun msg -> raise (Bad_req (code, msg))) fmt

let pp out self : unit =
  let pp_body out = function
    | `String s -> Format.fprintf out "%S" s
    | `Stream _ -> Format.pp_print_string out "<stream>"
    | `Writer _ -> Format.pp_print_string out "<writer>"
    | `Void -> ()
  in
  Format.fprintf out "{@[code=%d;@ headers=[@[%a@]];@ body=%a@]}" self.code
    Headers.pp self.headers pp_body self.body

let output_ ~bytes (oc : IO.Output.t) (self : t) : unit =
  (* double indirection:
     - print into [buffer] using [bprintf]
     - transfer to [buf_] so we can output from there *)
  let tmp_buffer = Buffer.create 32 in
  let buf = Buf.of_bytes bytes in

  (* write start of reply *)
  Printf.bprintf tmp_buffer "HTTP/1.1 %d %s\r\n" self.code
    (Response_code.descr self.code);
  Buf.add_buffer buf tmp_buffer;
  Buffer.clear tmp_buffer;

  let body, is_chunked =
    match self.body with
    | `String s when String.length s > 1024 * 500 ->
      (* chunk-encode large bodies *)
      `Writer (IO.Writer.of_string s), true
    | `String _ as b -> b, false
    | `Stream _ as b -> b, true
    | `Writer _ as b -> b, true
    | `Void as b -> b, false
  in
  let headers =
    if is_chunked then
      self.headers
      |> Headers.set "transfer-encoding" "chunked"
      |> Headers.remove "content-length"
    else
      self.headers
  in
  let self = { self with headers; body } in
  Log.debug (fun k ->
      k "t[%d]: output response: %s"
        (Thread.id @@ Thread.self ())
        (Format.asprintf "%a" pp { self with body = `String "<...>" }));

  (* write headers, using [buf] to batch writes *)
  List.iter
    (fun (k, v) ->
      Printf.bprintf tmp_buffer "%s: %s\r\n" k v;
      Buf.add_buffer buf tmp_buffer;
      Buffer.clear tmp_buffer)
    headers;

  IO.Output.output_buf oc buf;
  IO.Output.output_string oc "\r\n";
  Buf.clear buf;

  (match body with
  | `String "" | `Void -> ()
  | `String s -> IO.Output.output_string oc s
  | `Writer w ->
    (* use buffer to chunk encode [w] *)
    let oc' = IO.Output.chunk_encoding ~buf ~close_rec:false oc in
    (try
       IO.Writer.write oc' w;
       IO.Output.close oc'
     with e ->
       let bt = Printexc.get_raw_backtrace () in
       IO.Output.close oc';
       IO.Output.flush oc;
       Printexc.raise_with_backtrace e bt)
  | `Stream str ->
    (match IO.Input.output_chunked' ~buf oc str with
    | () ->
      Log.debug (fun k ->
          k "t[%d]: done outputing stream" (Thread.id @@ Thread.self ()));
      IO.Input.close str
    | exception e ->
      let bt = Printexc.get_raw_backtrace () in
      Log.error (fun k ->
          k "t[%d]: outputing stream failed with %s"
            (Thread.id @@ Thread.self ())
            (Printexc.to_string e));
      IO.Input.close str;
      IO.Output.flush oc;
      Printexc.raise_with_backtrace e bt));
  IO.Output.flush oc

module Private_ = struct
  let make_void_force_ = make_void_force_
  let output_ = output_
end