From aa3a172c27171f7e1ce200d928f85a88eef63842 Mon Sep 17 00:00:00 2001
From: Armin Preiml <apreiml@strohwolke.at>
Date: Thu, 23 Jan 2025 15:54:30 +0100
Subject: [PATCH] ev::dns: nomem changes

Signed-off-by: Armin Preiml <apreiml@strohwolke.at>

---
 ev/dns/dns.ha | 30 ++++++++++++++++++++++++------

diff --git a/ev/dns/dns.ha b/ev/dns/dns.ha
index b09d38975ceba0fea83cd0c38b8a241e85e181b6..52f86d35745005ddbebe41fa34d7ef128f9afd7f 100644
--- a/ev/dns/dns.ha
+++ b/ev/dns/dns.ha
@@ -16,7 +16,7 @@
 // Callback for a [[query]] operation.
 export type querycb = fn(
 	user: nullable *opaque,
-	r: (*dns::message | dns::error),
+	r: (*dns::message | dns::error | nomem),
 ) void;
 
 type qstate = struct {
@@ -56,7 +56,7 @@ 	query: *dns::message,
 	cb: *querycb,
 	user: nullable *opaque,
 	servers: ip::addr...
-) (ev::req | dns::error | net::error | errors::error) = {
+) (ev::req | nomem | dns::error | net::error | errors::error) = {
 	if (len(servers) == 0) {
 		const rconf = resolvconf::load();
 		servers = rconf.nameservers;
@@ -66,20 +66,31 @@ 		// Fall back to localhost
 		servers = [ip::LOCAL_V6, ip::LOCAL_V4];
 	};
 
+	let stateok = false;
 	const socket4 = ev::listen_udp(loop, ip::ANY_V4, 0)?;
+	defer if (!stateok) ev::close(socket4);
 	const socket6 = ev::listen_udp(loop, ip::ANY_V6, 0)?;
+	defer if (!stateok) ev::close(socket6);
+
 	const timeout = ev::newtimer(loop, &timeoutcb, time::clock::MONOTONIC)?;
+	defer if (!stateok) ev::close(timeout);
+
+	let rbuf: []u8 = alloc([0...], 512)?;
+	defer if (!stateok) free(rbuf);
+
 	let state = alloc(qstate {
 		loop = loop,
 		socket4 = socket4,
 		socket6 = socket6,
 		timer = timeout,
-		rbuf = alloc([0...], 512),
 		rid = query.header.id,
 		cb = cb,
+		rbuf = rbuf,
 		user = user,
 		...
-	});
+	})?;
+	defer if (!stateok) free(state);
+
 	state.qlen = dns::encode(state.query, query)?: u16;
 	ev::setuser(socket4, state);
 	ev::setuser(socket6, state);
@@ -101,6 +112,7 @@ 	};
 
 	state.r4 = ev::recvfrom(socket4, &qrecvcb, state.rbuf);
 	state.r6 = ev::recvfrom(socket6, &qrecvcb, state.rbuf);
+	stateok = true;
 	return ev::mkreq(&query_cancel, state);
 };
 
@@ -119,7 +131,7 @@ 	free(q.rbuf);
 	free(q);
 };
 
-fn query_complete(q: *qstate, r: (*dns::message | dns::error)) void = {
+fn query_complete(q: *qstate, r: (*dns::message | dns::error | nomem)) void = {
 	const cb = q.cb;
 	const user = q.user;
 	cb(user, r);
@@ -262,7 +274,13 @@ 	};
 
 	const rlen = endian::begetu16(q.zbuf);
 	q.rid = rlen;
-	q.rbuf = alloc([0...], rlen);
+	q.rbuf = match (alloc([0u8...], rlen)) {
+	case let rbuf: []u8 =>
+		yield rbuf;
+	case nomem =>
+		query_complete(q, nomem);
+		return;
+	};
 	q.r4 = ev::read(file, &qtcp_readdata_cb, q.rbuf);
 };
 

-- 
2.48.1