Merge pull request #84 from telekom-mms/feature/update-dnsproxy

Feature/update dnsproxy
telekom-mms · May 13, 2024 · 0bc82d1 · 0bc82d1
2 parents cef4fde + 5243520
commit 0bc82d1
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/internal/dnsproxy/proxy.go b/internal/dnsproxy/proxy.go
@@ -20,6 +20,14 @@ type Proxy struct {
 	closed  chan struct{}
 }
 
+// sendReport sends report over the reports channel.
+func (p *Proxy) sendReport(report *Report) {
+	select {
+	case p.reports <- report:
+	case <-p.done:
+	}
+}
+
 // handleRequest handles a dns client request.
 func (p *Proxy) handleRequest(w dns.ResponseWriter, r *dns.Msg) {
 	// make sure the client request is valid
@@ -86,7 +94,7 @@ func (p *Proxy) handleRequest(w dns.ResponseWriter, r *dns.Msg) {
 			return
 		}
 		report := NewReport(rr.Hdr.Name, rr.A, rr.Hdr.Ttl)
-		p.reports <- report
+		p.sendReport(report)
 		report.Wait()
 	}
 
@@ -99,7 +107,7 @@ func (p *Proxy) handleRequest(w dns.ResponseWriter, r *dns.Msg) {
 			return
 		}
 		report := NewReport(rr.Hdr.Name, rr.AAAA, rr.Hdr.Ttl)
-		p.reports <- report
+		p.sendReport(report)
 		report.Wait()
 	}
 

diff --git a/internal/dnsproxy/proxy_test.go b/internal/dnsproxy/proxy_test.go
@@ -91,12 +91,16 @@ func TestProxyHandleRequest(t *testing.T) {
 
 	// with watches in proxy
 	p.SetWatches([]string{"test.example.com."})
+	reportsDone := make(chan struct{})
 	go func() {
+		defer close(reportsDone)
 		for r := range p.Reports() {
 			r.Done()
 		}
 	}()
 	p.handleRequest(&responseWriter{}, &dns.Msg{Question: []dns.Question{{Name: "test.example.com."}}})
+	close(p.reports)
+	<-reportsDone
 }
 
 // TestProxyHandleRequest tests handleRequest of Proxy, DNS records.