src/ws.bend

bend-kube-runtime@0.2.0.0 · 203 lines · 29 definitions

1import Base
2
3# WebSocket
4# =========
5#
6# RFC 6455 framing, as Kubernetes' exec speaks it (exec.bend). Frames a
7# client sends are masked; frames it reads are decoded a byte at a time,
8# so a frame may arrive split over any number of recvs, and many frames
9# may arrive in one.
10
11type Frame is Data:
12  Frame{fin: Bool, op: U32, data: List<&2, U32>}
13
14# Where the decoder is.
15#   WHead: before a frame's first byte
16#   WLen: before the byte with the mask bit and the 7-bit length
17#   WExt: need more bytes of a 16- or 64-bit length
18#   WKey: need more bytes of the masking key
19#   WPay: left bytes of payload to go, key rotated to the next byte
20#   WBad: a frame the decoder refuses; it stays here
21type WSt is Data:
22  WHead{}
23  WLen{fin: Bool, op: U32}
24  WExt{fin: Bool, op: U32, masked: Bool, need: U32, len: U32}
25  WKey{fin: Bool, op: U32, len: U32, need: U32, key: List<&2, U32>}
26  WPay{fin: Bool, op: U32, left: U32, key: List<&2, U32>, acc: List<&2, U32>}
27  WBad{why: String}
28
29type Stepped is Data:
30  Stepped{st: WSt, got: Maybe<&2, Frame>}
31
32type Fed is Data:
33  Fed{st: WSt, frames: List<&2, Frame>}
34
35# Opcodes.
36def Ws.text() -> U32:
37  1
38
39def Ws.binary() -> U32:
40  2
41
42def Ws.close() -> U32:
43  8
44
45def Ws.ping() -> U32:
46  9
47
48def Ws.pong() -> U32:
49  10
50
51# The largest payload the decoder accepts.
52def Ws.max() -> U32:
53  16777216
54
55# Encoding
56# --------
57
58# The nth byte of x, counting from the low end.
59def Ws.byte(+x: U32, n: Nat) -> U32:
60  U32.mod(U32.shrn(x, Nat.mul(8n, n)), 256)
61
62def Ws.len.go(+n: U32, small: Bool, mid: Bool) -> List<&2, U32>:
63  match small mid:
64    case True{} True{}:
65      [U32.or(128, n)]
66    case True{} False{}:
67      [U32.or(128, n)]
68    case False{} True{}:
69      [254, Ws.byte(n, 1n), Ws.byte(n, 0n)]
70    case False{} False{}:
71      [255, 0, 0, 0, 0, Ws.byte(n, 3n), Ws.byte(n, 2n), Ws.byte(n, 1n), Ws.byte(n, 0n)]
72
73# The header from the second byte on: the mask bit (always set, as a client
74# must) and the length, in 7, 16 or 64 bits.
75def Ws.len(+n: U32) -> List<&2, U32>:
76  Ws.len.go(n, U32.is_lt(n, 126), U32.is_lt(n, 65536))
77
78# XORs bytes with the key a, b, c, d, rotating.
79def Ws.mask(bs: List<&2, U32>, +a: U32, +b: U32, +c: U32, +d: U32) -> List<&2, U32>:
80  match bs:
81    case Nil{}:
82      Nil{}
83    case Con{x, t}:
84      Con{U32.xor(x, a), Ws.mask(t, b, c, d, a)}
85
86# One final frame of op carrying data, masked with key.
87def Ws.frame(+op: U32, +data: List<&2, U32>, +key: U32) -> List<&2, U32>:
88  List.append(&2, U32,
89    Con{U32.or(128, op), Ws.len(U32.from_nat(List.length(&2, U32, data)))},
90    Con{Ws.byte(key, 3n), Con{Ws.byte(key, 2n), Con{Ws.byte(key, 1n), Con{Ws.byte(key, 0n),
91      Ws.mask(data, Ws.byte(key, 3n), Ws.byte(key, 2n), Ws.byte(key, 1n), Ws.byte(key, 0n))}}}})
92
93# Decoding
94# --------
95
96def Ws.unmask(key: List<&2, U32>, +b: U32) -> U32:
97  match key:
98    case Nil{}:
99      b
100    case Con{k, t}:
101      U32.xor(b, k)
102
103def Ws.rotate(key: List<&2, U32>) -> List<&2, U32>:
104  match key:
105    case Nil{}:
106      Nil{}
107    case Con{k, t}:
108      List.append(&2, U32, t, [k])
109
110def Ws.payload(+fin: Bool, +op: U32, +len: U32, key: List<&2, U32>, empty: Bool) -> Stepped:
111  match empty:
112    case True{}:
113      Stepped{WHead{}, Some{Frame{fin, op, Nil{}}}}
114    case False{}:
115      Stepped{WPay{fin, op, len, key, Nil{}}, None{}}
116
117def Ws.begin.go(+fin: Bool, +op: U32, +len: U32, big: Bool, masked: Bool) -> Stepped:
118  match big masked:
119    case True{} m:
120      Stepped{WBad{"frame of " ++ U32.show(len) ++ " bytes is too large"}, None{}}
121    case False{} True{}:
122      Stepped{WKey{fin, op, len, 4, Nil{}}, None{}}
123    case False{} False{}:
124      Ws.payload(fin, op, len, Nil{}, U32.is_zero(len))
125
126# The header is whole: the key follows if masked, then the payload.
127def Ws.begin(+fin: Bool, +op: U32, masked: Bool, +len: U32) -> Stepped:
128  Ws.begin.go(fin, op, len, U32.is_gt(len, Ws.max()), masked)
129
130def Ws.len7(+fin: Bool, +op: U32, masked: Bool, +n: U32, two: Bool, eight: Bool) -> Stepped:
131  match two eight:
132    case True{} e:
133      Stepped{WExt{fin, op, masked, 2, 0}, None{}}
134    case False{} True{}:
135      Stepped{WExt{fin, op, masked, 8, 0}, None{}}
136    case False{} False{}:
137      Ws.begin(fin, op, masked, n)
138
139def Ws.ext(+fin: Bool, +op: U32, masked: Bool, +need: U32, +len: U32, last: Bool) -> Stepped:
140  match last:
141    case True{}:
142      Ws.begin(fin, op, masked, len)
143    case False{}:
144      Stepped{WExt{fin, op, masked, need, len}, None{}}
145
146def Ws.key(+fin: Bool, +op: U32, +len: U32, +need: U32, +key: List<&2, U32>, last: Bool) -> Stepped:
147  match last:
148    case True{}:
149      Ws.payload(fin, op, len, key, U32.is_zero(len))
150    case False{}:
151      Stepped{WKey{fin, op, len, need, key}, None{}}
152
153def Ws.pay(+fin: Bool, +op: U32, +left: U32, key: List<&2, U32>, +acc: List<&2, U32>, last: Bool) ->
154  Stepped:
155  match last:
156    case True{}:
157      Stepped{WHead{}, Some{Frame{fin, op, List.reverse(&2, U32, acc)}}}
158    case False{}:
159      Stepped{WPay{fin, op, left, key, acc}, None{}}
160
161# One byte into the decoder.
162def Ws.step(st: WSt, +b: U32) -> Stepped:
163  match st:
164    case WHead{}:
165      Stepped{WLen{U32.is_ge(b, 128), U32.and(b, 15)}, None{}}
166    case WLen{+fin, +op}:
167      Ws.len7(fin, op, U32.is_ge(b, 128), U32.and(b, 127), U32.is_eq(U32.and(b, 127), 126),
168        U32.is_eq(U32.and(b, 127), 127))
169    case WExt{+fin, +op, masked, +need, +len}:
170      Ws.ext(fin, op, masked, U32.sub(need, 1), U32.add(U32.mul(len, 256), b),
171        U32.is_eq(need, 1))
172    case WKey{+fin, +op, +len, +need, +key}:
173      Ws.key(fin, op, len, U32.sub(need, 1), List.append(&2, U32, key, [b]), U32.is_eq(need, 1))
174    case WPay{+fin, +op, +left, +key, +acc}:
175      Ws.pay(fin, op, U32.sub(left, 1), Ws.rotate(key), Con{Ws.unmask(key, b), acc},
176        U32.is_eq(left, 1))
177    case WBad{why}:
178      Stepped{WBad{why}, None{}}
179
180def Ws.add(got: Maybe<&2, Frame>, out: List<&2, Frame>) -> List<&2, Frame>:
181  match got:
182    case None{}:
183      out
184    case Some{f}:
185      Con{f, out}
186
187def Ws.feed.go(bs: List<&2, U32>, ev: Stepped, out: List<&2, Frame>) -> Fed:
188  match bs ev:
189    case Nil{} Stepped{st, got}:
190      Fed{st, List.reverse(&2, Frame, Ws.add(got, out))}
191    case Con{+b, t} Stepped{st, got}:
192      Ws.feed.go(t, Ws.step(st, b), Ws.add(got, out))
193
194# Feeds bytes to the decoder: its next state and the frames they finished.
195def Ws.feed(st: WSt, bs: List<&2, U32>) -> Fed:
196  Ws.feed.go(bs, Stepped{st, None{}}, Nil{})
197
198def Ws.error(st: WSt) -> Maybe<&2, String>:
199  match st:
200    case WBad{why}:
201      Some{why}
202    case _:
203      None{}