// DNS server implementation.

package dns

import (
	
	
	
	
	
	
	
	
	
)

// Default maximum number of TCP queries before we close the socket.
const maxTCPQueries = 128

// aLongTimeAgo is a non-zero time, far in the past, used for
// immediate cancellation of network operations.
var aLongTimeAgo = time.Unix(1, 0)

// Handler is implemented by any value that implements ServeDNS.
type Handler interface {
	ServeDNS(w ResponseWriter, r *Msg)
}

// The HandlerFunc type is an adapter to allow the use of
// ordinary functions as DNS handlers.  If f is a function
// with the appropriate signature, HandlerFunc(f) is a
// Handler object that calls f.
type HandlerFunc func(ResponseWriter, *Msg)

// ServeDNS calls f(w, r).
func ( HandlerFunc) ( ResponseWriter,  *Msg) {
	(, )
}

// A ResponseWriter interface is used by an DNS handler to
// construct an DNS response.
type ResponseWriter interface {
	// LocalAddr returns the net.Addr of the server
	LocalAddr() net.Addr
	// RemoteAddr returns the net.Addr of the client that sent the current request.
	RemoteAddr() net.Addr
	// WriteMsg writes a reply back to the client.
	WriteMsg(*Msg) error
	// Write writes a raw buffer back to the client.
	Write([]byte) (int, error)
	// Close closes the connection.
	Close() error
	// TsigStatus returns the status of the Tsig.
	TsigStatus() error
	// TsigTimersOnly sets the tsig timers only boolean.
	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.
	Hijack()
}

// A ConnectionStater interface is used by a DNS Handler to access TLS connection state
// when available.
type ConnectionStater interface {
	ConnectionState() *tls.ConnectionState
}

type response struct {
	closed         bool // connection has been closed
	hijacked       bool // connection has been hijacked by handler
	tsigTimersOnly bool
	tsigStatus     error
	tsigRequestMAC string
	tsigProvider   TsigProvider
	udp            net.PacketConn // i/o connection if UDP was used
	tcp            net.Conn       // i/o connection if TCP was used
	udpSession     *SessionUDP    // oob data to get egress interface right
	pcSession      net.Addr       // address to use when writing to a generic net.PacketConn
	writer         Writer         // writer to output the raw DNS bits
}

// handleRefused returns a HandlerFunc that returns REFUSED for every request it gets.
func handleRefused( ResponseWriter,  *Msg) {
	 := new(Msg)
	.SetRcode(, RcodeRefused)
	.WriteMsg()
}

// HandleFailed returns a HandlerFunc that returns SERVFAIL for every request it gets.
// Deprecated: This function is going away.
func ( ResponseWriter,  *Msg) {
	 := new(Msg)
	.SetRcode(, RcodeServerFailure)
	// does not matter if this write fails
	.WriteMsg()
}

// ListenAndServe Starts a server on address and network specified Invoke handler
// for incoming queries.
func ( string,  string,  Handler) error {
	 := &Server{Addr: , Net: , Handler: }
	return .ListenAndServe()
}

// ListenAndServeTLS acts like http.ListenAndServeTLS, more information in
// http://golang.org/pkg/net/http/#ListenAndServeTLS
func (, ,  string,  Handler) error {
	,  := tls.LoadX509KeyPair(, )
	if  != nil {
		return 
	}

	 := tls.Config{
		Certificates: []tls.Certificate{},
	}

	 := &Server{
		Addr:      ,
		Net:       "tcp-tls",
		TLSConfig: &,
		Handler:   ,
	}

	return .ListenAndServe()
}

// ActivateAndServe activates a server with a listener from systemd,
// l and p should not both be non-nil.
// If both l and p are not nil only p will be used.
// Invoke handler for incoming queries.
func ( net.Listener,  net.PacketConn,  Handler) error {
	 := &Server{Listener: , PacketConn: , Handler: }
	return .ActivateAndServe()
}

// Writer writes raw DNS messages; each call to Write should send an entire message.
type Writer interface {
	io.Writer
}

// Reader reads raw DNS messages; each call to ReadTCP or ReadUDP should return an entire message.
type Reader interface {
	// ReadTCP reads a raw message from a TCP connection. Implementations may alter
	// connection properties, for example the read-deadline.
	ReadTCP(conn net.Conn, timeout time.Duration) ([]byte, error)
	// ReadUDP reads a raw message from a UDP connection. Implementations may alter
	// connection properties, for example the read-deadline.
	ReadUDP(conn *net.UDPConn, timeout time.Duration) ([]byte, *SessionUDP, error)
}

// PacketConnReader is an optional interface that Readers can implement to support using generic net.PacketConns.
type PacketConnReader interface {
	Reader

	// ReadPacketConn reads a raw message from a generic net.PacketConn UDP connection. Implementations may
	// alter connection properties, for example the read-deadline.
	ReadPacketConn(conn net.PacketConn, timeout time.Duration) ([]byte, net.Addr, error)
}

// defaultReader is an adapter for the Server struct that implements the Reader and
// PacketConnReader interfaces using the readTCP, readUDP and readPacketConn funcs
// of the embedded Server.
type defaultReader struct {
	*Server
}

var _ PacketConnReader = defaultReader{}

func ( defaultReader) ( net.Conn,  time.Duration) ([]byte, error) {
	return .readTCP(, )
}

func ( defaultReader) ( *net.UDPConn,  time.Duration) ([]byte, *SessionUDP, error) {
	return .readUDP(, )
}

func ( defaultReader) ( net.PacketConn,  time.Duration) ([]byte, net.Addr, error) {
	return .readPacketConn(, )
}

// DecorateReader is a decorator hook for extending or supplanting the functionality of a Reader.
// Implementations should never return a nil Reader.
// Readers should also implement the optional PacketConnReader interface.
// PacketConnReader is required to use a generic net.PacketConn.
type DecorateReader func(Reader) Reader

// DecorateWriter is a decorator hook for extending or supplanting the functionality of a Writer.
// Implementations should never return a nil Writer.
type DecorateWriter func(Writer) Writer

// MsgInvalidFunc is a listener hook for observing incoming messages that were discarded
// because they could not be parsed.
// Every message that is read by a Reader will eventually be provided to the Handler,
// rejected (or ignored) by the MsgAcceptFunc, or passed to this function.
type MsgInvalidFunc func(m []byte, err error)

func ( []byte,  error) {}

// A Server defines parameters for running an DNS server.
type Server struct {
	// Address to listen on, ":dns" if empty.
	Addr string
	// if "tcp" or "tcp-tls" (DNS over TLS) it will invoke a TCP listener, otherwise an UDP one
	Net string
	// TCP Listener to use, this is to aid in systemd's socket activation.
	Listener net.Listener
	// TLS connection configuration
	TLSConfig *tls.Config
	// UDP "Listener" to use, this is to aid in systemd's socket activation.
	PacketConn net.PacketConn
	// Handler to invoke, dns.DefaultServeMux if nil.
	Handler Handler
	// Default buffer size to use to read incoming UDP messages. If not set
	// it defaults to MinMsgSize (512 B).
	UDPSize int
	// The net.Conn.SetReadTimeout value for new connections, defaults to 2 * time.Second.
	ReadTimeout time.Duration
	// The net.Conn.SetWriteTimeout value for new connections, defaults to 2 * time.Second.
	WriteTimeout time.Duration
	// TCP idle timeout for multiple queries, if nil, defaults to 8 * time.Second (RFC 5966).
	IdleTimeout func() time.Duration
	// An implementation of the TsigProvider interface. If defined it replaces TsigSecret and is used for all TSIG operations.
	TsigProvider TsigProvider
	// Secret(s) for Tsig map[<zonename>]<base64 secret>. The zonename must be in canonical form (lowercase, fqdn, see RFC 4034 Section 6.2).
	TsigSecret map[string]string
	// If NotifyStartedFunc is set it is called once the server has started listening.
	NotifyStartedFunc func()
	// DecorateReader is optional, allows customization of the process that reads raw DNS messages.
	// The decorated reader must not mutate the data read from the conn.
	DecorateReader DecorateReader
	// DecorateWriter is optional, allows customization of the process that writes raw DNS messages.
	DecorateWriter DecorateWriter
	// Maximum number of TCP queries before we close the socket. Default is maxTCPQueries (unlimited if -1).
	MaxTCPQueries int
	// Whether to set the SO_REUSEPORT socket option, allowing multiple listeners to be bound to a single address.
	// It is only supported on certain GOOSes and when using ListenAndServe.
	ReusePort bool
	// Whether to set the SO_REUSEADDR socket option, allowing multiple listeners to be bound to a single address.
	// Crucially this allows binding when an existing server is listening on `0.0.0.0` or `::`.
	// It is only supported on certain GOOSes and when using ListenAndServe.
	ReuseAddr bool
	// AcceptMsgFunc will check the incoming message and will reject it early in the process.
	// By default DefaultMsgAcceptFunc will be used.
	MsgAcceptFunc MsgAcceptFunc
	// MsgInvalidFunc is optional, will be called if a message is received but cannot be parsed.
	MsgInvalidFunc MsgInvalidFunc

	// Shutdown handling
	lock     sync.RWMutex
	started  bool
	shutdown chan struct{}
	conns    map[net.Conn]struct{}

	// A pool for UDP message buffers.
	udpPool sync.Pool
}

func ( *Server) () TsigProvider {
	if .TsigProvider != nil {
		return .TsigProvider
	}
	if .TsigSecret != nil {
		return tsigSecretProvider(.TsigSecret)
	}
	return nil
}

func ( *Server) () bool {
	.lock.RLock()
	 := .started
	.lock.RUnlock()
	return 
}

func makeUDPBuffer( int) func() interface{} {
	return func() interface{} {
		return make([]byte, )
	}
}

func ( *Server) () {
	.shutdown = make(chan struct{})
	.conns = make(map[net.Conn]struct{})

	if .UDPSize == 0 {
		.UDPSize = MinMsgSize
	}
	if .MsgAcceptFunc == nil {
		.MsgAcceptFunc = DefaultMsgAcceptFunc
	}
	if .MsgInvalidFunc == nil {
		.MsgInvalidFunc = DefaultMsgInvalidFunc
	}
	if .Handler == nil {
		.Handler = DefaultServeMux
	}

	.udpPool.New = makeUDPBuffer(.UDPSize)
}

func unlockOnce( sync.Locker) func() {
	var  sync.Once
	return func() { .Do(.Unlock) }
}

// ListenAndServe starts a nameserver on the configured address in *Server.
func ( *Server) () error {
	 := unlockOnce(&.lock)
	.lock.Lock()
	defer ()

	if .started {
		return &Error{err: "server already started"}
	}

	 := .Addr
	if  == "" {
		 = ":domain"
	}

	.init()

	switch .Net {
	case "tcp", "tcp4", "tcp6":
		,  := listenTCP(.Net, , .ReusePort, .ReuseAddr)
		if  != nil {
			return 
		}
		.Listener = 
		.started = true
		()
		return .serveTCP()
	case "tcp-tls", "tcp4-tls", "tcp6-tls":
		if .TLSConfig == nil || (len(.TLSConfig.Certificates) == 0 && .TLSConfig.GetCertificate == nil) {
			return errors.New("dns: neither Certificates nor GetCertificate set in Config")
		}
		 := strings.TrimSuffix(.Net, "-tls")
		,  := listenTCP(, , .ReusePort, .ReuseAddr)
		if  != nil {
			return 
		}
		 = tls.NewListener(, .TLSConfig)
		.Listener = 
		.started = true
		()
		return .serveTCP()
	case "udp", "udp4", "udp6":
		,  := listenUDP(.Net, , .ReusePort, .ReuseAddr)
		if  != nil {
			return 
		}
		 := .(*net.UDPConn)
		if  := setUDPSocketOptions();  != nil {
			.Close()
			return 
		}
		.PacketConn = 
		.started = true
		()
		return .serveUDP()
	}
	return &Error{err: "bad network"}
}

// ActivateAndServe starts a nameserver with the PacketConn or Listener
// configured in *Server. Its main use is to start a server from systemd.
func ( *Server) () error {
	 := unlockOnce(&.lock)
	.lock.Lock()
	defer ()

	if .started {
		return &Error{err: "server already started"}
	}

	.init()

	if .PacketConn != nil {
		// Check PacketConn interface's type is valid and value
		// is not nil
		if ,  := .PacketConn.(*net.UDPConn);  &&  != nil {
			if  := setUDPSocketOptions();  != nil {
				return 
			}
		}
		.started = true
		()
		return .serveUDP(.PacketConn)
	}
	if .Listener != nil {
		.started = true
		()
		return .serveTCP(.Listener)
	}
	return &Error{err: "bad listeners"}
}

// Shutdown shuts down a server. After a call to Shutdown, ListenAndServe and
// ActivateAndServe will return.
func ( *Server) () error {
	return .ShutdownContext(context.Background())
}

// ShutdownContext shuts down a server. After a call to ShutdownContext,
// ListenAndServe and ActivateAndServe will return.
//
// A context.Context may be passed to limit how long to wait for connections
// to terminate.
func ( *Server) ( context.Context) error {
	.lock.Lock()
	if !.started {
		.lock.Unlock()
		return &Error{err: "server not started"}
	}

	.started = false

	if .PacketConn != nil {
		.PacketConn.SetReadDeadline(aLongTimeAgo) // Unblock reads
	}

	if .Listener != nil {
		.Listener.Close()
	}

	for  := range .conns {
		.SetReadDeadline(aLongTimeAgo) // Unblock reads
	}

	.lock.Unlock()

	if testShutdownNotify != nil {
		testShutdownNotify.Broadcast()
	}

	var  error
	select {
	case <-.shutdown:
	case <-.Done():
		 = .Err()
	}

	if .PacketConn != nil {
		.PacketConn.Close()
	}

	return 
}

var testShutdownNotify *sync.Cond

// getReadTimeout is a helper func to use system timeout if server did not intend to change it.
func ( *Server) () time.Duration {
	if .ReadTimeout != 0 {
		return .ReadTimeout
	}
	return dnsTimeout
}

// serveTCP starts a TCP listener for the server.
func ( *Server) ( net.Listener) error {
	defer .Close()

	if .NotifyStartedFunc != nil {
		.NotifyStartedFunc()
	}

	var  sync.WaitGroup
	defer func() {
		.Wait()
		close(.shutdown)
	}()

	for .isStarted() {
		,  := .Accept()
		if  != nil {
			if !.isStarted() {
				return nil
			}
			if ,  := .(net.Error);  && .Temporary() {
				continue
			}
			return 
		}
		.lock.Lock()
		// Track the connection to allow unblocking reads on shutdown.
		.conns[] = struct{}{}
		.lock.Unlock()
		.Add(1)
		go .serveTCPConn(&, )
	}

	return nil
}

// serveUDP starts a UDP listener for the server.
func ( *Server) ( net.PacketConn) error {
	defer .Close()

	 := Reader(defaultReader{})
	if .DecorateReader != nil {
		 = .DecorateReader()
	}

	,  := .(*net.UDPConn)
	,  := .(PacketConnReader)
	if ! && ! {
		return &Error{err: "PacketConnReader was not implemented on Reader returned from DecorateReader but is required for net.PacketConn"}
	}

	if .NotifyStartedFunc != nil {
		.NotifyStartedFunc()
	}

	var  sync.WaitGroup
	defer func() {
		.Wait()
		close(.shutdown)
	}()

	 := .getReadTimeout()
	// deadline is not used here
	for .isStarted() {
		var (
			    []byte
			  net.Addr
			 *SessionUDP
			  error
		)
		if  {
			, ,  = .ReadUDP(, )
		} else {
			, ,  = .ReadPacketConn(, )
		}
		if  != nil {
			if !.isStarted() {
				return nil
			}
			if ,  := .(net.Error);  && .Temporary() {
				continue
			}
			return 
		}
		if len() < headerSize {
			if cap() == .UDPSize {
				.udpPool.Put([:.UDPSize])
			}
			.MsgInvalidFunc(, ErrShortRead)
			continue
		}
		.Add(1)
		go .serveUDPPacket(&, , , , )
	}

	return nil
}

// Serve a new TCP connection.
func ( *Server) ( *sync.WaitGroup,  net.Conn) {
	 := &response{tsigProvider: .tsigProvider(), tcp: }
	if .DecorateWriter != nil {
		.writer = .DecorateWriter()
	} else {
		.writer = 
	}

	 := Reader(defaultReader{})
	if .DecorateReader != nil {
		 = .DecorateReader()
	}

	 := tcpIdleTimeout
	if .IdleTimeout != nil {
		 = .IdleTimeout()
	}

	 := .getReadTimeout()

	 := .MaxTCPQueries
	if  == 0 {
		 = maxTCPQueries
	}

	for  := 0; ( <  ||  == -1) && .isStarted(); ++ {
		,  := .ReadTCP(.tcp, )
		if  != nil {
			// TODO(tmthrgd): handle error
			break
		}
		.serveDNS(, )
		if .closed {
			break // Close() was called
		}
		if .hijacked {
			break // client will call Close() themselves
		}
		// The first read uses the read timeout, the rest use the
		// idle timeout.
		 = 
	}

	if !.hijacked {
		.Close()
	}

	.lock.Lock()
	delete(.conns, .tcp)
	.lock.Unlock()

	.Done()
}

// Serve a new UDP request.
func ( *Server) ( *sync.WaitGroup,  []byte,  net.PacketConn,  *SessionUDP,  net.Addr) {
	 := &response{tsigProvider: .tsigProvider(), udp: , udpSession: , pcSession: }
	if .DecorateWriter != nil {
		.writer = .DecorateWriter()
	} else {
		.writer = 
	}

	.serveDNS(, )
	.Done()
}

func ( *Server) ( []byte,  *response) {
	, ,  := unpackMsgHdr(, 0)
	if  != nil {
		.MsgInvalidFunc(, )
		// Let client hang, they are sending crap; any reply can be used to amplify.
		return
	}

	 := new(Msg)
	.setHdr()

	switch  := .MsgAcceptFunc();  {
	case MsgAccept:
		 := .unpack(, , )
		if  == nil {
			break
		}

		.MsgInvalidFunc(, )
		fallthrough
	case MsgReject, MsgRejectNotImplemented:
		 := .Opcode
		.SetRcodeFormatError()
		.Zero = false
		if  == MsgRejectNotImplemented {
			.Opcode = 
			.Rcode = RcodeNotImplemented
		}

		// Are we allowed to delete any OPT records here?
		.Ns, .Answer, .Extra = nil, nil, nil

		.WriteMsg()
		fallthrough
	case MsgIgnore:
		if .udp != nil && cap() == .UDPSize {
			.udpPool.Put([:.UDPSize])
		}

		return
	}

	.tsigStatus = nil
	if .tsigProvider != nil {
		if  := .IsTsig();  != nil {
			.tsigStatus = TsigVerifyWithProvider(, .tsigProvider, "", false)
			.tsigTimersOnly = false
			.tsigRequestMAC = .MAC
		}
	}

	if .udp != nil && cap() == .UDPSize {
		.udpPool.Put([:.UDPSize])
	}

	.Handler.ServeDNS(, ) // Writes back to the client
}

func ( *Server) ( net.Conn,  time.Duration) ([]byte, error) {
	// If we race with ShutdownContext, the read deadline may
	// have been set in the distant past to unblock the read
	// below. We must not override it, otherwise we may block
	// ShutdownContext.
	.lock.RLock()
	if .started {
		.SetReadDeadline(time.Now().Add())
	}
	.lock.RUnlock()

	var  uint16
	if  := binary.Read(, binary.BigEndian, &);  != nil {
		return nil, 
	}

	 := make([]byte, )
	if ,  := io.ReadFull(, );  != nil {
		return nil, 
	}

	return , nil
}

func ( *Server) ( *net.UDPConn,  time.Duration) ([]byte, *SessionUDP, error) {
	.lock.RLock()
	if .started {
		// See the comment in readTCP above.
		.SetReadDeadline(time.Now().Add())
	}
	.lock.RUnlock()

	 := .udpPool.Get().([]byte)
	, ,  := ReadFromSessionUDP(, )
	if  != nil {
		.udpPool.Put()
		return nil, nil, 
	}
	 = [:]
	return , , nil
}

func ( *Server) ( net.PacketConn,  time.Duration) ([]byte, net.Addr, error) {
	.lock.RLock()
	if .started {
		// See the comment in readTCP above.
		.SetReadDeadline(time.Now().Add())
	}
	.lock.RUnlock()

	 := .udpPool.Get().([]byte)
	, ,  := .ReadFrom()
	if  != nil {
		.udpPool.Put()
		return nil, nil, 
	}
	 = [:]
	return , , nil
}

// WriteMsg implements the ResponseWriter.WriteMsg method.
func ( *response) ( *Msg) ( error) {
	if .closed {
		return &Error{err: "WriteMsg called after Close"}
	}

	var  []byte
	if .tsigProvider != nil { // if no provider, dont check for the tsig (which is a longer check)
		if  := .IsTsig();  != nil {
			, .tsigRequestMAC,  = TsigGenerateWithProvider(, .tsigProvider, .tsigRequestMAC, .tsigTimersOnly)
			if  != nil {
				return 
			}
			_,  = .writer.Write()
			return 
		}
	}
	,  = .Pack()
	if  != nil {
		return 
	}
	_,  = .writer.Write()
	return 
}

// Write implements the ResponseWriter.Write method.
func ( *response) ( []byte) (int, error) {
	if .closed {
		return 0, &Error{err: "Write called after Close"}
	}

	switch {
	case .udp != nil:
		if ,  := .udp.(*net.UDPConn);  {
			return WriteToSessionUDP(, , .udpSession)
		}
		return .udp.WriteTo(, .pcSession)
	case .tcp != nil:
		if len() > MaxMsgSize {
			return 0, &Error{err: "message too large"}
		}

		 := make([]byte, 2+len())
		binary.BigEndian.PutUint16(, uint16(len()))
		copy([2:], )
		return .tcp.Write()
	default:
		panic("dns: internal error: udp and tcp both nil")
	}
}

// LocalAddr implements the ResponseWriter.LocalAddr method.
func ( *response) () net.Addr {
	switch {
	case .udp != nil:
		return .udp.LocalAddr()
	case .tcp != nil:
		return .tcp.LocalAddr()
	default:
		panic("dns: internal error: udp and tcp both nil")
	}
}

// RemoteAddr implements the ResponseWriter.RemoteAddr method.
func ( *response) () net.Addr {
	switch {
	case .udpSession != nil:
		return .udpSession.RemoteAddr()
	case .pcSession != nil:
		return .pcSession
	case .tcp != nil:
		return .tcp.RemoteAddr()
	default:
		panic("dns: internal error: udpSession, pcSession and tcp are all nil")
	}
}

// TsigStatus implements the ResponseWriter.TsigStatus method.
func ( *response) () error { return .tsigStatus }

// TsigTimersOnly implements the ResponseWriter.TsigTimersOnly method.
func ( *response) ( bool) { .tsigTimersOnly =  }

// Hijack implements the ResponseWriter.Hijack method.
func ( *response) () { .hijacked = true }

// Close implements the ResponseWriter.Close method
func ( *response) () error {
	if .closed {
		return &Error{err: "connection already closed"}
	}
	.closed = true

	switch {
	case .udp != nil:
		// Can't close the udp conn, as that is actually the listener.
		return nil
	case .tcp != nil:
		return .tcp.Close()
	default:
		panic("dns: internal error: udp and tcp both nil")
	}
}

// ConnectionState() implements the ConnectionStater.ConnectionState() interface.
func ( *response) () *tls.ConnectionState {
	type  interface {
		() tls.ConnectionState
	}
	if ,  := .tcp.();  {
		 := .()
		return &
	}
	return nil
}