~pepe/trevor

trevor/trevor/init.janet -rw-r--r-- 5.0 KiB
7389a261Josef Pospíšil Give supervisor close a day ago
                                                                                
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
165
166
167
168
169
170
171
# Parts of this code are from halo2.
# Thank you very much Sean!

(use chidi)
(use janetls)
(use marble)

(defn- key [buf]
  (-?>> buf
        (peg/match '(* (thru "Sec-WebSocket-Key") ": " '(to "\r\n")))
        first))

(def- magic-string "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")

(defn- handshake [buf]
  (def key (key buf))
  (switching-protocols
    (->>
      (string key magic-string)
      (md/digest :sha1)
      hex/decode
      string
      base64/encode)))

(defn- unmask [mask masked-data]
  (def data @"")
  (loop [[i byte] :pairs masked-data]
    (buffer/push data (bxor byte (mask (% i 4)))))
  data)

(defn- str->num [s nb]
  ((peg/match ~(uint ,nb) (string/reverse s)) 0))

(defn response [opc message]
  ```
  Returns the string of websocket response with opcode `opc`
  and `message`.
  ```
  (var dl (length message))
  (if (> dl 125)
    (set dl (string/from-bytes 2r01111110 (brshift dl 8) dl)))
  (buffer/push @"" (bor 2r10000000 opc) dl message))

(defn text [message]
  ```
  Returns a text message string.
  ```
  (response 0x1 message))

(defn binary [message]
  ```
  Returns a binary message string.
  ```
  (response 0x2 message))

(defn default-supervisor
  ```
  Default supervisor which is used when you do not supply your own.
  It expects channel where to take the events and handling table
  for handling new connections.
  ```
  [chan handling]
  (var last-connection nil)
  (forever
    (match (ev/take chan)
      [:close connection] (:close connection)
      [:error fiber]
      (let [err (fiber/last-value fiber)
            conn ((fiber/getenv fiber) :conn)]
        (unless (one-of err
                        "Connection reset by peer"
                        "stream is closed")
          (eprint err)
          (debug/stacktrace fiber)
          (protect (:write conn (text err)))
          (:close conn)))
      [:conn connection]
      (ev/go
        (fiber/new
          (fn handling-connection []
            (setdyn :conn connection)
            (handling connection)) :tp) nil chan))))

(defn- make-socket [connection handler]
  (def Socket
    @{:write
      (fn write [self msg]
        (match
          (protect
            (:write connection msg))
          [true _] (protect (:flush connection))
          [false err]
          (do
            (ev/give-supervisor :close connection)
            (:closed self))))
      :check
      (fn check [&] true)
      :close
      (fn close [self msg]
        (:write self (response 0x8 msg))
        (ev/give-supervisor :close connection))})
  (make Socket handler))

(defn on-connection
  ```
  A handler for the websockets connection.
  It is compatible with the chidi and can be used
  with its `server/start` function as `on-connection-fn`
  argument.
  Its only argument is `handler` table, which should have
  at least `:connect`, `:read` and `:closed` methods.
  Methods are called when client connects, sends a message
  and closes the connection respectively.
  Optionaly you can provide `:check` method, which is called
  with the initial request and connection, and returns true
  if we can proceed with the connection, or false to decline it.
  You can also write to the connection, as it can be used for
  authentication for example.
  Handler is enriched with `:write` method for writting to
  the connection and can be used for this outside of the Trevor.
  Handler is enriched with `:close` method for closing
  the connection.
  ```
  [handler]

  (def buff-size 1024)
  (assert (and (table? handler)
               (handler :connect)
               (handler :read)
               (handler :closed))
          "Handler is not valid")

  (fn on-connection [connection]
    (def handling (make-socket connection (merge handler)))
    (def req (:read connection buff-size))
    (when (and req
               (:check handling req connection)
               ((protect (:write handling (handshake req))) 0))
      (:flush connection)
      (:connect handling req)
      (def msg (buffer/new 125))
      (forever
        (def cbytes (:read connection 2))
        (unless cbytes
          (ev/give-supervisor :close connection)
          (break))
        (def [fb sb] cbytes)
        (def fin (band 2r10000000 fb))
        (def opc (band 2r00001111 fb))
        (def msk? (band 2r10000000 sb))
        (when (zero? msk?)
          (ev/give-supervisor :close connection)
          (break))
        (var pln (band 2r01111111 sb))
        (case pln
          126 (set pln (str->num (:read connection 2) 2))
          127 (set pln (str->num (:read connection 8) 8)))
        (if (pos? pln)
          (buffer/push msg (unmask (:read connection 4)
                                   (:read connection pln))))
        (case opc
          0x9 (:write handling (response 0xA msg))
          0x8 (do
                (:write connection (response 0x8 msg))
                (ev/give-supervisor :close connection)
                (break))
          (when (pos? fin)
            (:read handling opc msg)
            (buffer/clear msg)))))
    (:closed handling)
    (ev/give-supervisor :close connection)))