Merge pull request #48 from wolf-joe/linux_style_hosts

feat: support linux style hosts
wolf-joe · Mar 30, 2023 · a1cc733 · a1cc733
2 parents da94791 + 1146c17
commit a1cc733
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 4 deletions.
diff --git a/hosts/hosts.go b/hosts/hosts.go
@@ -3,14 +3,15 @@ package hosts
 import (
 	"bufio"
 	"fmt"
-	"github.com/miekg/dns"
-	"github.com/sirupsen/logrus"
-	"github.com/wolf-joe/ts-dns/config"
 	"net"
 	"os"
 	"regexp"
 	"strings"
 	"unicode"
+
+	"github.com/miekg/dns"
+	"github.com/sirupsen/logrus"
+	"github.com/wolf-joe/ts-dns/config"
 )
 
 // region interface
@@ -73,6 +74,13 @@ func NewDNSHosts(conf config.Conf) (IDNSHosts, error) {
 			if len(parts) < 2 {
 				continue
 			}
+			if ip := buildIPInfo(parts[0]); ip != zeroIP {
+				// linux style hosts file
+				for _, domain := range parts[1:] {
+					domainMap[domain] = ip
+				}
+				continue
+			}
 			if err = load(parts[0], parts[1]); err != nil {
 				return nil, fmt.Errorf("load hosts file %q error: %w", filename, err)
 			}

diff --git a/hosts/hosts_test.go b/hosts/hosts_test.go
@@ -1,11 +1,12 @@
 package hosts
 
 import (
+	"testing"
+
 	"github.com/miekg/dns"
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/wolf-joe/ts-dns/config"
-	"testing"
 )
 
 func buildReq(host string, qType uint16) *dns.Msg {
@@ -18,6 +19,41 @@ func buildReq(host string, qType uint16) *dns.Msg {
 	return msg
 }
 
+func TestLinuxHosts(t *testing.T) {
+	logrus.SetLevel(logrus.DebugLevel)
+	cfg := config.Conf{HostsFiles: []string{
+		"testdata/linux_style.txt",
+	}}
+	r, err := NewDNSHosts(cfg)
+	assert.Nil(t, err)
+	assert.NotNil(t, r)
+
+	cases := []struct {
+		host  string
+		query uint16
+		isNil bool
+		resp  string
+	}{
+		{"x.cn", dns.TypeA, true, ""},
+		{"z.cn", dns.TypeA, false, "z.cn.	0	IN	A	1.1.2.2"},
+		{"b.cn", dns.TypeA, false, "b.cn.	0	IN	A	1.1.3.3"},
+		{"b", dns.TypeA, false, "b.	0	IN	A	1.1.3.3"},
+		{"a.b.cn", dns.TypeA, true, ""},
+	}
+	for _, c := range cases {
+		t.Log(c)
+		resp := r.Get(buildReq(c.host, c.query))
+		assert.Nil(t, err)
+		if c.isNil {
+			assert.Nil(t, resp)
+		} else {
+			assert.NotNil(t, resp)
+			assert.Equal(t, 1, len(resp.Answer))
+			assert.Equal(t, c.resp, resp.Answer[0].String())
+		}
+	}
+}
+
 func TestNewHostReader(t *testing.T) {
 	logrus.SetLevel(logrus.DebugLevel)
 	cfg := config.Conf{Hosts: map[string]string{

diff --git a/hosts/testdata/linux_style.txt b/hosts/testdata/linux_style.txt
@@ -0,0 +1,3 @@
+# 1.1.1.1 x.cn
+1.1.2.2 z.cn
+1.1.3.3 b.cn b