MOD: Add hosts plugin

cmd/apexd/main.go

@@ -1,9 +1,11 @@
 package main
 
 import (
+	"github.com/oif/apex/pkg/config"
 	engine "github.com/oif/apex/pkg/engine/v1"
 	"github.com/oif/apex/plugin/cache"
 	"github.com/oif/apex/plugin/gdns"
+	"github.com/oif/apex/plugin/hosts"
 	"github.com/oif/apex/plugin/mupstream"
 	"github.com/oif/apex/plugin/statistics"
 
@@ -20,6 +22,15 @@ func main() {
 		plugin.ConfigFilePath = "statistics.toml"
 		return plugin
 	}())
+	s.RegisterPlugins(func() *hosts.Plugin {
+		plugin := new(hosts.Plugin)
+		plugin.Config = config.Hosts{
+			Enable:         true,
+			UpdateInterval: 3600,
+			Address:        "https://raw.githubusercontent.com/racaljk/hosts/master/hosts",
+		}
+		return plugin
+	}())
 	s.RegisterPlugins(func() *cache.Plugin {
 		plugin := new(cache.Plugin)
 		plugin.CacheSize = 1024

plugin/hosts/hosts.go

@@ -0,0 +1,45 @@
+package hosts
+
+import (
+	"bufio"
+	"net"
+	"net/http"
+	"strings"
+	"time"
+)
+
+func parseHosts(url string) (map[string]net.IP, error) {
+	client := http.Client{
+		Timeout: 5 * time.Second,
+	}
+	resp, err := client.Get(url)
+	if err != nil {
+		return nil, err
+	}
+
+	hosts := make(map[string]net.IP)
+
+	lines := bufio.NewReader(resp.Body)
+	for {
+		line, err := lines.ReadString('\n')
+		if err != nil {
+			break
+		}
+		if len(line) < 2 || line[0] == 35 || line[0] == 32 { // if first character is '#' or space
+			continue
+		}
+		if strings.ContainsAny(line, ":") { // ignore ipv6
+			continue
+		}
+		//  split line
+		var pair []string
+		pair = strings.SplitN(line, "\t", 2)
+		if pair[0] == "" || pair[1] == "" {
+			continue
+		}
+		pair[1] = strings.TrimRight(pair[1], "\n")
+		pair[1] = strings.TrimSpace(pair[1])
+		hosts[pair[1]+"."] = net.ParseIP(pair[0])
+	}
+	return hosts, nil
+}

plugin/hosts/hosts_test.go

@@ -0,0 +1,10 @@
+package hosts
+
+import "testing"
+
+func TestLoadHosts(t *testing.T) {
+	_, err := parseHosts("https://raw.githubusercontent.com/racaljk/hosts/master/hosts")
+	if err != nil {
+		t.Fatal(err)
+	}
+}

plugin/hosts/plugin.go

@@ -0,0 +1,127 @@
+package hosts
+
+import (
+	"log"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/oif/apex/pkg/config"
+	plugin "github.com/oif/apex/pkg/plugin/v1"
+
+	"github.com/Sirupsen/logrus"
+	"github.com/miekg/dns"
+)
+
+// PluginName for plugin
+const PluginName = "Hosts Plugin"
+
+// Plugin implements pkg/plugin/v1
+type Plugin struct {
+	Config config.Hosts
+	hosts  map[string]net.IP
+	lock   sync.RWMutex
+}
+
+// Name return the name of this plugin
+func (p *Plugin) Name() string {
+	return PluginName
+}
+
+// Initialize Plugin
+func (p *Plugin) Initialize() error {
+	if p.Config.Enable {
+		p.Config.Check() // check config
+
+		go func() {
+			for {
+				startAt := time.Now().UnixNano()
+				if err := p.hostsUpdater(); err != nil {
+					logrus.WithFields(logrus.Fields{
+						"err": err,
+					}).Error("Update hosts failed")
+				} else {
+					logrus.WithFields(logrus.Fields{
+						"time_usage":  time.Now().UnixNano() - startAt,
+						"hosts_count": len(p.hosts),
+					}).Debug("Update hosts success")
+				}
+				time.Sleep(time.Duration(p.Config.UpdateInterval) * time.Second)
+			}
+		}()
+	}
+	return nil
+}
+
+func (p *Plugin) hostsUpdater() error {
+	newHosts, err := parseHosts(p.Config.Address)
+	if err != nil {
+		return err
+	}
+
+	// update hosts
+	p.lock.Lock()
+	p.hosts = newHosts
+	p.lock.Unlock()
+
+	return nil
+}
+
+func (p *Plugin) getHosts(name string) (net.IP, bool) {
+	p.lock.RLock()
+	ip, ok := p.hosts[name]
+	p.lock.RUnlock()
+	return ip, ok
+}
+
+func (p *Plugin) Warmup(c *plugin.Context)                   {}
+func (p *Plugin) AfterResponse(c *plugin.Context, err error) {}
+
+func (p *Plugin) Patch(c *plugin.Context) {
+	// for i, question := range c.Msg.Question {
+	// 	if question.Qtype != dns.TypeA { // only hosts type A
+	// 		continue
+	// 	}
+	// 	if c.Msg.Answer == nil {
+	// 		log.Print("nil")
+	// 	}
+	// 	if ip, ok := p.getHosts(question.Name); ok { // match hosts
+	// 		c.Msg.Answer[i] = &dns.A{
+	// 			Hdr: dns.RR_Header{
+	// 				Name:   question.Name,
+	// 				Rrtype: dns.TypeA,
+	// 				Ttl:    300,
+	// 			},
+	// 			A: ip,
+	// 		}
+	// 		c.Abort()
+	// 		logrus.WithFields(logrus.Fields{
+	// 			"name": question.Name,
+	// 		}).Debug("Match hosts")
+	// 	}
+	// }
+
+	if c.Msg.Question[0].Qtype != dns.TypeA { // only hosts type A
+		return
+	}
+	if c.Msg.Answer == nil {
+		log.Print("nil")
+	}
+	if ip, ok := p.getHosts(c.Msg.Question[0].Name); ok { // match hosts
+		c.Msg.Answer = []dns.RR{
+			&dns.A{
+				Hdr: dns.RR_Header{
+					Name:   c.Msg.Question[0].Name,
+					Rrtype: dns.TypeA,
+					Class:  c.Msg.Question[0].Qclass,
+					Ttl:    300,
+				},
+				A: ip,
+			},
+		}
+		c.Abort()
+		logrus.WithFields(logrus.Fields{
+			"name": c.Msg.Question[0].Name,
+		}).Debug("Match hosts")
+	}
+}