dns

commit 521607802a026aae130395eb147b81d480b94e34
parent 554ced8329b2f5a7abf40a2131cf5e869c1e17d4
Author: Oliver Lowe <o@olowe.co>
Date:   Thu, 16 Dec 2021 20:11:37 +1100

dns: Separate sending messages from sending bytes

That made it easy to then let us ResponseWriter to implement io.Writer

Diffstat:
M
dns.go
 | 
25
++++++++++++-------------
M
server.go
 | 
21
++++++++++++++++++---

2 files changed, 30 insertions(+), 16 deletions(-)
diff --git a/dns.go b/dns.go
@@ -152,7 +152,7 @@ func ExchangeTLS(msg dnsmessage.Message, addr string) (dnsmessage.Message, error
 }
 
 func exchange(msg dnsmessage.Message, conn net.Conn) (dnsmessage.Message, error) {
-	if err := send(msg, conn); err != nil {
+	if err := sendMsg(msg, conn); err != nil {
 		return dnsmessage.Message{}, err
 	}
 	rmsg, err := receive(conn)
@@ -165,31 +165,30 @@ func exchange(msg dnsmessage.Message, conn net.Conn) (dnsmessage.Message, error)
 	return rmsg, nil
 }
 
-func send(msg dnsmessage.Message, conn net.Conn) error {
+func sendMsg(msg dnsmessage.Message, conn net.Conn) error {
 	packed, err := msg.Pack()
 	if err != nil {
 		return err
 	}
+	_, err = send(packed, conn)
+	return err
+}
+
+func send(p []byte, conn net.Conn) (int, error) {
 	if _, ok := conn.(net.PacketConn); ok {
-		if _, err := conn.Write(packed); err != nil {
-			return err
-		}
-		return nil
+		return conn.Write(p)
 	}
 	// DNS over TCP requires you to prepend the message with a
 	// 2-octet length field.
-	l := len(packed)
+	l := len(p)
 	m := make([]byte, 2+l)
 	m[0] = byte(l >> 8)
 	m[1] = byte(l)
-	copy(m[2:], packed)
-	if _, err := conn.Write(m); err != nil {
-		return err
-	}
-	return nil
+	copy(m[2:], p)
+	return conn.Write(m)
 }
 
-func sendPacket(msg dnsmessage.Message, conn net.PacketConn, addr net.Addr) error {
+func sendMsgTo(msg dnsmessage.Message, conn net.PacketConn, addr net.Addr) error {
 	packed, err := msg.Pack()
 	if err != nil {
 		return err
diff --git a/server.go b/server.go
@@ -5,6 +5,12 @@ import (
 	"net"
 )
 
+type Zone struct {
+	Name dnsmessage.Name
+	SOA dnsmessage.SOAResource
+	Resources []dnsmessage.Resource
+}
+
 // Server contains settings for running a DNS server. An empty Server
 // with a nil Handler is a valid configuration.
 type Server struct {
@@ -21,16 +27,25 @@ type response struct {
 	conn  net.Conn
 }
 
+func (r *response) Write(p []byte) (n int, err error) {
+	if r.pconn != nil {
+		return r.pconn.WriteTo(p, r.raddr)
+	}
+	return send(p, r.conn)
+}
+
 func (r *response) WriteMsg(msg dnsmessage.Message) error {
 	if r.pconn != nil {
-		return sendPacket(msg, r.pconn, r.raddr)
+		return sendMsgTo(msg, r.pconn, r.raddr)
 	}
-	return send(msg, r.conn)
+	return sendMsg(msg, r.conn)
 }
 
 // The ResponseWriter interface is used by a Handler to reply to
 // DNS requests.
 type ResponseWriter interface {
+	// Write writes the data to the underlying connection as a DNS response.
+	Write(p []byte) (n int, err error)
 	// WriteMsg writes the DNS message to the connection.
 	WriteMsg(dnsmessage.Message) error
 }
@@ -53,7 +68,7 @@ func (srv *Server) ServePacket(conn net.PacketConn) error {
 			var msg dnsmessage.Message
 			if err := msg.Unpack(buf[:n]); err != nil {
 				msg.Header.RCode = dnsmessage.RCodeRefused
-				sendPacket(msg, conn, raddr)
+				sendMsgTo(msg, conn, raddr)
 				return
 			}
 			resp := &response{raddr: raddr, pconn: conn}