x/internal/util/dns/conn.go

package dns

import (
	"errors"
	"net"
	"sync"
	"sync/atomic"

	"github.com/go-gost/core/common/bufpool"
	mdns "github.com/miekg/dns"
)

const (
	defaultDomain  = "gost.local."
	defaultUDPSize = 4096
	readBufferSize = 65535

	edns0PayloadOption uint16 = mdns.EDNS0LOCALSTART
)

var (
	ErrInvalidPacket = errors.New("dns: invalid packet")
	ErrShortBuffer   = errors.New("dns: short buffer")
)

type options struct {
	domain  string
	udpSize uint16
}

// Option configures a DNS packet wrapper.
type Option func(*options)

// DomainOption sets the DNS question name used by wrapped packets.
func DomainOption(domain string) Option {
	return func(o *options) {
		if domain != "" {
			o.domain = mdns.Fqdn(domain)
		}
	}
}

// UDPSizeOption sets the EDNS0 UDP payload size advertised in wrapped packets.
func UDPSizeOption(size int) Option {
	return func(o *options) {
		if size > 0 && size <= readBufferSize {
			o.udpSize = uint16(size)
		}
	}
}

func parseOptions(opts ...Option) options {
	o := options{
		domain:  defaultDomain,
		udpSize: defaultUDPSize,
	}
	for _, opt := range opts {
		opt(&o)
	}
	return o
}

type clientConn struct {
	net.PacketConn
	id      uint32
	options options
}

// ClientConn wraps outbound payloads as DNS queries and unwraps DNS responses.
func ClientConn(conn net.PacketConn, opts ...Option) net.PacketConn {
	return &clientConn{
		PacketConn: conn,
		options:    parseOptions(opts...),
	}
}

func (c *clientConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
	buf := bufpool.Get(readBufferSize)
	defer bufpool.Put(buf)

	for {
		n, addr, err = c.PacketConn.ReadFrom(buf)
		if err != nil {
			return
		}

		m := mdns.Msg{}
		if err = m.Unpack(buf[:n]); err != nil {
			return 0, addr, err
		}
		if !m.Response {
			continue
		}

		payload, ok := payloadFromMessage(&m)
		if !ok {
			continue
		}
		if len(b) < len(payload) {
			return 0, addr, ErrShortBuffer
		}
		n = copy(b, payload)
		return
	}
}

func (c *clientConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
	id := uint16(atomic.AddUint32(&c.id, 1))
	wb, err := packQuery(id, c.options.domain, c.options.udpSize, b)
	if err != nil {
		return 0, err
	}
	_, err = c.PacketConn.WriteTo(wb, addr)
	if err != nil {
		return 0, err
	}
	return len(b), nil
}

type queryState struct {
	id       uint16
	question []mdns.Question
}

type serverConn struct {
	net.PacketConn
	mu      sync.RWMutex
	queries map[string]queryState
	options options
}

// ServerConn unwraps DNS queries as inbound payloads and wraps outbound
// payloads as DNS responses.
func ServerConn(conn net.PacketConn, opts ...Option) net.PacketConn {
	return &serverConn{
		PacketConn: conn,
		queries:    make(map[string]queryState),
		options:    parseOptions(opts...),
	}
}

func (c *serverConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
	buf := bufpool.Get(readBufferSize)
	defer bufpool.Put(buf)

	for {
		n, addr, err = c.PacketConn.ReadFrom(buf)
		if err != nil {
			return
		}

		m := mdns.Msg{}
		if err = m.Unpack(buf[:n]); err != nil {
			return 0, addr, err
		}
		if m.Response || !validQuestion(&m, c.options.domain) {
			continue
		}

		payload, ok := payloadFromMessage(&m)
		if !ok {
			continue
		}
		if len(b) < len(payload) {
			return 0, addr, ErrShortBuffer
		}

		c.storeQuery(addr, &m)
		n = copy(b, payload)
		return
	}
}

func (c *serverConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
	state := c.loadQuery(addr)
	wb, err := packResponse(state.id, state.question, c.options.domain, c.options.udpSize, b)
	if err != nil {
		return 0, err
	}
	_, err = c.PacketConn.WriteTo(wb, addr)
	if err != nil {
		return 0, err
	}
	return len(b), nil
}

func (c *serverConn) storeQuery(addr net.Addr, m *mdns.Msg) {
	if addr == nil {
		return
	}

	state := queryState{
		id:       m.Id,
		question: append([]mdns.Question(nil), m.Question...),
	}

	c.mu.Lock()
	c.queries[addr.String()] = state
	c.mu.Unlock()
}

func (c *serverConn) loadQuery(addr net.Addr) queryState {
	if addr == nil {
		return queryState{}
	}

	c.mu.RLock()
	state := c.queries[addr.String()]
	c.mu.RUnlock()
	return state
}

func packQuery(id uint16, domain string, udpSize uint16, payload []byte) ([]byte, error) {
	m := mdns.Msg{}
	m.SetQuestion(domain, mdns.TypeNULL)
	m.Id = id
	m.RecursionDesired = false
	addPayloadOpt(&m, udpSize, payload)
	return m.Pack()
}

func packResponse(id uint16, question []mdns.Question, domain string, udpSize uint16, payload []byte) ([]byte, error) {
	m := mdns.Msg{}
	m.Id = id
	m.Response = true
	m.Authoritative = true
	m.Rcode = mdns.RcodeSuccess
	if len(question) > 0 {
		m.Question = append([]mdns.Question(nil), question...)
	} else {
		m.SetQuestion(domain, mdns.TypeNULL)
		m.Id = id
		m.Response = true
		m.Authoritative = true
	}
	addPayloadOpt(&m, udpSize, payload)
	return m.Pack()
}

func addPayloadOpt(m *mdns.Msg, udpSize uint16, payload []byte) {
	opt := &mdns.OPT{
		Hdr: mdns.RR_Header{
			Name:   ".",
			Rrtype: mdns.TypeOPT,
			Class:  udpSize,
		},
		Option: []mdns.EDNS0{
			&mdns.EDNS0_LOCAL{
				Code: edns0PayloadOption,
				Data: append([]byte(nil), payload...),
			},
		},
	}
	m.Extra = append(m.Extra, opt)
}

func payloadFromMessage(m *mdns.Msg) ([]byte, bool) {
	for _, rr := range m.Extra {
		opt, ok := rr.(*mdns.OPT)
		if !ok {
			continue
		}
		for _, option := range opt.Option {
			local, ok := option.(*mdns.EDNS0_LOCAL)
			if !ok || local.Code != edns0PayloadOption {
				continue
			}
			return append([]byte(nil), local.Data...), true
		}
	}
	return nil, false
}

func validQuestion(m *mdns.Msg, domain string) bool {
	if len(m.Question) == 0 {
		return false
	}
	for _, q := range m.Question {
		if q.Qtype == mdns.TypeNULL && mdns.CanonicalName(q.Name) == mdns.CanonicalName(domain) {
			return true
		}
	}
	return false
}