m-dns/server/server.go

package server

import (
	"net"
	"os"
	"os/signal"
	"strings"
	"sync"
	"time"

	"git.giaever.org/joachimmg/go-log.git/log"
	"git.giaever.org/joachimmg/m-dns/config"
	"git.giaever.org/joachimmg/m-dns/errors"
	"git.giaever.org/joachimmg/m-dns/zone"
	"github.com/miekg/dns"
)

type MDnsServer struct {
	sync.Mutex
	zone.Zone
	*net.Interface

	ipv4 *net.UDPConn
	ipv6 *net.UDPConn

	running bool
	runCh   chan struct{}
}

func New(z zone.Zone, iface *net.Interface) (*MDnsServer, error) {
	m := new(MDnsServer)

	if m == nil {
		log.Traceln(errors.Server, errors.OutOfMemory)
		return nil, errors.OutOfMemory
	}

	if z == nil {
		log.Traceln(errors.Server, errors.ServerIsNil)
		return nil, errors.ServerIsNil
	}

	m.Zone = z
	m.Interface = iface

	var err error

	if m.ipv4, err = net.ListenMulticastUDP("udp4", m.Interface, config.MdnsIPv4Addr); err != nil {
		log.Infoln(errors.Server, config.MdnsIPv4, err)
	}

	if m.ipv6, err = net.ListenMulticastUDP("udp6", m.Interface, config.MdnsIPv6Addr); err != nil {
		log.Infoln(errors.Server, config.MdnsIPv6, err)
	}

	if m.ipv4 == nil && m.ipv6 == nil {
		log.Traceln(errors.Server, errors.ServerNoListenersStarted)
		return nil, errors.ServerNoListenersStarted
	}

	m.runCh = make(chan struct{})
	m.running = true

	go m.recv(m.ipv4)
	go m.recv(m.ipv6)

	return m, nil
}

func (m *MDnsServer) shutdownListener() {
	log.Traceln("Shutdown listener set on ctrl+x")
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		for range c {
			m.Shutdown()
		}
	}()
}

func (m *MDnsServer) Shutdown() {
	log.Traceln("Shutting down MDNS-server")
	m.Lock()
	defer m.Unlock()

	if !m.running {
		return
	}

	if m.ipv4 != nil {
		m.ipv4.Close()
	}

	if m.ipv6 != nil {
		m.ipv6.Close()
	}

	m.running = false
	close(m.runCh)
}

func (m *MDnsServer) Daemon() {
	log.Traceln("Daemon running.")
	go m.shutdownListener()
	<-m.runCh
	log.Traceln("Daemon ending.")
}

func (m *MDnsServer) recv(i *net.UDPConn) {
	if i == nil {
		return
	}

	buf := make([]byte, config.BufSize)

	for m.running {
		m.Lock()
		i.SetReadDeadline(time.Now().Add(config.BufReadDeadline * time.Second))
		_, addr, err := i.ReadFromUDP(buf)
		m.Unlock()

		if err != nil {
			if !strings.Contains(err.Error(), "i/o timeout") {
				log.Traceln(errors.Server, err)
			}
			continue
		}

		go m.handlePacket(buf, addr)

	}
}

func (m *MDnsServer) send(msg *dns.Msg, addr *net.UDPAddr) (int, error) {

	if msg == nil {
		return 0, nil
	}

	buf, err := msg.Pack()

	if err != nil {
		return 0, err
	}

	if addr.IP.To4() != nil {
		return m.ipv4.WriteToUDP(buf, addr)
	}

	return m.ipv6.WriteToUDP(buf, addr)
}

func (m *MDnsServer) handlePacket(p []byte, addr *net.UDPAddr) {
	msg := new(dns.Msg)

	if err := msg.Unpack(p); err != nil {
		log.Warningln(errors.Server, addr, err)
	}

	umsg, mmsg, err := m.handleMsg(msg)

	if err != nil {
		log.Warningln(errors.Server, addr, err)
		return
	}

	if n, err := m.send(umsg, addr); err != nil {
		log.Warningln(errors.Server, "Wrote", n, err)
	}

	if n, err := m.send(mmsg, addr); err != nil {
		log.Warningln(errors.Server, "Wrote", n, err)
	}
}

func (m *MDnsServer) handleMsg(msg *dns.Msg) (*dns.Msg, *dns.Msg, error) {
	if msg.Opcode != dns.OpcodeQuery {
		log.Traceln(errors.Server, errors.ServerReceivedNonQueryOpcode)
		return nil, nil, errors.ServerReceivedNonQueryOpcode
	}

	if msg.Rcode != 0 {
		log.Traceln(errors.Server, errors.ServerReceivedNonZeroRcode)
		return nil, nil, errors.ServerReceivedNonZeroRcode
	}

	if msg.Truncated {
		log.Traceln(errors.Server, errors.ServerReceivedTruncatedSet)
		return nil, nil, errors.ServerReceivedTruncatedSet
	}

	var uAnswer, mAnswer []dns.RR

	for _, q := range msg.Question {
		uRecords, mRecords := m.handleQuestion(q)
		uAnswer = append(uAnswer, uRecords...)
		mAnswer = append(mAnswer, mRecords...)
	}

	return m.handleResponse(msg, true, uAnswer), m.handleResponse(msg, false, mAnswer), nil
}

func (m *MDnsServer) handleQuestion(question dns.Question) ([]dns.RR, []dns.RR) {
	r := m.Records(question)

	if len(r) == 0 {
		return nil, nil
	}

	for i, rec := range r {
		log.Traceln(errors.Server, "Record", i, rec)
	}

	if question.Qclass&(1<<15) != 0 || config.ForceUnicast {
		return r, nil
	}

	return nil, r
}

func (m *MDnsServer) handleResponse(msg *dns.Msg, uni bool, ans []dns.RR) *dns.Msg {
	id := uint16(0)
	if uni {
		id = msg.Id
	}

	if len(ans) == 0 {
		return nil
	}

	return &dns.Msg{
		MsgHdr: dns.MsgHdr{
			Id:            id,
			Response:      true,
			Opcode:        dns.OpcodeQuery,
			Authoritative: true,
		},
		Compress: true,
		Answer:   ans,
	}
}