Source file helpers.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
(* This file is part of Dream, released under the MIT license. See LICENSE.md
   for details, or visit https://github.com/aantron/dream.

   Copyright 2021 Anton Bachin *)



module Formats = Dream_pure.Formats
module Message = Dream_pure.Message
module Status = Dream_pure.Status
module Stream = Dream_pure.Stream



let client_field =
  Message.new_field
    ~name:"dream.client"
    ~show_value:(fun client -> client)
    ()

(* TODO What should be reported when the client address is missing? This is a
   sign of local testing. *)
let client request =
  match Message.field request client_field with
  | None -> "127.0.0.1:0"
  | Some client -> client

let set_client request client =
  Message.set_field request client_field client



let tls_field =
  Message.new_field
    ~name:"dream.tls"
    ~show_value:string_of_bool
    ()

let tls request =
  match Message.field request tls_field with
  | Some true -> true
  | _ -> false

let set_tls request tls =
  Message.set_field request tls_field tls



let request ~client ~method_ ~target ~tls ~headers server_stream =
  let request =
    Message.request ~method_ ~target ~headers Stream.null server_stream in
  set_client request client;
  set_tls request tls;
  request

let request_with_body ?method_ ?target ?headers body =
  Message.request ?method_ ?target ?headers Stream.null (Stream.string body)



let response_with_body ?status ?code ?headers body =
  let response =
    Message.response ?status ?code ?headers Stream.null Stream.null in
  Message.set_body response body;
  response

let respond ?status ?code ?headers body =
  Lwt.return (response_with_body ?status ?code ?headers body)

let html ?status ?code ?headers body =
  let response = response_with_body ?status ?code ?headers body in
  Message.set_header response "Content-Type" Formats.text_html;
  Lwt.return response

let json ?status ?code ?headers body =
  let response = response_with_body ?status ?code ?headers body in
  Message.set_header response "Content-Type" Formats.application_json;
  Lwt.return response

(* TODO Actually use the request and extract the site prefix. *)
let redirect ?status ?code ?headers _request location =
  let status = (status :> Status.redirection option) in
  let status =
    match status, code with
    | None, None -> Some (`See_Other)
    | _ -> status
  in
  let response = response_with_body ?status ?code ?headers "" in
  Message.set_header response "Location" location;
  Lwt.return response

let stream ?status ?code ?headers ?(close = true) callback =
  let reader, writer = Stream.pipe () in
  let client_stream = Stream.stream reader Stream.no_writer
  and server_stream = Stream.stream Stream.no_reader writer in
  let response =
    Message.response ?status ?code ?headers client_stream server_stream in

  (* TODO Make sure the request id is propagated to the callback. *)
  Lwt.async (fun () ->
    if close then
      match%lwt callback server_stream with
      | () ->
        Message.close server_stream
      | exception exn ->
        let%lwt () = Message.close server_stream in
        raise exn
    else
      callback server_stream);

  Lwt.return response

let empty ?headers status =
  respond ?headers ~status ""

let not_found _ =
  respond ~status:`Not_Found ""



let websocket ?headers ?(close = true) callback =
  let response =
    Message.response
      ~status:`Switching_Protocols ?headers Stream.empty Stream.null in
  let websocket = Message.create_websocket response in

  (* TODO Make sure the request id is propagated to the callback. *)
  Lwt.async (fun () ->
    if close then
      match%lwt callback websocket with
      | () ->
        Message.close_websocket websocket
      | exception exn ->
        let%lwt () = Message.close_websocket websocket ~code:1005 in
        raise exn
    else
      callback websocket);

  Lwt.return response

let receive (_, server_stream) =
  Message.receive server_stream

let receive_fragment (_, server_stream) =
  Message.receive_fragment server_stream

let send ?text_or_binary ?end_of_message (_, server_stream) data =
  Message.send ?text_or_binary ?end_of_message server_stream data