From 74dbfccc11c4ce9e88fb0822d946e1d80e0f06be Mon Sep 17 00:00:00 2001
From: Miek Gieben <miek@miek.nl>
Date: Wed, 28 Nov 2018 18:37:11 +0000
Subject: [PATCH] Code Review

Signed-off-by: Miek Gieben <miek@miek.nl>
---
 acceptfunc.go | 31 ++++++++++++++++++++-----------
 server.go     |  9 +++++----
 2 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/acceptfunc.go b/acceptfunc.go
index 5f680e2b..57d8333c 100644
--- a/acceptfunc.go
+++ b/acceptfunc.go
@@ -3,7 +3,7 @@ package dns
 // MsgAcceptFunc is used early in the server code to accept or reject a message with RcodeFormatError.
 // There are to booleans to be returned, once signaling the rejection and another to signal if
 // a reply is to be send back (you want to prevent DNS ping-pong and not reply to a response for instance).
-type MsgAcceptFunc func(dh Header) (accept bool, respond bool)
+type MsgAcceptFunc func(dh Header) MsgAcceptAction
 
 // DefaultMsgAcceptFunc checks the request and will reject if:
 //
@@ -14,33 +14,42 @@ type MsgAcceptFunc func(dh Header) (accept bool, respond bool)
 // * has more than 0 RRs in the Answer section
 // * has more than 0 RRs in the Authority section
 // * has more than 2 RRs in the Additional section
-var DefaultMsgAcceptFunc = defaultMsgAcceptFunc
+var DefaultMsgAcceptFunc MsgAcceptFunc = defaultMsgAcceptFunc
 
-var defaultMsgAcceptFunc = func(dh Header) (bool, bool) {
+// MsgAcceptAction represents the action to be taken.
+type MsgAcceptAction int
+
+const (
+	MsgAccept MsgAcceptAction = iota
+	MsgReject
+	MsgIgnore
+)
+
+var defaultMsgAcceptFunc = func(dh Header) MsgAcceptAction {
 	if isResponse := dh.Bits&_QR != 0; isResponse {
-		return false, false
+		return MsgIgnore
 	}
 
 	// Don't allow dynamic updates, because then the sections can contain a whole bunch of RRs.
 	opcode := int(dh.Bits>>11) & 0xF
 	if opcode != OpcodeQuery && opcode != OpcodeNotify {
-		return false, true
+		return MsgReject
 	}
 
 	if isZero := dh.Bits&_Z != 0; isZero {
-		return false, true
+		return MsgReject
 	}
 	if dh.Qdcount != 1 {
-		return false, true
+		return MsgReject
 	}
 	if dh.Ancount != 0 {
-		return false, true
+		return MsgReject
 	}
 	if dh.Nscount != 0 {
-		return false, true
+		return MsgReject
 	}
 	if dh.Arcount > 2 {
-		return false, true
+		return MsgReject
 	}
-	return true, true
+	return MsgAccept
 }
diff --git a/server.go b/server.go
index 0cf1bd92..c5a0a4d6 100644
--- a/server.go
+++ b/server.go
@@ -214,8 +214,8 @@ type Server struct {
 	// Whether to set the SO_REUSEPORT socket option, allowing multiple listeners to be bound to a single address.
 	// It is only supported on go1.11+ and when using ListenAndServe.
 	ReusePort bool
-	// AcceptMsgFunc will check the incoming message and will reject it early in the process. This function must be
-	// defined. By default DefaultMsgAcceptFunc will be used.
+	// AcceptMsgFunc will check the incoming message and will reject it early in the process.
+	// By default DefaultMsgAcceptFunc will be used.
 	MsgAcceptFunc MsgAcceptFunc
 
 	// UDP packet or TCP connection queue
@@ -642,10 +642,11 @@ func (srv *Server) serveDNS(w *response) {
 	req := new(Msg)
 	req.setHdr(dh)
 
-	if accept, respond := srv.MsgAcceptFunc(dh); !accept {
-		if !respond {
+	if action := srv.MsgAcceptFunc(dh); action != MsgAccept {
+		if action == MsgIgnore {
 			return
 		}
+
 		req.SetRcodeFormatError(req)
 		// Are we allowed to delete any OPT records here?
 		req.Ns, req.Answer, req.Extra = nil, nil, nil