1 files changed, 104 insertions, 0 deletions
diff --git a/internal/server/server.go b/internal/server/server.go
new file mode 100644
index 0000000..01bda09
--- /dev/null
+++ b/internal/server/server.go
@@ -0,0 +1,104 @@
+package server
+
+import (
+	"context"
+	"log/slog"
+	"time"
+
+	"github.com/miekg/dns"
+)
+
+type Server struct {
+	addr string
+	logger *slog.Logger
+	inner *dns.Server
+}
+
+func New(addr string, logger *slog.Logger) (*Server, error) {
+	mux := dns.NewServeMux()
+	mux.HandleFunc(".", handleQuery)
+
+	inner := &dns.Server{
+		Addr: addr,
+		Net: "udp",
+		Handler: mux,
+		UDPSize: 4096,
+		ReadTimeout: 5 * time.Second,
+		WriteTimeout: 5 * time.Second,
+	}
+	return &Server{
+		addr: addr,
+		logger: logger,
+		inner: inner,
+	}, nil
+}
+
+func (s *Server) Run(ctx context.Context) error {
+	errCh := make(chan error, 1)
+	go func() {
+		s.logger.Info("udp listener active", "addr", s.addr)
+		errCh <- s.inner.ListenAndServe()
+	}()
+
+	select {
+	case <-ctx.Done():
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5 *time.Second)
+		defer cancel()
+		if err := s.inner.ShutdownContext(shutdownCtx); err != nil {
+			s.logger.Error("graceful shutdown failed", "err", err)
+			return err
+		}
+		return ctx.Err()
+	case err := <-errCh:
+		return err
+	}
+}
+
+func (s *Server) Close() error {
+	return s.inner.Shutdown()
+}
+
+func handleQuery(w dns.ResponseWriter, req *dns.Msg) {
+	if len(req.Question) == 0 {
+		m := new(dns.Msg)
+		m.SetRcode(req, dns.RcodeFormatError)
+		_ = w.WriteMsg(m)
+		return
+	}
+
+	q := req.Question[0]
+	resp := new(dns.Msg)
+	resp.SetReply(req)
+	resp.Authoritative = false
+
+	if q.Name == "example.com." && q.Qtype == dns.TypeA {
+		resp.Answer = []dns.RR{
+			&dns.A{
+				Hdr: dns.RR_Header{
+					Name: q.Name,
+					Rrtype: dns.TypeA,
+					Class: dns.ClassINET,
+					Ttl: 60,
+				},
+				A: []byte{127,0,0,1},
+			},
+		}
+	} else {
+		resp.Rcode = dns.RcodeNameError
+	}
+
+	if err := w.WriteMsg(resp); err != nil {
+		slog.Error("write response failed",
+		"err", err,
+		"qname", q.Name,
+		"qtype", dns.TypeToString[q.Qtype],
+	)
+	return
+}
+slog.Info("query served",
+"qname", q.Name,
+"qtype", dns.TypeToString[q.Qtype],
+"rcode", dns.RcodeToString[resp.Rcode],
+"client", w.RemoteAddr().String(),
+)
+}