Fix resolv ...

* Fix mbedtls_ssl_read() is not called repeatedly until no data is
  returned
* Impl exit code of proone-resolv
* Add test for resolv

3 files changed, 194 insertions, 114 deletions

diff --git a/src/proone-resolv.c b/src/proone-resolv.c
index d27df4f..94759b8 100644
--- a/src/proone-resolv.c
+++ b/src/proone-resolv.c
@@ -38,6 +38,11 @@ pth_cond_t prm_cond = PTH_COND_INIT;
 prne_pth_cv_t prm_cv = { &prm_lock, &prm_cond, false };
 
 prne_resolv_t *resolv = NULL;
+struct {
+	bool parse_err;
+	bool query_err;
+	bool proc;
+} exec_result;
 
 static void proc_prompt_line (char *line, const size_t line_len) {
 	static regmatch_t rm[3];
@@ -93,6 +98,7 @@ static void proc_prompt_line (char *line, const size_t line_len) {
 				(prne_llist_element_t)prm) != NULL);
 		}
 		else {
+			exec_result.query_err = true;
 			perror("* Queue failed");
 			prne_resolv_free_prm(prm);
 			prne_free(prm);
@@ -101,6 +107,7 @@ static void proc_prompt_line (char *line, const size_t line_len) {
 	else if (line_len > 0 &&
 		regexec(&empty_line_regex, line, 0, NULL, 0) != 0)
 	{
+		exec_result.parse_err = true;
 		fprintf(stderr, "* Line not recognised.\n");
 	}
 }
@@ -148,6 +155,7 @@ static void *stdin_wkr_entry (void *ctx) {
 			else {
 				line_buf_cnt = 0;
 				if (!missed_line) {
+					exec_result.parse_err = true;
 					fprintf(stderr, "* Line too long!\n");
 				}
 				missed_line = true;
@@ -194,8 +202,12 @@ static void *stdout_wkr_entry (void *ctx) {
 				if (prm->fut->qr == PRNE_RESOLV_QR_OK ||
 					prm->fut->qr == PRNE_RESOLV_QR_STATUS)
 				{
+					exec_result.proc = true;
 					status_str = prne_resolv_rcode_tostr(prm->fut->status);
 				}
+				else {
+					exec_result.query_err = true;
+				}
 				if (status_str == NULL) {
 					status_str = "";
 				}
@@ -299,10 +311,10 @@ int main (void) {
 		&prmpt_regex,
 		"(A|AAAA|TXT)\\s+([a-z0-9\\-\\.]+)",
 		REG_ICASE | REG_EXTENDED) == 0);
-	// org regex: ^\s+$
+	// org regex: (^[#;].*)|(^(\s+)?$)
 	prne_assert(regcomp(
 		&empty_line_regex,
-		"^\\s+$",
+		"(^[#;].*)|(^(\\s+)?$)",
 		REG_NOSUB | REG_EXTENDED) == 0);
 	prne_mbedtls_entropy_init(&entropy);
 	mbedtls_ctr_drbg_init(&rnd);
@@ -363,5 +375,14 @@ int main (void) {
 	}
 	prne_free_llist(&prm_list);
 
-	return 0;
+	if (exec_result.proc) {
+		if (exec_result.parse_err || exec_result.query_err) {
+			return 3;
+		}
+		return 0;
+	}
+	if (exec_result.parse_err) {
+		return 2;
+	}
+	return 1;
 }
diff --git a/src/resolv.c b/src/resolv.c
index c841e4e..817c9f8 100644
--- a/src/resolv.c
+++ b/src/resolv.c
@@ -1134,139 +1134,160 @@ static void resolv_proc_expired (prne_resolv_t *ctx) {
 	}
 }
 
+static bool resolv_proc_out (prne_resolv_t *ctx) {
+	int f_ret;
+	bool ret = false;
+
+	while (ctx->iobuf[1].len > 0) {
+		f_ret = mbedtls_ssl_write(
+			&ctx->ssl.ctx,
+			ctx->iobuf[1].m,
+			ctx->iobuf[1].len);
+		if (f_ret == MBEDTLS_ERR_SSL_WANT_READ ||
+			f_ret == MBEDTLS_ERR_SSL_WANT_WRITE)
+		{
+			break;
+		}
+		if (f_ret <= 0) {
+			// we don't renegotiate with terrorists.
+			resolv_close_sck(ctx, &RESOLV_CONN_ERR_PAUSE, true);
+			break;
+		}
+		prne_iobuf_shift(&ctx->iobuf[1], -f_ret);
+		ret = true;
+	}
+
+	return ret;
+}
+
+static bool resolv_proc_in (prne_resolv_t *ctx) {
+	size_t pos, msg_len;
+	bool err_flag = false, ret = false;
+
+	pos = 0;
+	while (true) {
+		if (pos + 1 >= ctx->iobuf[0].len) {
+			break;
+		}
+		msg_len = prne_recmb_msb16(
+			ctx->iobuf[0].m[pos],
+			ctx->iobuf[0].m[pos + 1]);
+		if (msg_len > 512) { // unimplemented.
+			prne_dbgpf(
+				"* [resolv_wkr] Protocol error: "
+				"received %zu bytes long msg.\n"
+				"* [resolv_wkr] Dropping connection!\n",
+				msg_len);
+			// try to get qid
+			if (ctx->iobuf[0].len > pos + 4) {
+				const uint16_t qid = prne_recmb_msb16(
+					ctx->iobuf[0].m[pos + 2],
+					ctx->iobuf[0].m[pos + 3]);
+				const prne_imap_tuple_t *tpl = prne_imap_lookup(
+					&ctx->qid_map,
+					qid);
+
+				if (tpl->val != (prne_imap_val_type_t)NULL) {
+					query_entry_t *qent = (query_entry_t*)tpl->val;
+					qent->fut.qr = PRNE_RESOLV_QR_IMPL;
+					resolv_disown_qent(qent);
+				}
+				prne_imap_erase(&ctx->qid_map, qid);
+			}
+			resolv_close_sck(ctx, &RESOLV_CONN_ERR_PAUSE, true);
+			break;
+		}
+		if (pos + 1 + msg_len >= ctx->iobuf[0].len) {
+			break;
+		}
+
+		ret |= resolv_proc_dns_msg(
+			ctx,
+			ctx->iobuf[0].m + pos + 2,
+			msg_len,
+			&err_flag);
+		if (err_flag) {
+			resolv_close_sck(ctx, &RESOLV_CONN_ERR_PAUSE, true);
+			return false;
+		}
+		pos += 2 + msg_len;
+	}
+
+	prne_iobuf_shift(&ctx->iobuf[0], -pos);
+	return ret;
+}
+
 static void resolv_proc_q (prne_resolv_t *ctx) {
-	int pollret, ret;
-	short pfd_events;
-	bool proc;
+	int ret;
+	bool proc = false;
 	/*
 	* The server could be sending gibberish response messages that look legit.
 	* Timeout the loop when we don't receive response we recognise in time.
 	*/
 	pth_event_t ev = NULL;
 
-LOOP:
-	proc = false;
-	while (ctx->qlist.size > 0 || ctx->qid_map.size > 0) {
+	while (true) {
 		resolv_proc_expired(ctx);
-
-		if (ctx->iobuf[1].len > 0 || ctx->qid_map.size < RESOLV_PIPELINE_SIZE) {
-			pfd_events = POLLIN | POLLOUT;
-		}
-		else {
-			pfd_events = POLLIN;
+		if (ctx->qlist.size == 0 && ctx->qid_map.size == 0) {
+			break;
 		}
 
+		// connect
+		if (ctx->act_sck_pfd.fd < 0) {
+			proc = false;
+		}
 		if (!resolv_ensure_conn(ctx)) {
-			goto LOOP;
+			continue;
 		}
 
-		if (proc || ev == NULL) {
-			pth_event_free(ev, FALSE);
-			ev = pth_event(
-				PTH_EVENT_TIME,
-				prne_pth_tstimeout(RESOLV_SCK_OP_TIMEOUT));
+		// write
+		if (resolv_send_dns_msgs(ctx)) {
+			proc |= resolv_proc_out(ctx);
 		}
-		proc = false;
-
-		ctx->act_sck_pfd.events = pfd_events;
-		prne_assert(ev != NULL); // fatal without timeout
-		pollret = prne_pth_poll(&ctx->act_sck_pfd, 1, -1, ev);
-
-		if (pth_event_status(ev) != PTH_STATUS_PENDING) {
-			resolv_close_sck(ctx, NULL, true);
+		if (ctx->act_sck_pfd.fd < 0) {
+			continue;
 		}
-		else if (pollret > 0) {
-			if (ctx->act_sck_pfd.revents & (POLLERR | POLLHUP | POLLNVAL)) {
-				resolv_close_sck(ctx, &RESOLV_CONN_ERR_PAUSE, true);
-				goto LOOP;
-			}
-			if (ctx->act_sck_pfd.revents & POLLIN) {
-				size_t pos, msg_len;
-				bool err_flag = false;
-
-				ret = mbedtls_ssl_read(
-					&ctx->ssl.ctx,
-					ctx->iobuf[0].m + ctx->iobuf[0].len,
-					ctx->iobuf[0].avail);
-				if (ret <= 0) {
-					// we don't renegotiate with terrorists.
-					resolv_close_sck(ctx, &RESOLV_CONN_ERR_PAUSE, true);
-					goto LOOP;
-				}
-				prne_iobuf_shift(&ctx->iobuf[0], ret);
-
-				pos = 0;
-				while (true) {
-					if (pos + 1 >= ctx->iobuf[0].len) {
-						break;
-					}
-					msg_len = prne_recmb_msb16(
-						ctx->iobuf[0].m[pos],
-						ctx->iobuf[0].m[pos + 1]);
-					if (msg_len > 512) { // unimplemented.
-						prne_dbgpf(
-							"* [resolv_wkr] Protocol error: "
-							"received %zu bytes long msg.\n"
-							"* [resolv_wkr] Dropping connection!\n",
-							msg_len);
-						// try to get qid
-						if (ctx->iobuf[0].len > pos + 4) {
-							const uint16_t qid = prne_recmb_msb16(
-								ctx->iobuf[0].m[pos + 2],
-								ctx->iobuf[0].m[pos + 3]);
-							const prne_imap_tuple_t *tpl = prne_imap_lookup(
-								&ctx->qid_map,
-								qid);
-
-							if (tpl->val != (prne_imap_val_type_t)NULL) {
-								query_entry_t *qent = (query_entry_t*)tpl->val;
-								qent->fut.qr = PRNE_RESOLV_QR_IMPL;
-								resolv_disown_qent(qent);
-							}
-							prne_imap_erase(&ctx->qid_map, qid);
-						}
-						resolv_close_sck(ctx, &RESOLV_CONN_ERR_PAUSE, true);
-						goto LOOP;
-					}
-					if (pos + 1 + msg_len >= ctx->iobuf[0].len) {
-						break;
-					}
-
-					proc |= resolv_proc_dns_msg(
-						ctx,
-						ctx->iobuf[0].m + pos + 2,
-						msg_len,
-						&err_flag);
-					if (err_flag) {
-						resolv_close_sck(ctx, &RESOLV_CONN_ERR_PAUSE, true);
-						goto LOOP;
-					}
-					pos += 2 + msg_len;
-				}
 
-				prne_iobuf_shift(&ctx->iobuf[0], -pos);
+		// read
+		ret = mbedtls_ssl_read(
+			&ctx->ssl.ctx,
+			ctx->iobuf[0].m + ctx->iobuf[0].len,
+			ctx->iobuf[0].avail);
+		if (ret == MBEDTLS_ERR_SSL_WANT_READ ||
+			ret == MBEDTLS_ERR_SSL_WANT_WRITE)
+		{
+			if (proc || ev == NULL) {
+				// timer reset
+				pth_event_free(ev, FALSE);
+				ev = pth_event(
+					PTH_EVENT_TIME,
+					prne_pth_tstimeout(RESOLV_SCK_OP_TIMEOUT));
 			}
-
-			if ((ctx->act_sck_pfd.revents & POLLOUT) && ctx->iobuf[1].len > 0) {
-				ret = mbedtls_ssl_write(
-					&ctx->ssl.ctx,
-					ctx->iobuf[1].m,
-					ctx->iobuf[1].len);
-				if (ret <= 0) {
-					// we don't renegotiate with terrorists.
-					resolv_close_sck(ctx, &RESOLV_CONN_ERR_PAUSE, true);
-					goto LOOP;
-				}
-				prne_iobuf_shift(&ctx->iobuf[1], -ret);
+			if (ctx->iobuf[1].len > 0 || ret == MBEDTLS_ERR_SSL_WANT_WRITE) {
+				ctx->act_sck_pfd.events = POLLIN | POLLOUT;
+			}
+			else {
+				ctx->act_sck_pfd.events = POLLIN;
 			}
 
-			if (ctx->iobuf[1].len == 0) {
-				proc |= resolv_send_dns_msgs(ctx);
+			prne_assert(ev != NULL); // fatal without timeout
+			ret = prne_pth_poll(&ctx->act_sck_pfd, 1, -1, ev);
+
+			if (pth_event_status(ev) != PTH_STATUS_PENDING) {
+				resolv_close_sck(ctx, NULL, true);
+			}
+			else if (ret < 0) {
+				resolv_close_sck(ctx, &RESOLV_RSRC_ERR_PAUSE, true);
 			}
+			continue;
 		}
-		else {
-			resolv_close_sck(ctx, &RESOLV_RSRC_ERR_PAUSE, true);
+		else if (ret <= 0) {
+			resolv_close_sck(ctx, &RESOLV_CONN_ERR_PAUSE, true);
+			continue;
 		}
+		prne_iobuf_shift(&ctx->iobuf[0], ret);
+
+		proc |= resolv_proc_in(ctx);
 	}
 
 	pth_event_free(ev, FALSE);
diff --git a/src/test-resolv.sh b/src/test-resolv.sh
new file mode 100755
index 0000000..c10d305
--- /dev/null
+++ b/src/test-resolv.sh
@@ -0,0 +1,38 @@
+#!/bin/bash
+assert_ec () {
+	echo -n "$3: " >&2
+	if [ $1 -ne $2 ]; then
+		echo "FAIL (expected=$2, returned=$1)" >&2
+		exit 1
+	else
+		echo "OK" >&2
+	fi
+}
+SUBJECT_EXEC="./proone-resolv"
+
+echo "a example.com" | "$SUBJECT_EXEC"
+assert_ec $? 0 "Single NOERROR execution"
+
+echo "a example.test" | "$SUBJECT_EXEC"
+assert_ec $? 0 "Single NXDOMAIN execution"
+
+cat << EOF | "$SUBJECT_EXEC"
+; Queue more than RESOLV_PIPELINE_SIZE(4)
+a example.com
+aaaa example.com
+a www.example.com
+aaaa www.example.com
+txt kernel.org
+a www.google.com
+aaaa www.google.com
+txt example.com
+a www.kernel.org
+aaaa www.kernel.org
+EOF
+assert_ec $? 0 "Queue congestion"
+
+cat << EOF | "$SUBJECT_EXEC"
+aaaa example.com
+txt example.test
+EOF
+assert_ec $? 0 "Mixed result"