consul/command/agent/dns.go

package agent

import (
	"fmt"
	"github.com/miekg/dns"
	"io"
	"log"
	"time"
)

// DNSServer is used to wrap an Agent and expose various
// service discovery endpoints using a DNS interface.
type DNSServer struct {
	agent      *Agent
	dnsHandler *dns.ServeMux
	dnsServer  *dns.Server
	logger     *log.Logger
}

// NewDNSServer starts a new DNS server to provide an agent interface
func NewDNSServer(agent *Agent, logOutput io.Writer, bind string) (*DNSServer, error) {
	// Construct the DNS components
	mux := dns.NewServeMux()

	// Setup the server
	server := &dns.Server{
		Addr:    bind,
		Net:     "udp",
		Handler: mux,
		UDPSize: 65535,
	}

	// Create the server
	srv := &DNSServer{
		agent:      agent,
		dnsHandler: mux,
		dnsServer:  server,
		logger:     log.New(logOutput, "", log.LstdFlags),
	}

	// Register mux handlers
	mux.HandleFunc("consul.", srv.handleConsul)

	// Async start the DNS Server, handle a potential error
	errCh := make(chan error, 1)
	go func() {
		srv.logger.Printf("[INFO] dns: starting server at %v", bind)
		err := server.ListenAndServe()
		srv.logger.Printf("[ERR] dns: error starting server: %v", err)
		errCh <- err
	}()

	// Check the server is running, do a test lookup
	checkCh := make(chan error, 1)
	go func() {
		// This is jank, but we have no way to edge trigger on
		// the start of our server, so we just wait and hope it is up.
		time.Sleep(50 * time.Millisecond)

		m := new(dns.Msg)
		m.SetQuestion("_test.consul.", dns.TypeANY)

		c := new(dns.Client)
		in, _, err := c.Exchange(m, bind)
		if err != nil {
			checkCh <- err
			return
		}

		srv.logger.Printf("resp %#v", in)
		if len(in.Answer) == 0 {
			checkCh <- fmt.Errorf("no response to test message")
			return
		}
		close(checkCh)
	}()

	// Wait for either the check, listen error, or timeout
	select {
	case e := <-errCh:
		return srv, e
	case e := <-checkCh:
		return srv, e
	case <-time.After(time.Second):
		return srv, fmt.Errorf("timeout setting up DNS server")
	}
	return srv, nil
}

// handleConsul is used to handle DNS queries in the ".consul." domain
func (d *DNSServer) handleConsul(resp dns.ResponseWriter, req *dns.Msg) {
	q := req.Question[0]
	d.logger.Printf("[DEBUG] dns: request for %v", q)

	if q.Qtype != dns.TypeANY && q.Qtype != dns.TypeTXT {
		return
	}

	// Always respond with TXT "ok"
	m := new(dns.Msg)
	m.SetReply(req)
	header := dns.RR_Header{Name: q.Name, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 0}
	txt := &dns.TXT{header, []string{"ok"}}
	m.Answer = append(m.Answer, txt)
	resp.WriteMsg(m)
}