1	# DNS A-record lookup over UDP. Source: https://github.com/paymog/bend-kit/tree/main/dns
2	import Base
3	# bend-kit-wire@0.4.0.2
4	import 0xb481bb64f1136053b0a6b2c6da6277ab/wire.bend as Wire
5	
6	# DNS A-record lookup (RFC 1035) over UDP. Messages are byte strings.
7	#   import bend-kit-dns@0.3.1.1/dns.bend as Dns
8	
9	type Cur is Data:
10	  Cur{v: U32, rest: String}
11	
12	type RR is Data:
13	  RRBad{}
14	  RRA{ip: String}
15	  RRSkip{rest: String}
16	
17	def u8(s: String) -> Maybe<&2, Cur>:
18	  match s:
19	    case SNil{}:
20	      None{}
21	    case SCon{Chr{c}, t}:
22	      Some{Cur{c, t}}
23	
24	def u16.lo(hi: U32, m: Maybe<&2, Cur>) -> Maybe<&2, Cur>:
25	  match m:
26	    case None{}:
27	      None{}
28	    case Some{Cur{lo, r}}:
29	      Some{Cur{(hi * 256 + lo : U32), r}}
30	
31	def u16.go(m: Maybe<&2, Cur>) -> Maybe<&2, Cur>:
32	  match m:
33	    case None{}:
34	      None{}
35	    case Some{Cur{hi, r}}:
36	      u16.lo(hi, u8(r))
37	
38	def u16(s: String) -> Maybe<&2, Cur>:
39	  u16.go(u8(s))
40	
41	def skip.if(+s: String, +n: Nat, short: Bool) -> Maybe<&2, String>:
42	  match short:
43	    case True{}:
44	      None{}
45	    case False{}:
46	      Some{String.drop(s, n)}
47	
48	def skip(+s: String, +n: Nat) -> Maybe<&2, String>:
49	  skip.if(s, n, Nat.is_lt(String.length(s), n))
50	
51	def then(m: Maybe<&2, String>, n: Nat) -> Maybe<&2, String>:
52	  match m:
53	    case None{}:
54	      None{}
55	    case Some{r}:
56	      skip(r, n)
57	
58	def push16(+n: U32, acc: String) -> String:
59	  SCon{Chr{U32.and(n, 255)}, SCon{Chr{U32.shrn(n, 8n)}, acc}}
60	
61	# Query
62	# -----
63	
64	def label.ok(+l: String) -> Bool:
65	  +n = String.length(l)
66	  Bool.and(Nat.is_lt(0n, n), Nat.is_le(n, 63n))
67	
68	def label.ascii(s: String) -> Bool:
69	  match s:
70	    case SNil{}:
71	      True{}
72	    case SCon{Chr{c}, t}:
73	      Bool.and(U32.is_lt(c, 128), label.ascii(t))
74	
75	def labels.ok(xs: List<&2, String>) -> Bool:
76	  match xs:
77	    case Nil{}:
78	      True{}
79	    case Con{+l, t}:
80	      Bool.and(Bool.and(label.ok(l), label.ascii(l)), labels.ok(t))
81	
82	def labels.put(xs: List<&2, String>, acc: String) -> String:
83	  match xs:
84	    case Nil{}:
85	      acc
86	    case Con{+l, t}:
87	      labels.put(t, String.reverse(l) ++ SCon{Chr{U32.from_nat(String.length(l))}, acc})
88	
89	def query.if(+id: U32, xs: List<&2, String>, ok: Bool) -> Maybe<&2, String>:
90	  match ok:
91	    case False{}:
92	      None{}
93	    case True{}:
94	      # header: id, RD, QDCOUNT 1; question: QNAME, QTYPE A, QCLASS IN
95	      Some{String.reverse(push16(1, push16(1, SCon{Chr{0}, labels.put(xs, push16(0, push16(0, push16(0, push16(1, push16(256, push16(id, SNil{})))))))})))}
96	
97	def query.labels(+id: U32, +xs: List<&2, String>) -> Maybe<&2, String>:
98	  query.if(id, xs, labels.ok(xs))
99	
100	def strip_dot.if(+s: String, dot: Bool) -> String:
101	  match dot:
102	    case True{}:
103	      String.reverse(String.drop(String.reverse(s), 1n))
104	    case False{}:
105	      s
106	
107	def strip_dot(+s: String) -> String:
108	  strip_dot.if(s, String.ends_with(s, "."))
109	
110	# RFC 1035 §4.1: a standard query for the A record of name.
111	def query(id: U32, name: String) -> Maybe<&2, String>:
112	  query.labels(id, String.split(strip_dot(name), '.'))
113	
114	# Answer
115	# ------
116	
117	# §4.1.4: a name is labels ending in a zero octet, or a two-octet pointer.
118	def name.skip(fuel: Nat, s: String) -> Maybe<&2, String>:
119	  match fuel:
120	    case 0n:
121	      None{}
122	    case 1n+f:
123	      match s:
124	        case SNil{}:
125	          None{}
126	        case SCon{Chr{+l}, +t}:
127	          Bool.pick(Maybe<&2, String>, U32.is_eq(l, 0), Some{t}, Bool.pick(Maybe<&2, String>, U32.is_le(192, l), skip(t, 1n), Bool.pick(Maybe<&2, String>, U32.is_le(64, l), None{}, name.skip(f, String.drop(t, U32.to_nat(l))))))
128	
129	def name.then(m: Maybe<&2, String>, n: Nat) -> Maybe<&2, String>:
130	  match m:
131	    case None{}:
132	      None{}
133	    case Some{+r}:
134	      then(name.skip(String.length(r), r), n)
135	
136	def dotted(s: String) -> String:
137	  match s:
138	    case SNil{}:
139	      ""
140	    case SCon{Chr{c}, t}:
141	      match t:
142	        case SNil{}:
143	          U32.show(c)
144	        case SCon{_, _}:
145	          U32.show(c) ++ "." ++ dotted(t)
146	
147	def rr.f(+r: String, +len: U32, a: Bool) -> RR:
148	  match a:
149	    case True{}:
150	      RRA{dotted(String.take(r, 4n))}
151	    case False{}:
152	      RRSkip{String.drop(r, U32.to_nat(len))}
153	
154	def rr.data(+r: String, +len: U32, a: Bool, short: Bool) -> RR:
155	  match short:
156	    case True{}:
157	      RRBad{}
158	    case False{}:
159	      rr.f(r, len, a)
160	
161	def rr.e(+ty: U32, +cl: U32, m: Maybe<&2, Cur>) -> RR:
162	  match m:
163	    case None{}:
164	      RRBad{}
165	    case Some{Cur{+len, +r}}:
166	      rr.data(r, len, Bool.and(Bool.and(U32.is_eq(ty, 1), U32.is_eq(cl, 1)), U32.is_eq(len, 4)), Nat.is_lt(String.length(r), U32.to_nat(len)))
167	
168	def rr.d(ty: U32, cl: U32, m: Maybe<&2, String>) -> RR:
169	  match m:
170	    case None{}:
171	      RRBad{}
172	    case Some{r}:
173	      rr.e(ty, cl, u16(r))
174	
175	def rr.c(ty: U32, m: Maybe<&2, Cur>) -> RR:
176	  match m:
177	    case None{}:
178	      RRBad{}
179	    case Some{Cur{cl, r}}:
180	      rr.d(ty, cl, skip(r, 4n))
181	
182	def rr.b(m: Maybe<&2, Cur>) -> RR:
183	  match m:
184	    case None{}:
185	      RRBad{}
186	    case Some{Cur{ty, r}}:
187	      rr.c(ty, u16(r))
188	
189	def rr.a(m: Maybe<&2, String>) -> RR:
190	  match m:
191	    case None{}:
192	      RRBad{}
193	    case Some{r}:
194	      rr.b(u16(r))
195	
196	# §4.1.3: NAME, TYPE, CLASS, TTL, RDLENGTH, RDATA.
197	def rr(+s: String) -> RR:
198	  rr.a(name.skip(String.length(s), s))
199	
200	# fuel = records left. The first A record wins; CNAMEs before it are skipped.
201	def answers.scan(fuel: Nat, r: RR) -> Maybe<&2, String>:
202	  match fuel:
203	    case 0n:
204	      None{}
205	    case 1n+f:
206	      match r:
207	        case RRBad{}:
208	          None{}
209	        case RRA{ip}:
210	          Some{ip}
211	        case RRSkip{s}:
212	          answers.scan(f, rr(s))
213	
214	def questions(fuel: Nat, m: Maybe<&2, String>) -> Maybe<&2, String>:
215	  match fuel:
216	    case 0n:
217	      m
218	    case 1n+f:
219	      questions(f, name.then(m, 4n))
220	
221	def body.go(an: U32, m: Maybe<&2, String>) -> Maybe<&2, String>:
222	  match m:
223	    case None{}:
224	      None{}
225	    case Some{s}:
226	      answers.scan(U32.to_nat(an), rr(s))
227	
228	def body.an(qd: U32, m: Maybe<&2, Cur>) -> Maybe<&2, String>:
229	  match m:
230	    case None{}:
231	      None{}
232	    case Some{Cur{an, r}}:
233	      body.go(an, questions(U32.to_nat(qd), skip(r, 4n)))
234	
235	def body.qd(m: Maybe<&2, Cur>) -> Maybe<&2, String>:
236	  match m:
237	    case None{}:
238	      None{}
239	    case Some{Cur{qd, r}}:
240	      body.an(qd, u16(r))
241	
242	# §4.1.1: QR set, opcode QUERY, not truncated, RCODE 0.
243	def flags.ok(+f: U32) -> Bool:
244	  Bool.and(Bool.and(U32.is_eq(U32.and(f, 32768), 32768), U32.is_eq(U32.and(f, 30720), 0)), Bool.and(U32.is_eq(U32.and(f, 512), 0), U32.is_eq(U32.and(f, 15), 0)))
245	
246	def head.flags.if(r: String, ok: Bool) -> Maybe<&2, String>:
247	  match ok:
248	    case False{}:
249	      None{}
250	    case True{}:
251	      body.qd(u16(r))
252	
253	def head.flags(m: Maybe<&2, Cur>) -> Maybe<&2, String>:
254	  match m:
255	    case None{}:
256	      None{}
257	    case Some{Cur{+f, r}}:
258	      head.flags.if(r, flags.ok(f))
259	
260	def head.id.if(r: String, ok: Bool) -> Maybe<&2, String>:
261	  match ok:
262	    case False{}:
263	      None{}
264	    case True{}:
265	      head.flags(u16(r))
266	
267	def head.id(id: U32, m: Maybe<&2, Cur>) -> Maybe<&2, String>:
268	  match m:
269	    case None{}:
270	      None{}
271	    case Some{Cur{got, r}}:
272	      head.id.if(r, U32.is_eq(id, got))
273	
274	# The IPv4 address the response gives for the query with this id.
275	def answer(id: U32, msg: String) -> Maybe<&2, String>:
276	  head.id(id, u16(msg))
277	
278	# resolv.conf
279	# -----------
280	
281	def nameservers.add(+line: String, rest: List<&2, String>, hit: Bool) -> List<&2, String>:
282	  match hit:
283	    case False{}:
284	      rest
285	    case True{}:
286	      Con{String.trim(String.drop(line, 10n)), rest}
287	
288	def nameservers.go(xs: List<&2, String>) -> List<&2, String>:
289	  match xs:
290	    case Nil{}:
291	      Nil{}
292	    case Con{+line, t}:
293	      nameservers.add(line, nameservers.go(t), String.starts_with(line, "nameserver"))
294	
295	# Every nameserver line, in file order.
296	def nameservers(conf: String) -> List<&2, String>:
297	  nameservers.go(String.lines(conf))
298	
299	def nameserver.first(xs: List<&2, String>) -> Maybe<&2, String>:
300	  match xs:
301	    case Nil{}:
302	      None{}
303	    case Con{n, t}:
304	      Some{n}
305	
306	# The first nameserver line of a resolv.conf.
307	def nameserver(conf: String) -> Maybe<&2, String>:
308	  nameserver.first(nameservers(conf))
309	
310	# Resolve
311	# -------
312	
313	def none() -> IO(Maybe<&2, String>):
314	  IO.pure(Maybe<&2, String>, None{})
315	
316	def ipv4.go(s: String) -> Bool:
317	  match s:
318	    case SNil{}:
319	      True{}
320	    case SCon{Chr{+c}, t}:
321	      Bool.and(Bool.or(U32.is_eq(c, 46), Bool.and(U32.is_le(48, c), U32.is_le(c, 57))), ipv4.go(t))
322	
323	# Digits and dots; TCP.connect rejects anything else that is not an address.
324	def ipv4(+s: String) -> Bool:
325	  Bool.and(Bool.not(String.is_empty(s)), ipv4.go(s))
326	
327	# One attempt's outcome: try again, or this answer (None: no address).
328	type Try is Data:
329	  Again{}
330	  Got{ip: Maybe<&2, String>}
331	
332	def try.pure(s: Socket, t: Try) -> IO(Socket & Try):
333	  IO.pure(Socket & Try, (s, t))
334	
335	# A datagram from anyone but the nameserver's port 53 is ignored (retried).
336	def try.from(s: Socket, +ns: String, id: U32, hpd: String & U32 & String) -> IO(Socket & Try):
337	  (h, +p, d) = hpd
338	  try.pure(s, Bool.pick(Try, Bool.and(String.eq(h, ns), U32.is_eq(p, 53)), Got{answer(id, d)}, Again{}))
339	
340	def try.back(ns: String, id: U32, m: Socket & Result<&1, &1, U32 & String, String & U32 & String>) -> IO(Socket & Try):
341	  (s, r) = m
342	  match r:
343	    case Fail{e}:
344	      try.pure(s, Again{})
345	    case Done{hpd}:
346	      try.from(s, ns, id, hpd)
347	
348	# §4.2.1 and resolv.conf defaults: 5 s per attempt.
349	def try.sent(ns: String, id: U32, m: Socket & Result<&1, &1, U32 & String, Unit>) -> IO(Socket & Try):
350	  (s, r) = m
351	  match r:
352	    case Fail{e}:
353	      try.pure(s, Again{})
354	    case Done{u}:
355	      do IO<Socket & Try>:
356	        back : Socket & Result<&1, &1, U32 & String, String & U32 & String> <- Wire.recv_from(s, 512, 5000)
357	        try.back(ns, id, back)
358	
359	def try.once(s: Socket, +ns: String, id: U32, q: String) -> IO(Socket & Try):
360	  do IO<Socket & Try>:
361	    sent : Socket & Result<&1, &1, U32 & String, Unit> <- Wire.send_to(s, ns, 53, q)
362	    try.sent(ns, id, sent)
363	
364	def try.close(s: Socket, r: Maybe<&2, String>) -> IO(Maybe<&2, String>):
365	  do IO<Maybe<&2, String>>:
366	    Socket.close(s)
367	    return r
368	
369	def tries.end(st: Socket & Try) -> IO(Maybe<&2, String>):
370	  (s, t) = st
371	  match t:
372	    case Got{ip}:
373	      try.close(s, ip)
374	    case Again{}:
375	      try.close(s, None{})
376	
377	# fuel: attempts left (resolv.conf's default is 2).
378	def tries(fuel: Nat, +ns: String, +id: U32, +q: String, st: Socket & Try) -> IO(Maybe<&2, String>):
379	  match fuel:
380	    case 0n:
381	      tries.end(st)
382	    case 1n+f:
383	      (s, t) = st
384	      match t:
385	        case Got{ip}:
386	          try.close(s, ip)
387	        case Again{}:
388	          do IO<Maybe<&2, String>>:
389	            next : Socket & Try <- try.once(s, ns, id, q)
390	            tries(f, ns, id, q, next)
391	
392	def resolve.sock(ns: String, id: U32, q: String, r: Result<&1, &1, U32 & String, Socket>) -> IO(Maybe<&2, String>):
393	  match r:
394	    case Fail{e}:
395	      none()
396	    case Done{s}:
397	      tries(2n, ns, id, q, (s, Again{}))
398	
399	def resolve.q(ns: String, id: U32, q: Maybe<&2, String>) -> IO(Maybe<&2, String>):
400	  match q:
401	    case None{}:
402	      none()
403	    case Some{bytes}:
404	      do IO<Maybe<&2, String>>:
405	        u : Result<&1, &1, U32 & String, Socket> <- UDP.bind(0)
406	        resolve.sock(ns, id, bytes, u)
407	
408	# §7.3: a random id makes forged answers harder to land.
409	def resolve.id(name: String, ns: String, r: Result<&1, &1, U32 & String, U32>) -> IO(Maybe<&2, String>):
410	  match r:
411	    case Fail{e}:
412	      none()
413	    case Done{x}:
414	      +id = U32.and(x, 65535)
415	      resolve.q(ns, id, query(id, name))
416	
417	def resolve.ns(name: String, ns: Maybe<&2, String>) -> IO(Maybe<&2, String>):
418	  match ns:
419	    case None{}:
420	      none()
421	    case Some{n}:
422	      do IO<Maybe<&2, String>>:
423	        r : Result<&1, &1, U32 & String, U32> <- IO.random_u32()
424	        resolve.id(name, n, r)
425	
426	# Ask one nameserver. Tests use this; resolve reads /etc/resolv.conf.
427	def resolve.at(+host: String, ns: String) -> IO(Maybe<&2, String>):
428	  resolve.ns(host, Some{ns})
429	
430	def conf.text(r: Result<&1, &1, U32 & String, String>) -> String:
431	  match r:
432	    case Fail{e}:
433	      ""
434	    case Done{s}:
435	      s
436	
437	# ponytail: 3 nameservers; a longer resolv.conf ignores the rest
438	def resolve.n3(+name: String, xs: List<&2, String>) -> IO(Maybe<&2, String>):
439	  match xs:
440	    case Nil{}:
441	      none()
442	    case Con{+ns, t}:
443	      resolve.at(name, ns)
444	
445	def resolve.n2b(+name: String, rest: List<&2, String>, ip: Maybe<&2, String>) -> IO(Maybe<&2, String>):
446	  match ip:
447	    case Some{s}:
448	      IO.pure(Maybe<&2, String>, Some{s})
449	    case None{}:
450	      resolve.n3(name, rest)
451	
452	def resolve.n2(+name: String, xs: List<&2, String>) -> IO(Maybe<&2, String>):
453	  match xs:
454	    case Nil{}:
455	      none()
456	    case Con{+ns, t}:
457	      do IO<Maybe<&2, String>>:
458	        ip : Maybe<&2, String> <- resolve.at(name, ns)
459	        resolve.n2b(name, t, ip)
460	
461	def resolve.n1b(+name: String, rest: List<&2, String>, ip: Maybe<&2, String>) -> IO(Maybe<&2, String>):
462	  match ip:
463	    case Some{s}:
464	      IO.pure(Maybe<&2, String>, Some{s})
465	    case None{}:
466	      resolve.n2(name, rest)
467	
468	def resolve.list(+name: String, xs: List<&2, String>) -> IO(Maybe<&2, String>):
469	  match xs:
470	    case Nil{}:
471	      none()
472	    case Con{+ns, t}:
473	      do IO<Maybe<&2, String>>:
474	        ip : Maybe<&2, String> <- resolve.at(name, ns)
475	        resolve.n1b(name, t, ip)
476	
477	def resolve.read(name: String, m: File & Result<&1, &1, U32 & String, String>) -> IO(Maybe<&2, String>):
478	  (f, r) = m
479	  do IO<Maybe<&2, String>>:
480	    File.close(f)
481	    resolve.list(name, nameservers(conf.text(r)))
482	
483	def resolve.conf(name: String, r: Result<&1, &1, U32 & String, File>) -> IO(Maybe<&2, String>):
484	  match r:
485	    case Fail{e}:
486	      none()
487	    case Done{f}:
488	      do IO<Maybe<&2, String>>:
489	        m : File & Result<&1, &1, U32 & String, String> <- File.read(f, 65536)
490	        resolve.read(name, m)
491	
492	def resolve.dns(name: String) -> IO(Maybe<&2, String>):
493	  do IO<Maybe<&2, String>>:
494	    f : Result<&1, &1, U32 & String, File> <- File.open("/etc/resolv.conf", "r")
495	    resolve.conf(name, f)
496	
497	
498	def hosts.pick(found: Maybe<&2, String>, now: Maybe<&2, String>) -> Maybe<&2, String>:
499	  match found:
500	    case Some{s}:
501	      Some{s}
502	    case None{}:
503	      now
504	
505	def hosts.hit(+ip: String, hit: Bool) -> Maybe<&2, String>:
506	  match hit:
507	    case True{}:
508	      Some{ip}
509	    case False{}:
510	      None{}
511	
512	def hosts.names(xs: List<&2, String>, +name: String, +ip: String, found: Maybe<&2, String>) -> Maybe<&2, String>:
513	  match xs:
514	    case Nil{}:
515	      found
516	    case Con{+n, t}:
517	      hosts.names(t, name, ip, hosts.pick(found, hosts.hit(ip, String.eq(String.to_lower(n), name))))
518	
519	def hosts.ip(+name: String, +ip: String, ns: List<&2, String>, ok: Bool) -> Maybe<&2, String>:
520	  match ok:
521	    case False{}:
522	      None{}
523	    case True{}:
524	      hosts.names(ns, name, ip, None{})
525	
526	def hosts.line(+name: String, xs: List<&2, String>) -> Maybe<&2, String>:
527	  match xs:
528	    case Nil{}:
529	      None{}
530	    case Con{+ip, ns}:
531	      hosts.ip(name, ip, ns, ipv4(ip))
532	
533	def hosts.sp(+c: U32, tab: Bool) -> U32:
534	  match tab:
535	    case True{}:
536	      32
537	    case False{}:
538	      c
539	
540	def hosts.flat(s: String) -> String:
541	  match s:
542	    case SNil{}:
543	      ""
544	    case SCon{Chr{+c}, t}:
545	      SCon{Chr{hosts.sp(c, U32.is_eq(c, 9))}, hosts.flat(t)}
546	
547	def hosts.scan(xs: List<&2, String>, +name: String, found: Maybe<&2, String>) -> Maybe<&2, String>:
548	  match xs:
549	    case Nil{}:
550	      found
551	    case Con{line, t}:
552	      hosts.scan(t, name, hosts.pick(found, hosts.line(name, String.split(hosts.flat(line), ' '))))
553	
554	# First IPv4 address for name in an /etc/hosts file. The name is already lowercase.
555	def hosts(+name: String, text: String) -> Maybe<&2, String>:
556	  hosts.scan(String.lines(text), name, None{})
557	
558	def resolve.from(+name: String, m: Maybe<&2, String>) -> IO(Maybe<&2, String>):
559	  match m:
560	    case Some{ip}:
561	      IO.pure(Maybe<&2, String>, Some{ip})
562	    case None{}:
563	      resolve.dns(name)
564	
565	def resolve.text(+name: String, r: Result<&1, &1, U32 & String, String>) -> IO(Maybe<&2, String>):
566	  match r:
567	    case Fail{e}:
568	      resolve.dns(name)
569	    case Done{s}:
570	      resolve.from(name, hosts(name, s))
571	
572	def resolve.opened(+name: String, m: File & Result<&1, &1, U32 & String, String>) -> IO(Maybe<&2, String>):
573	  (f, r) = m
574	  do IO<Maybe<&2, String>>:
575	    File.close(f)
576	    resolve.text(name, r)
577	
578	def resolve.hosts(+name: String, r: Result<&1, &1, U32 & String, File>) -> IO(Maybe<&2, String>):
579	  match r:
580	    case Fail{e}:
581	      resolve.dns(name)
582	    case Done{f}:
583	      do IO<Maybe<&2, String>>:
584	        m : File & Result<&1, &1, U32 & String, String> <- File.read(f, 65536)
585	        resolve.opened(name, m)
586	
587	# An IPv4 address for host: literals as is, localhost (RFC 6761), /etc/hosts, else DNS.
588	# ponytail: no AAAA, no TCP retry on TC
589	def resolve.ask(+name: String) -> IO(Maybe<&2, String>):
590	  do IO<Maybe<&2, String>>:
591	    f : Result<&1, &1, U32 & String, File> <- File.open("/etc/hosts", "r")
592	    resolve.hosts(name, f)
593	
594	def resolve(+host: String) -> IO(Maybe<&2, String>):
595	  Bool.pick(IO(Maybe<&2, String>), ipv4(host), IO.pure(Maybe<&2, String>, Some{host}), Bool.pick(IO(Maybe<&2, String>), String.eq(host, "localhost"), IO.pure(Maybe<&2, String>, Some{"127.0.0.1"}), resolve.ask(String.to_lower(host))))