From 2c61900dd5efd81a5351513a20fa65580c8a6616 Mon Sep 17 00:00:00 2001
From: radhitya <alif@radhitya.org>
Date: Tue, 23 Jun 2026 05:36:25 +0700
Subject: make if error for unsupported query type

---
 internal/resolver/resolver.go      |  8 ++++++--
 internal/resolver/resolver_test.go | 21 +++++++++++++++++----
 internal/server/handler.go         | 32 ++++++++++++++++++++++++++------
 internal/server/server_test.go     | 23 ++++++++++++++++++++---
 4 files changed, 69 insertions(+), 15 deletions(-)

(limited to 'internal')

diff --git a/internal/resolver/resolver.go b/internal/resolver/resolver.go
index c7a0694..3cdf846 100644
--- a/internal/resolver/resolver.go
+++ b/internal/resolver/resolver.go
@@ -11,8 +11,9 @@ import (
 )
 
 var (
-	ErrMaxDelegations = errors.New("max delegations exceeded")
-	ErrNoServers      = errors.New("no nameservers available")
+	ErrMaxDelegations  = errors.New("max delegations exceeded")
+	ErrNoServers       = errors.New("no nameservers available")
+	ErrUnsupportedType = errors.New("unsupported query type")
 )
 
 type Resolver struct {
@@ -251,6 +252,9 @@ func (r *Resolver) exchangeWithRetries(ctx context.Context, servers []string,
 	qname string, qtype uint16, rd bool) (*dns.Msg, error) {
 
 	msg := dns.NewMsg(qname, qtype)
+	if msg == nil {
+		return nil, fmt.Errorf("%w %d for %s", ErrUnsupportedType, qtype, qname)
+	}
 	msg.UDPSize = 4096
 	msg.RecursionDesired = rd
 
diff --git a/internal/resolver/resolver_test.go b/internal/resolver/resolver_test.go
index daa8a98..3c44cb4 100644
--- a/internal/resolver/resolver_test.go
+++ b/internal/resolver/resolver_test.go
@@ -1,17 +1,30 @@
 package resolver
 
 import (
+	"codeberg.org/miekg/dns"
+	"codeberg.org/miekg/dns/dnsutil"
+	"codeberg.org/miekg/dns/rdata"
 	"context"
+	"errors"
 	"io"
 	"net/netip"
 	"testing"
 	"time"
-
-	"codeberg.org/miekg/dns"
-	"codeberg.org/miekg/dns/dnsutil"
-	"codeberg.org/miekg/dns/rdata"
 )
 
+func TestLookupUnknownType(t *testing.T) {
+	r := New(WithForwarders([]string{"127.0.0.1:1"}), WithTimeout(50*time.Millisecond))
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	_, err := r.Lookup(ctx, "example.com.", 0xaa58)
+	if err == nil {
+		t.Fatal("expected error for unsupported qtype")
+	}
+	if !errors.Is(err, ErrUnsupportedType) {
+		t.Fatalf("expected ErrUnsupportedType, got %v", err)
+	}
+}
 func startTestServer(t *testing.T, addr string, handler dns.Handler) *dns.Server {
 	t.Helper()
 	ready := make(chan struct{})
diff --git a/internal/server/handler.go b/internal/server/handler.go
index 4516468..46699d6 100644
--- a/internal/server/handler.go
+++ b/internal/server/handler.go
@@ -1,20 +1,36 @@
 package server
 
 import (
+	"codeberg.org/miekg/dns"
+	"codeberg.org/miekg/dns/rdata"
 	"context"
+	"errors"
 	"io"
+	"linum/internal/blocklist"
+	"linum/internal/cache"
+	"linum/internal/resolver"
 	"log/slog"
 	"net"
 	"net/netip"
+	"runtime/debug"
 	"time"
-
-	"codeberg.org/miekg/dns"
-	"codeberg.org/miekg/dns/rdata"
-	"linum/internal/blocklist"
-	"linum/internal/cache"
 )
 
 func (s *Server) handleQuery(ctx context.Context, w dns.ResponseWriter, req *dns.Msg) {
+	defer func() {
+		if r := recover(); r != nil {
+			slog.Error("panic in query handler", "recover", r, "stack", string(debug.Stack()))
+			m := new(dns.Msg)
+			m.Response = true
+			m.Rcode = dns.RcodeServerFailure
+			if req != nil {
+				m.ID = req.ID
+				m.Question = req.Question
+			}
+			io.Copy(w, m)
+		}
+	}()
+
 	clientIP := parseClientIP(w.RemoteAddr())
 	if !s.isAllowed(clientIP) {
 		slog.Warn("query denied by ACL", "client", clientIP)
@@ -92,13 +108,17 @@ func (s *Server) buildResponse(req *dns.Msg) (*dns.Msg, bool) {
 
 	reply, err := s.resolver.Lookup(ctx, qname, qtype)
 	if err != nil {
+		rcode := dns.RcodeServerFailure
+		if errors.Is(err, resolver.ErrUnsupportedType) {
+			rcode = dns.RcodeNotImplemented
+		}
 		slog.Error("resolution failed",
 			"err", err,
 			"qname", qname,
 			"qtype", dns.TypeToString[qtype],
 		)
 		m := new(dns.Msg)
-		m.Rcode = dns.RcodeServerFailure
+		m.Rcode = uint16(rcode)
 		m.Response = true
 		m.ID = req.ID
 		m.Question = req.Question
diff --git a/internal/server/server_test.go b/internal/server/server_test.go
index 2340acd..93e2a69 100644
--- a/internal/server/server_test.go
+++ b/internal/server/server_test.go
@@ -1,15 +1,32 @@
 package server
 
 import (
+	"codeberg.org/miekg/dns"
+	"codeberg.org/miekg/dns/rdata"
 	"context"
+	"linum/internal/resolver"
 	"log/slog"
 	"testing"
 	"time"
-
-	"codeberg.org/miekg/dns"
-	"linum/internal/resolver"
 )
 
+func TestBuildResponseUnknownType(t *testing.T) {
+	s := testServer(t)
+	m := new(dns.Msg)
+	m.Question = []dns.RR{&dns.RFC3597{
+		Hdr:     dns.Header{Name: "example.com.", Class: dns.ClassINET},
+		RFC3597: rdata.RFC3597{RRType: 0xaa58},
+	}}
+
+	resp, _ := s.buildResponse(m)
+	if resp == nil {
+		t.Fatal("buildResponse returned nil")
+	}
+	if resp.Rcode != dns.RcodeNotImplemented {
+		t.Errorf("expected NOTIMPL, got %d", resp.Rcode)
+	}
+}
+
 func testServer(t *testing.T) *Server {
 	t.Helper()
 	r := resolver.New(
-- 
cgit v1.2.3