From 399fafb679369ba01a418ab6c264b3cb3fe0c5b1 Mon Sep 17 00:00:00 2001
From: Chris Chapman <cchapman@hashicorp.com>
Date: Fri, 30 Sep 2022 14:51:12 -0700
Subject: [PATCH] Making suggested changes

---
 agent/agent.go                                |  6 ++---
 agent/grpc-external/services/dns/server.go    | 23 +++++++++----------
 .../grpc-external/services/dns/server_test.go |  2 +-
 3 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index 38ca5e4c3..ff41544ef 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -924,9 +924,9 @@ func (a *Agent) listenAndServeDNS() error {
 	s, _ := NewDNSServer(a)
 
 	grpcDNS.NewServer(grpcDNS.Config{
-		Logger:       a.logger.Named("grpc-api.dns"),
-		DNSServeMux:  s.mux,
-		LocalAddress: grpcDNS.Local{IP: net.IPv4(127, 0, 0, 1), Port: a.config.GRPCPort},
+		Logger:      a.logger.Named("grpc-api.dns"),
+		DNSServeMux: s.mux,
+		LocalAddr:   grpcDNS.LocalAddr{IP: net.IPv4(127, 0, 0, 1), Port: a.config.GRPCPort},
 	}).Register(a.externalGRPCServer)
 
 	a.dnsServers = append(a.dnsServers, s)
diff --git a/agent/grpc-external/services/dns/server.go b/agent/grpc-external/services/dns/server.go
index cef3a1073..bdcf38296 100644
--- a/agent/grpc-external/services/dns/server.go
+++ b/agent/grpc-external/services/dns/server.go
@@ -15,15 +15,15 @@ import (
 	"github.com/hashicorp/consul/proto-public/pbdns"
 )
 
-type Local struct {
+type LocalAddr struct {
 	IP   net.IP
 	Port int
 }
 
 type Config struct {
-	Logger       hclog.Logger
-	DNSServeMux  *dns.ServeMux
-	LocalAddress Local
+	Logger      hclog.Logger
+	DNSServeMux *dns.ServeMux
+	LocalAddr   LocalAddr
 }
 
 type Server struct {
@@ -40,7 +40,7 @@ func (s *Server) Register(grpcServer *grpc.Server) {
 
 // BufferResponseWriter writes a DNS response to a byte buffer.
 type BufferResponseWriter struct {
-	ResponseBuffer []byte
+	responseBuffer []byte
 	LocalAddress   net.Addr
 	RemoteAddress  net.Addr
 	Logger         hclog.Logger
@@ -64,15 +64,14 @@ func (b *BufferResponseWriter) WriteMsg(m *dns.Msg) error {
 		b.Logger.Error("error packing message", "err", err)
 		return err
 	}
-	b.ResponseBuffer = msgBytes
+	b.responseBuffer = msgBytes
 	return nil
 }
 
 // Write writes a raw buffer back to the client.
 func (b *BufferResponseWriter) Write(m []byte) (int, error) {
-	b.Logger.Info("Write was called")
-	copy(b.ResponseBuffer, m)
-	return len(b.ResponseBuffer), nil
+	b.Logger.Debug("Write was called")
+	return copy(b.responseBuffer, m), nil
 }
 
 // Close closes the connection.
@@ -108,11 +107,11 @@ func (s *Server) Query(ctx context.Context, req *pbdns.QueryRequest) (*pbdns.Que
 	switch req.GetProtocol() {
 	case pbdns.Protocol_PROTOCOL_TCP:
 		remote = pr.Addr
-		local = &net.TCPAddr{IP: s.LocalAddress.IP, Port: s.LocalAddress.Port}
+		local = &net.TCPAddr{IP: s.LocalAddr.IP, Port: s.LocalAddr.Port}
 	case pbdns.Protocol_PROTOCOL_UDP:
 		remoteAddr := pr.Addr.(*net.TCPAddr)
 		remote = &net.UDPAddr{IP: remoteAddr.IP, Port: remoteAddr.Port}
-		local = &net.UDPAddr{IP: s.LocalAddress.IP, Port: s.LocalAddress.Port}
+		local = &net.UDPAddr{IP: s.LocalAddr.IP, Port: s.LocalAddr.Port}
 	default:
 		return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("error protocol type not set: %v", req.GetProtocol()))
 	}
@@ -131,7 +130,7 @@ func (s *Server) Query(ctx context.Context, req *pbdns.QueryRequest) (*pbdns.Que
 	}
 	s.DNSServeMux.ServeDNS(respWriter, msg)
 
-	queryResponse := &pbdns.QueryResponse{Msg: respWriter.ResponseBuffer}
+	queryResponse := &pbdns.QueryResponse{Msg: respWriter.responseBuffer}
 
 	return queryResponse, nil
 }
diff --git a/agent/grpc-external/services/dns/server_test.go b/agent/grpc-external/services/dns/server_test.go
index e8107b46e..b477fed49 100644
--- a/agent/grpc-external/services/dns/server_test.go
+++ b/agent/grpc-external/services/dns/server_test.go
@@ -58,7 +58,7 @@ func (s *DNSTestSuite) TestProxy_Success() {
 	server := NewServer(Config{
 		Logger:      hclog.Default(),
 		DNSServeMux: mux,
-		LocalAddress: Local{
+		LocalAddr: LocalAddr{
 			net.IPv4(127, 0, 0, 1),
 			0,
 		},