feat: add dns black list support (#11)

* feat: add dns black list support

* dns server refactor

* fix the api calling

---------

Co-authored-by: Rick <[email protected]>

Author: LinuxSuRen

README.md

@@ -16,6 +16,7 @@ atest-collector dns --simple-config config.yaml
 
 Below is an example of a simple DNS config:
 ```yaml
-atest.com: 127.0.0.1
-www.atest.com: 127.0.0.1
+simple:
+    atest.com: 127.0.0.1
+    www.atest.com: 127.0.0.1
 ```

cmd/dns.go

@@ -1,5 +1,5 @@
 /*
-Copyright 2024 LinuxSuRen.
+Copyright 2024-2025 LinuxSuRen.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -18,16 +18,9 @@ package cmd
 
 import (
 	"fmt"
-	"github.com/google/gopacket"
-	"github.com/google/gopacket/layers"
 	adns "github.com/linuxsuren/atest-ext-collector/pkg/dns"
-	"github.com/miekg/dns"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
-	"gopkg.in/yaml.v3"
-	"net"
-	"os"
-	"strings"
 )
 
 func createDNSCmd() (cmd *cobra.Command) {
@@ -51,6 +44,8 @@ type dnsOptions struct {
 
 	// inner fields
 	cacheHandler adns.DNSCache
+	server       adns.Server
+	config       *adns.DNSConfig
 }
 
 func (o *dnsOptions) setFlags(flags *pflag.FlagSet) {
@@ -68,107 +63,34 @@ func (o *dnsOptions) preRunE(_ *cobra.Command, _ []string) (err error) {
 	}
 
 	if o.simpleConfig != "" {
-		var data []byte
-		if data, err = os.ReadFile(o.simpleConfig); err == nil {
-			records := make(map[string]string)
-			if err = yaml.Unmarshal(data, &records); err == nil {
-				o.cacheHandler.Init(records)
-			}
+		if o.config, err = adns.ParseFromFile(o.simpleConfig); err != nil {
+			return
 		}
-	}
-	return
-}
 
-func (o *dnsOptions) runE(cmd *cobra.Command, args []string) (err error) {
-	addr := net.UDPAddr{
-		Port: o.port,
-		IP:   net.ParseIP("0.0.0.0"),
+		o.cacheHandler.Init(o.config.Simple)
+		o.cacheHandler.GetWildcardCache().Init(o.config.Wildcard)
 	}
-	var u *net.UDPConn
-	if u, err = net.ListenUDP("udp", &addr); err != nil {
-		return
+	if o.config == nil {
+		o.config = &adns.DNSConfig{}
 	}
 
-	cmd.Println("DNS server is ready!")
-	go func() {
-		if err := adns.NewHTTPServer(o.httpPort, o.cacheHandler).Start(); err != nil {
-			panic(err)
-		}
-	}()
-	// Wait to get request on that port
-	for {
-		tmp := make([]byte, 1024)
-		_, addr, _ := u.ReadFrom(tmp)
-		clientAddr := addr
-		packet := gopacket.NewPacket(tmp, layers.LayerTypeDNS, gopacket.Default)
-		dnsPacket := packet.Layer(layers.LayerTypeDNS)
-		tcp, _ := dnsPacket.(*layers.DNS)
-		if !o.serveDNS(u, clientAddr, tcp) {
-			o.cacheDNS(string(tcp.Questions[0].Name))
-		}
+	if o.upstream != "" {
+		o.config.Upstream = o.upstream
+	}
+	if o.port > 0 {
+		o.config.Port = o.port
 	}
+	o.server = adns.NewDNSServer(o.config, o.cacheHandler)
 	return
 }
 
-func (o *dnsOptions) cacheDNS(name string) {
-	client := dns.Client{}
-	var m dns.Msg
-	m.SetQuestion(name+".", dns.TypeA)
-	reply, _, err := client.Exchange(&m, o.upstream)
-	if err != nil {
-		fmt.Println("failed query domain", name, err)
-		return
-	}
-	if reply.Rcode == dns.RcodeSuccess && len(reply.Answer) > 0 {
-		if a, ok := reply.Answer[0].(*dns.A); ok {
-			o.cacheHandler.Put(strings.TrimSuffix(reply.Question[0].Name, "."), a.A.String())
-			fmt.Println("cache new record", reply.Question[0].Name, "==", a.A.String())
-			fmt.Println("total cache item count", o.cacheHandler.Size())
-			return
-		} else if cname, ok := reply.Answer[0].(*dns.CNAME); ok {
-			fmt.Println(name, "==", strings.TrimSuffix(cname.Hdr.Name, "."), "==", strings.TrimSuffix(cname.Target, "."))
-			if ip := o.cacheHandler.LookupIP(strings.TrimSuffix(cname.Target, ".")); ip != "" {
-				o.cacheHandler.Put(strings.TrimSuffix(cname.Hdr.Name, "."), ip)
-			} else {
-				o.cacheDNS(strings.TrimSuffix(cname.Target, "."))
-			}
+func (o *dnsOptions) runE(cmd *cobra.Command, args []string) (err error) {
+	go func() {
+		if err := adns.NewHTTPServer(o.httpPort, o.upstream, o.cacheHandler).Start(); err != nil {
+			panic(err)
 		}
-		fmt.Println("unknown record", reply.Question[0].Name, ",", reply.Answer[0].String())
-	}
-}
+	}()
 
-func (o *dnsOptions) serveDNS(u *net.UDPConn, clientAddr net.Addr, request *layers.DNS) (resolved bool) {
-	replyMess := request
-	var dnsAnswer layers.DNSResourceRecord
-	dnsAnswer.Type = layers.DNSTypeA
-	var ip string
-	var err error
-	resolved = true
-	ip = o.cacheHandler.LookupIP(string(request.Questions[0].Name))
-	if ip == "" {
-		//fmt.Printf("cannot found: %s\n", request.Questions[0].Name)
-		resolved = false
-		return
-		//Todo: Log no data present for the IP and handle:todo
-	}
-	a, _, _ := net.ParseCIDR(ip + "/24")
-	dnsAnswer.Type = layers.DNSTypeA
-	dnsAnswer.IP = a
-	dnsAnswer.Name = []byte(request.Questions[0].Name)
-	fmt.Println(string(request.Questions[0].Name), "===", ip)
-	dnsAnswer.Class = layers.DNSClassIN
-	replyMess.QR = true
-	replyMess.ANCount = 1
-	replyMess.OpCode = layers.DNSOpCodeNotify
-	replyMess.AA = true
-	replyMess.Answers = append(replyMess.Answers, dnsAnswer)
-	replyMess.ResponseCode = layers.DNSResponseCodeNoErr
-	buf := gopacket.NewSerializeBuffer()
-	opts := gopacket.SerializeOptions{} // See SerializeOptions for more details.
-	err = replyMess.SerializeTo(buf, opts)
-	if err != nil {
-		panic(err)
-	}
-	u.WriteTo(buf.Bytes(), clientAddr)
+	err = o.server.Start()
 	return
 }

pkg/dns/config.go

@@ -0,0 +1,45 @@
+/*
+Copyright 2025 LinuxSuRen.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package dns
+
+import (
+	"gopkg.in/yaml.v3"
+	"os"
+)
+
+type DNSConfig struct {
+	Wildcard      map[string]string `yaml:"wildcard"`
+	Simple        map[string]string `yaml:"simple"`
+	Black         []string          `yaml:"black"`
+	WildcardBlack []string          `yaml:"wildcard_black"`
+	Upstream      string            `yaml:"upstream"`
+	Port          int               `yaml:"port"`
+}
+
+func ParseFromFile(file string) (config *DNSConfig, err error) {
+	var data []byte
+	if data, err = os.ReadFile(file); err == nil {
+		config, err = ParseFromBuffer(data)
+	}
+	return
+}
+
+func ParseFromBuffer(buffer []byte) (config *DNSConfig, err error) {
+	config = &DNSConfig{}
+	err = yaml.Unmarshal(buffer, &config)
+	return
+}

pkg/dns/config_test.go

@@ -0,0 +1,52 @@
+/*
+Copyright 2025 LinuxSuRen.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package dns_test
+
+import (
+	"github.com/linuxsuren/atest-ext-collector/pkg/dns"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+import _ "embed"
+
+func TestParse(t *testing.T) {
+	defaultConfig := &dns.DNSConfig{
+		Wildcard: map[string]string{
+			"*.def.com": "0.0.0.2",
+		},
+		Simple: map[string]string{
+			"abc.com": "0.0.0.1",
+		},
+		Black:         []string{"ghi.com"},
+		WildcardBlack: []string{"*.jkl.com"},
+		Upstream:      "8.8.8.8",
+		Port:          53,
+	}
+
+	config, err := dns.ParseFromBuffer(configYaml)
+	assert.NoError(t, err)
+	assert.Equal(t, defaultConfig, config)
+
+	config, err = dns.ParseFromFile("testdata/config.yaml")
+	assert.NoError(t, err)
+	assert.Equal(t, defaultConfig, config)
+}
+
+var (
+	//go:embed testdata/config.yaml
+	configYaml []byte
+)

pkg/dns/data/index.html

@@ -5,6 +5,9 @@
         function showAddForm() {
             document.getElementById('addFormZone').classList.remove('hide');
         }
+        function showAddBlackDomain() {
+            document.getElementById('addBlackDomain').classList.remove('hide');
+        }
     </script>
     <style>
         .hide {
@@ -15,13 +18,25 @@
 <body>
 
 <div>
-    <button onclick="showAddForm()">Add a new record</button> There are {{.size}} records.
-    <div class="hide" id="addFormZone">
-        <form action="/add" method="post">
-            Domain: <input name="domain"/>
-            IP: <input name="ip"/>
-            <button type="submit">Submit</button>
-        </form>
+    <div>Upstream DNS Server: {{.upstream}}</div>
+    <div>
+        <button onclick="showAddForm()">Add a new record</button> There are {{.size}} records.
+        <div class="hide" id="addFormZone">
+            <form action="/add" method="post">
+                Domain: <input name="domain"/>
+                IP: <input name="ip"/>
+                <button type="submit">Submit</button>
+            </form>
+        </div>
+    </div>
+    <div>
+        <button onclick="showAddBlackDomain()">Add black domain</button>
+        <div class="hide" id="addBlackDomain">
+            <form action="/addBlack" method="post">
+                Domain: <input name="domain"/>
+                <button type="submit">Submit</button>
+            </form>
+        </div>
     </div>
 </div>
 
@@ -39,5 +54,19 @@
     </tr>
     {{end}}
 </table>
+
+<div>Black list</div>
+<table border="1">
+    <tr>
+        <th>Domain</th>
+        <th>Operation</th>
+    </tr>
+    {{range .black}}
+    <tr>
+        <td>{{ . }}</td>
+        <td><a href="/removeBlack?domain={{.}}">Remove</a></td>
+    </tr>
+    {{end}}
+</table>
 </body>
 </html>