update hosts

pkg/chain/router.go

@@ -15,7 +15,7 @@ import (
 type Router struct {
 	Retries  int
 	Chain    *Chain
-	Hosts    *hosts.Hosts
+	Hosts    hosts.HostMapper
 	Resolver resolver.Resolver
 	Logger   logger.Logger
 }
@@ -78,9 +78,11 @@ func (r *Router) resolve(ctx context.Context, addr string) (string, error) {
 		return "", err
 	}
 
-	if ip := r.Hosts.Lookup(host); ip != nil {
-		r.Logger.Debugf("hit hosts: %s -> %s", host, ip)
-		return net.JoinHostPort(ip.String(), port), nil
+	if r.Hosts != nil {
+		if ip := r.Hosts.Lookup(host); ip != nil {
+			r.Logger.Debugf("hit hosts: %s -> %s", host, ip)
+			return net.JoinHostPort(ip.String(), port), nil
+		}
 	}
 
 	if r.Resolver != nil {

pkg/handler/http/handler.go

@@ -139,22 +139,6 @@ func (h *httpHandler) handleRequest(ctx context.Context, conn net.Conn, req *htt
 		resp.Header = http.Header{}
 	}
 
-	/*
-		if !Can("tcp", host, h.options.Whitelist, h.options.Blacklist) {
-			log.Logf("[http] %s - %s : Unauthorized to tcp connect to %s",
-				conn.RemoteAddr(), conn.LocalAddr(), host)
-			resp.StatusCode = http.StatusForbidden
-
-			if Debug {
-				dump, _ := httputil.DumpResponse(resp, false)
-				log.Logf("[http] %s <- %s\n%s", conn.RemoteAddr(), conn.LocalAddr(), string(dump))
-			}
-
-			resp.Write(conn)
-			return
-		}
-	*/
-
 	if h.bypass != nil && h.bypass.Contains(addr) {
 		resp.StatusCode = http.StatusForbidden
 

pkg/hosts/hosts.go

@@ -4,34 +4,37 @@ import (
 	"net"
 )
 
-// Host is a static mapping from hostname to IP.
-type Host struct {
+// HostMapper is a mapping from hostname to IP.
+type HostMapper interface {
+	Lookup(host string) net.IP
+}
+
+type host struct {
 	IP       net.IP
 	Hostname string
 	Aliases  []string
 }
 
-// NewHost creates a Host.
-func NewHost(ip net.IP, hostname string, aliases ...string) Host {
-	return Host{
-		IP:       ip,
-		Hostname: hostname,
-		Aliases:  aliases,
-	}
-}
-
 // Hosts is a static table lookup for hostnames.
 // For each host a single line should be present with the following information:
 // IP_address canonical_hostname [aliases...]
 // Fields of the entry are separated by any number of blanks and/or tab characters.
 // Text from a "#" character until the end of the line is a comment, and is ignored.
 type Hosts struct {
-	hosts []Host
+	mappings []host
 }
 
-// AddHost adds host(s) to the host table.
-func (h *Hosts) AddHost(host ...Host) {
-	h.hosts = append(h.hosts, host...)
+func NewHosts() *Hosts {
+	return &Hosts{}
+}
+
+// Map maps ip to hostname or aliases.
+func (h *Hosts) Map(ip net.IP, hostname string, aliases ...string) {
+	h.mappings = append(h.mappings, host{
+		IP:       ip,
+		Hostname: hostname,
+		Aliases:  aliases,
+	})
 }
 
 // Lookup searches the IP address corresponds to the given host from the host table.
@@ -40,7 +43,7 @@ func (h *Hosts) Lookup(host string) (ip net.IP) {
 		return
 	}
 
-	for _, h := range h.hosts {
+	for _, h := range h.mappings {
 		if h.Hostname == host {
 			ip = h.IP
 			break

pkg/internal/util/resolver/cache.go

@@ -51,22 +51,12 @@ func (c *Cache) Load(key CacheKey) *dns.Msg {
 		return nil
 	}
 
-	elapsed := time.Since(item.ts)
-	if item.ttl > 0 {
-		if elapsed > item.ttl {
-			c.m.Delete(key)
-			return nil
-		}
-	} else {
-		for _, rr := range item.msg.Answer {
-			if elapsed > time.Duration(rr.Header().Ttl)*time.Second {
-				c.m.Delete(key)
-				return nil
-			}
-		}
+	if time.Since(item.ts) > item.ttl {
+		c.m.Delete(key)
+		return nil
 	}
 
-	c.logger.Debugf("resolver cache hit %s", key)
+	c.logger.Debugf("resolver cache hit: %s", key)
 
 	return item.msg.Copy()
 }
@@ -76,11 +66,23 @@ func (c *Cache) Store(key CacheKey, mr *dns.Msg, ttl time.Duration) {
 		return
 	}
 
+	if ttl == 0 {
+		for _, answer := range mr.Answer {
+			v := time.Duration(answer.Header().Ttl) * time.Second
+			if ttl == 0 || ttl > v {
+				ttl = v
+			}
+		}
+	}
+	if ttl == 0 {
+		ttl = 30 * time.Second
+	}
+
 	c.m.Store(key, &cacheItem{
 		msg: mr.Copy(),
 		ts:  time.Now(),
 		ttl: ttl,
 	})
 
-	c.logger.Debugf("resolver cache store %s", key)
+	c.logger.Debugf("resolver cache store: %s, ttl: %v", key, ttl)
 }