Merge pull request #46 from wolf-joe/ipv6_ipset

support ipv6 ipset

config/config.go

@@ -43,6 +43,7 @@ type Group struct {
 	TCPPingPort int  `toml:"tcp_ping_port"`
 
 	IPSet    string `toml:"ipset"`
+	IPSet6   string `toml:"ipset6"`
 	IPSetTTL int    `toml:"ipset_ttl"`
 
 	Redirector string `toml:"redirector"`

outbound/groups.go

@@ -158,12 +158,19 @@ func BuildGroups(globalConf config.Conf) (map[string]IGroup, error) {
 		}
 		g.callers = callers
 		// ipset
-		if conf.IPSet != "" {
-			is, err := ipset.New(conf.IPSet, "hash:ip", &ipset.Params{Timeout: conf.IPSetTTL})
+		if name := conf.IPSet; name != "" {
+			is, err := ipset.New(name, "hash:ip", &ipset.Params{Timeout: conf.IPSetTTL})
 			if err != nil {
-				return nil, fmt.Errorf("build ipset %q failed: %w", conf.IPSet, err)
+				return nil, fmt.Errorf("build ipset %q failed: %w", name, err)
 			}
-			g.ipSet = is
+			g.ipSet = ipSetWrapper{is}
+		}
+		if name := conf.IPSet6; name != "" {
+			is, err := ipset.New(name, "hash:ip", &ipset.Params{Timeout: conf.IPSetTTL, HashFamily: "inet6"})
+			if err != nil {
+				return nil, fmt.Errorf("build ipset %q failed: %w", name, err)
+			}
+			g.ipSet6 = ipSetWrapper{is}
 		}
 		groups[name] = g
 	}
@@ -194,7 +201,8 @@ type groupImpl struct {
 	fastestIP   bool // 是否对响应中的IP地址进行测速，找出ping值最低的IP地址
 	tcpPingPort int  // 是否使用tcp ping
 
-	ipSet *ipset.IPSet // 将响应中的IPv4地址加入ipset
+	ipSet  iIPSet // 将响应中的IPv4地址加入ipset
+	ipSet6 iIPSet // 将响应中的IPv4地址加入ipset
 
 	stopCh  chan struct{}
 	stopped chan struct{}
@@ -362,14 +370,23 @@ doPing:
 }
 
 func (g *groupImpl) PostProcess(_ *dns.Msg, resp *dns.Msg) {
-	if resp == nil || g.ipSet == nil {
+	if resp == nil {
 		return
 	}
 	for _, answer := range resp.Answer {
-		if a, ok := answer.(*dns.A); ok {
-			if err := g.ipSet.Add(a.A.String(), g.ipSet.Timeout); err != nil {
-				logrus.Warnf("add %s to ipset<%s> failed: %+v", a.A, g.ipSet.Name, err)
-			}
+		var val string
+		var target iIPSet
+		switch tv := answer.(type) {
+		case *dns.A:
+			val, target = tv.A.String(), g.ipSet
+		case *dns.AAAA:
+			val, target = tv.AAAA.String(), g.ipSet6
+		}
+		if val == "" || target == nil {
+			continue
+		}
+		if err := target.Add(val, target.GetTimeout()); err != nil {
+			logrus.Warnf("add %s to ipset<%s> failed: %+v", val, target.GetName(), err)
 		}
 	}
 }

outbound/groups_test.go

@@ -65,3 +65,28 @@ func TestDisableIPv6(t *testing.T) {
 	})
 	assert.Nil(t, resp)
 }
+
+func TestPostProcess(t *testing.T) {
+	var v4val, v6val string
+	group := &groupImpl{
+		ipSet: MockIPSet{
+			Name:    "",
+			Timeout: 0,
+			MockAdd: func(val string, _ int) error { v4val = val; return nil },
+		},
+		ipSet6: MockIPSet{
+			Name:    "",
+			Timeout: 0,
+			MockAdd: func(val string, _ int) error { v6val = val; return nil },
+		},
+	}
+	rr, err := dns.NewRR("z.cn 0 IN A 1.1.1.1")
+	assert.Nil(t, err)
+	group.PostProcess(nil, &dns.Msg{Answer: []dns.RR{rr}})
+	assert.Equal(t, "1.1.1.1", v4val)
+
+	rr, err = dns.NewRR("z.cn 0 IN AAAA ff80::1")
+	assert.Nil(t, err)
+	group.PostProcess(nil, &dns.Msg{Answer: []dns.RR{rr}})
+	assert.Equal(t, "ff80::1", v6val)
+}

outbound/ipset.go

@@ -0,0 +1,31 @@
+package outbound
+
+import "github.com/wolf-joe/go-ipset/ipset"
+
+type iIPSet interface {
+	Add(entry string, timeout int) error
+	GetName() string
+	GetTimeout() int
+}
+
+type ipSetWrapper struct {
+	*ipset.IPSet
+}
+
+func (i ipSetWrapper) GetName() string { return i.Name }
+func (i ipSetWrapper) GetTimeout() int { return i.Timeout }
+
+var (
+	_ iIPSet = ipSetWrapper{}
+	_ iIPSet = MockIPSet{}
+)
+
+type MockIPSet struct {
+	Name    string
+	Timeout int
+	MockAdd func(entry string, timeout int) error
+}
+
+func (i MockIPSet) GetName() string                     { return i.Name }
+func (i MockIPSet) GetTimeout() int                     { return i.Timeout }
+func (i MockIPSet) Add(entry string, timeout int) error { return i.MockAdd(entry, timeout) }

ts-dns-full.toml

@@ -44,6 +44,7 @@ max_ttl = 86400  # 最大ttl，单位为秒
 
   # 警告：进程启动时会覆盖已有同名IPSet
   ipset = "blocked"  # 目标IPSet名称，该组所有域名的ipv4解析结果将加入到该IPSet中
+  ipset6 = "blocked6"  # 目标IPSet名称，该组所有域名的ipv6解析结果将加入到该IPSet中
   ipset_ttl = 86400 # ipset记录超时时间，单位为秒，推荐设置以避免ipset记录过多
 
   # 可选自定义分组，用于其它情况