From fc888600fdb0f695daad4a09adce55a01ec82c06 Mon Sep 17 00:00:00 2001
From: Johan Brandhorst-Satzkorn <johan.brandhorst@gmail.com>
Date: Tue, 9 Jul 2024 09:40:31 -0700
Subject: [PATCH] Add ExposesUnderlyingConns interface

The interface can be used by DNS handlers to access the underlying TCP
or UDP connection. This can be useful if more information is needed
about the request, such as to what port it was made.

Fixes #1487
---
 server.go      | 20 ++++++++++++++++++++
 server_test.go |  3 +++
 2 files changed, 23 insertions(+)

diff --git a/server.go b/server.go
index 81580d1e5..d9c9243a8 100644
--- a/server.go
+++ b/server.go
@@ -79,6 +79,26 @@ type response struct {
 	writer         Writer         // writer to output the raw DNS bits
 }
 
+// IncomingConn exposes the underlying TCP connection, if it is used.
+// For UDP connections, it returns nil.
+func (r *response) IncomingConn() net.Conn {
+	return r.tcp
+}
+
+// IncomingPacketConn exposes the underlying UDP connection, if it is used.
+// For TCP connections, it returns nil.
+func (r *response) IncomingPacketConn() net.PacketConn {
+	return r.udp
+}
+
+// The ExposesUnderlyingConns interface is used by a DNS Handler to access underlying connections.
+type ExposesUnderlyingConns interface {
+	// IncomingConn is only valid for TCP connections. It is otherwise nil.
+	IncomingConn() net.Conn
+	// IncomingPackageConn is only valid for UDP connections. It is otherwise nil.
+	IncomingPacketConn() net.PacketConn
+}
+
 // handleRefused returns a HandlerFunc that returns REFUSED for every request it gets.
 func handleRefused(w ResponseWriter, r *Msg) {
 	m := new(Msg)
diff --git a/server_test.go b/server_test.go
index 4fc2af329..3c21c06b8 100644
--- a/server_test.go
+++ b/server_test.go
@@ -17,6 +17,9 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
+// Assert that the response exposes the underlying connections.
+var _ ExposesUnderlyingConns = &response{}
+
 func HelloServer(w ResponseWriter, req *Msg) {
 	m := new(Msg)
 	m.SetReply(req)