Adding grpc handler for dns proxy

agent/grpc-external/services/dns/server.go

@@ -0,0 +1,136 @@
+package dns
+
+import (
+	"context"
+	"fmt"
+	"net"
+
+	"github.com/hashicorp/go-hclog"
+	"github.com/miekg/dns"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/status"
+
+	"github.com/hashicorp/consul/proto-public/pbdns"
+)
+
+type Local struct {
+	IP   net.IP
+	Port int
+}
+
+type Config struct {
+	Logger       hclog.Logger
+	DNSServeMux  *dns.ServeMux
+	LocalAddress Local
+}
+
+type Server struct {
+	Config
+}
+
+func NewServer(cfg Config) *Server {
+	return &Server{cfg}
+}
+
+func (s *Server) Register(grpcServer *grpc.Server) {
+	pbdns.RegisterDNSServiceServer(grpcServer, s)
+}
+
+// BufferResponseWriter writes a DNS response to a byte buffer.
+type BufferResponseWriter struct {
+	ResponseBuffer []byte
+	LocalAddress   net.Addr
+	RemoteAddress  net.Addr
+	Logger         hclog.Logger
+}
+
+// LocalAddr returns the net.Addr of the server
+func (b *BufferResponseWriter) LocalAddr() net.Addr {
+	return b.LocalAddress
+}
+
+// RemoteAddr returns the net.Addr of the client that sent the current request.
+func (b *BufferResponseWriter) RemoteAddr() net.Addr {
+	return b.RemoteAddress
+}
+
+// WriteMsg writes a reply back to the client.
+func (b *BufferResponseWriter) WriteMsg(m *dns.Msg) error {
+	// Pack message to bytes first.
+	msgBytes, err := m.Pack()
+	if err != nil {
+		b.Logger.Error("error packing message", "err", err)
+		return err
+	}
+	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
+}
+
+// Close closes the connection.
+func (b *BufferResponseWriter) Close() error {
+	// There's nothing for us to do here as we don't handle the connection.
+	return nil
+}
+
+// TsigStatus returns the status of the Tsig.
+func (b *BufferResponseWriter) TsigStatus() error {
+	// TSIG doesn't apply to this response writer.
+	return nil
+}
+
+// TsigTimersOnly sets the tsig timers only boolean.
+func (b *BufferResponseWriter) TsigTimersOnly(bool) {}
+
+// Hijack lets the caller take over the connection.
+// After a call to Hijack(), the DNS package will not do anything with the connection. {
+func (b *BufferResponseWriter) Hijack() {}
+
+func (s *Server) Query(ctx context.Context, req *pbdns.QueryRequest) (*pbdns.QueryResponse, error) {
+	pr, ok := peer.FromContext(ctx)
+	if !ok {
+		return nil, fmt.Errorf("error retrieving peer information from context")
+	}
+
+	var local net.Addr
+	var remote net.Addr
+	// We do this so that we switch to udp/tcp when handling the request since it will be proxied
+	// through consul through gRPC and we need to 'fake' the protocol to get the correct response
+	switch req.GetProtocol() {
+	case pbdns.Protocol_PROTOCOL_TCP:
+		remote = pr.Addr
+		local = &net.TCPAddr{IP: s.LocalAddress.IP, Port: s.LocalAddress.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}
+	default:
+		return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("error protocol type not set: %v", req.GetProtocol()))
+	}
+
+	respWriter := &BufferResponseWriter{
+		LocalAddress:  local,
+		RemoteAddress: remote,
+		Logger:        s.Logger,
+	}
+
+	msg := &dns.Msg{}
+	err := msg.Unpack(req.Msg)
+	if err != nil {
+		s.Logger.Error("error unpacking message", "err", err)
+		return nil, status.Error(codes.Internal, fmt.Sprintf("failure decoding dns request: %s", err.Error()))
+	}
+	s.DNSServeMux.ServeDNS(respWriter, msg)
+
+	queryResponse := &pbdns.QueryResponse{Msg: respWriter.ResponseBuffer}
+
+	return queryResponse, nil
+}

agent/grpc-external/services/dns/server_test.go

@@ -0,0 +1,127 @@
+package dns
+
+import (
+	"context"
+	"errors"
+	"net"
+	"testing"
+
+	"github.com/hashicorp/go-hclog"
+	"github.com/miekg/dns"
+	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/suite"
+	"google.golang.org/grpc"
+
+	"github.com/hashicorp/consul/agent/grpc-external/testutils"
+	"github.com/hashicorp/consul/proto-public/pbdns"
+)
+
+var txtRR = []string{"Hello world"}
+
+func helloServer(w dns.ResponseWriter, req *dns.Msg) {
+	m := new(dns.Msg)
+	m.SetReply(req)
+
+	m.Extra = make([]dns.RR, 1)
+	m.Extra[0] = &dns.TXT{
+		Hdr: dns.RR_Header{Name: m.Question[0].Name, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 0},
+		Txt: txtRR,
+	}
+	w.WriteMsg(m)
+}
+
+func testClient(t *testing.T, server *Server) pbdns.DNSServiceClient {
+	t.Helper()
+
+	addr := testutils.RunTestServer(t, server)
+
+	conn, err := grpc.DialContext(context.Background(), addr.String(), grpc.WithInsecure())
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		require.NoError(t, conn.Close())
+	})
+
+	return pbdns.NewDNSServiceClient(conn)
+}
+
+type DNSTestSuite struct {
+	suite.Suite
+}
+
+func TestDNS_suite(t *testing.T) {
+	suite.Run(t, new(DNSTestSuite))
+}
+
+func (s *DNSTestSuite) TestProxy_Success() {
+	mux := dns.NewServeMux()
+	mux.Handle(".", dns.HandlerFunc(helloServer))
+	server := NewServer(Config{
+		Logger:      hclog.Default(),
+		DNSServeMux: mux,
+		LocalAddress: Local{
+			net.IPv4(127, 0, 0, 1),
+			0,
+		},
+	})
+
+	client := testClient(s.T(), server)
+
+	testCases := map[string]struct {
+		question    string
+		clientQuery func(qR *pbdns.QueryRequest)
+		expectedErr error
+	}{
+
+		"happy path udp": {
+			question: "abc.com.",
+			clientQuery: func(qR *pbdns.QueryRequest) {
+				qR.Protocol = pbdns.Protocol_PROTOCOL_UDP
+			},
+		},
+		"happy path tcp": {
+			question: "abc.com.",
+			clientQuery: func(qR *pbdns.QueryRequest) {
+				qR.Protocol = pbdns.Protocol_PROTOCOL_TCP
+			},
+		},
+		"No protocol set": {
+			question:    "abc.com.",
+			clientQuery: func(qR *pbdns.QueryRequest) {},
+			expectedErr: errors.New("error protocol type not set: PROTOCOL_UNSET_UNSPECIFIED"),
+		},
+		"Invalid question": {
+			question: "notvalid",
+			clientQuery: func(qR *pbdns.QueryRequest) {
+				qR.Protocol = pbdns.Protocol_PROTOCOL_UDP
+			},
+			expectedErr: errors.New("failure decoding dns request"),
+		},
+	}
+
+	for name, tc := range testCases {
+		s.Run(name, func() {
+			req := dns.Msg{}
+			req.SetQuestion(tc.question, dns.TypeA)
+
+			bytes, _ := req.Pack()
+
+			clientReq := &pbdns.QueryRequest{Msg: bytes}
+			tc.clientQuery(clientReq)
+			clientResp, err := client.Query(context.Background(), clientReq)
+			if tc.expectedErr != nil {
+				s.Require().Error(err, "no errror calling gRPC endpoint")
+				s.Require().ErrorContains(err, tc.expectedErr.Error())
+			} else {
+				s.Require().NoError(err, "error calling gRPC endpoint")
+
+				resp := clientResp.GetMsg()
+				var dnsResp dns.Msg
+
+				err = dnsResp.Unpack(resp)
+				s.Require().NoError(err, "error unpacking dns response")
+				rr := dnsResp.Extra[0].(*dns.TXT)
+				s.Require().EqualValues(rr.Txt, txtRR)
+			}
+		})
+	}
+}