From 1023a2e0a2733eecaf6a5a9394df31fc3339b171 Mon Sep 17 00:00:00 2001
From: Cottand <nico.dc@outlook.com>
Date: Mon, 6 Nov 2023 21:27:25 +0000
Subject: [PATCH] fix accidental pipelining of http doh requests

---
 doh.go        | 19 +++++++++++--------
 grimd_test.go |  5 ++---
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/doh.go b/doh.go
index a8b470e..83225aa 100644
--- a/doh.go
+++ b/doh.go
@@ -12,7 +12,6 @@ import (
 	stdlog "log"
 	"net"
 	"net/http"
-	"strconv"
 	"time"
 )
 
@@ -109,9 +108,10 @@ func (s *ServerHTTPS) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	var writer = DohResponseWriter{remoteAddr: r.RemoteAddr, host: r.Host, delegate: w}
+	var writer = DohResponseWriter{remoteAddr: r.RemoteAddr, host: r.Host, delegate: w, completed: make(chan empty)}
 	s.handler.ServeDNS(&writer, msg)
-	if writer.err != nil {
+	_, ok := <-writer.completed
+	if writer.err != nil || ok != true {
 		return
 	}
 
@@ -170,7 +170,7 @@ func getRequestToMsg(req *http.Request) (*dns.Msg, error) {
 }
 
 func base64ToMsg(b64 string) (*dns.Msg, error) {
-	buf, err := b64Enc.DecodeString(b64)
+	buf, err := base64.RawURLEncoding.DecodeString(b64)
 	if err != nil {
 		return nil, err
 	}
@@ -181,16 +181,16 @@ func base64ToMsg(b64 string) (*dns.Msg, error) {
 	return m, err
 }
 
-var b64Enc = base64.RawURLEncoding
-
-var _ dns.ResponseWriter = &DohResponseWriter{}
+type empty struct{}
 
+// DohResponseWriter implements dns.ResponseWriter
 type DohResponseWriter struct {
 	msg        *dns.Msg
 	remoteAddr string
 	delegate   http.ResponseWriter
 	host       string
 	err        error
+	completed  chan empty
 }
 
 // See section 4.2.1 of RFC 8484.
@@ -215,13 +215,16 @@ func (w *DohResponseWriter) RemoteAddr() net.Addr {
 }
 
 func (w *DohResponseWriter) WriteMsg(msg *dns.Msg) error {
+	defer func() {
+		w.completed <- empty{}
+		close(w.completed)
+	}()
 	w.msg = msg
 	buf, err := msg.Pack()
 	if err != nil {
 		w.handleErr(err)
 		return err
 	}
-	w.delegate.Header().Set("Content-Length", strconv.Itoa(len(buf)))
 	w.delegate.Header().Set("Content-Type", mimeTypeDOH)
 	_, err = w.Write(buf)
 	if err != nil {
diff --git a/grimd_test.go b/grimd_test.go
index 10da850..bfa1042 100644
--- a/grimd_test.go
+++ b/grimd_test.go
@@ -172,13 +172,12 @@ func TestDohIntegration(t *testing.T) {
 	})
 }
 
+// TestDohAsProxy checks that DoH works for non-custom records
 func TestDohAsProxy(t *testing.T) {
-	t.SkipNow() // see https://github.com/coredns/coredns/issues/4547
 	dohBind := "localhost:8181"
 	integrationTest(func(c *Config) {
 		c.DnsOverHttpServer.Bind = dohBind
 		c.DnsOverHttpServer.Enabled = true
-		//c.CustomDNSRecords = []string{"example.com          IN  A       10.10.0.1 "}
 	}, func(_ *dns.Client, _ string) {
 		resp, err := http.Get("http://" + dohBind + "/dns-query?dns=AAABAAABAAAAAAAAA3d3dwdleGFtcGxlA2NvbQAAAQAB")
 
@@ -194,7 +193,7 @@ func TestDohAsProxy(t *testing.T) {
 		msg := dns.Msg{}
 		err = msg.Unpack(respPacket)
 		if err != nil {
-			t.Fatalf("unexpected error during lookup %v", err)
+			t.Fatalf("unexpected error during lookup %v (response len=%vB)", err, len(respPacket))
 		}
 
 		if len(msg.Answer) < 1 {