golang的dns缓存, 含流量统计

ddns.go

package ddns

import (
	"log"
	"net"
	"sync"
)

var (
	dnsLock  = &sync.Mutex{}
	dnsCache = map[string]string{}
)

func FetchDns(hostport string) string {
	dnsLock.Lock()
	defer dnsLock.Unlock()
	host, port, _ := net.SplitHostPort(hostport)
	ip := dnsCache[host]
	if ip != "" {
		return ip + ":" + port
	}
	ips, err := net.LookupIP(host)
	if err != nil {
		log.Println(hostport, err)
		return ""
	}
	if len(ips) == 0 {
		return ""
	}
	ip = ips[0].String()
	dnsCache[host] = ip
	//log.Println(hostport, ip, port)
	return ip + ":" + port
}

dnet.go

package enet

import (
	"danoo.com/util/ddns"
	"net"
	"net/http"
	"net/http/httputil"
	"sync"
	"time"
)

type NetTraffic struct {
	Read  int64 `json:"read"`
	Write int64 `json:"write"`
	lock  *sync.Mutex
}

func (netTraffic *NetTraffic) AddRead(n int) {
	//netTraffic.lock.Lock()
	//defer netTraffic.lock.Unlock()
	netTraffic.Read += int64(n)
}

func (netTraffic *NetTraffic) AddWrite(n int) {
	//netTraffic.lock.Lock()
	//defer netTraffic.lock.Unlock()
	netTraffic.Write += int64(n)
}

type ConnProxy struct {
	conn       net.Conn
	netTraffic *NetTraffic
}

func (cproxy *ConnProxy) Read(b []byte) (n int, err error) {
	n, err = cproxy.conn.Read(b)
	if n > 0 {
		cproxy.netTraffic.AddRead(n)
	}
	return
}

func (cproxy *ConnProxy) Write(b []byte) (n int, err error) {
	n, err = cproxy.conn.Write(b)
	if n > 0 {
		cproxy.netTraffic.AddWrite(n)
	}
	return
}

func (cproxy *ConnProxy) Close() error {
	return cproxy.conn.Close()
}

func (cproxy *ConnProxy) LocalAddr() net.Addr {
	return cproxy.conn.LocalAddr()
}

func (cproxy *ConnProxy) RemoteAddr() net.Addr {
	return cproxy.conn.RemoteAddr()
}

func (cproxy *ConnProxy) SetDeadline(t time.Time) error {
	return cproxy.conn.SetDeadline(t)
}

func (cproxy *ConnProxy) SetReadDeadline(t time.Time) error {
	return cproxy.conn.SetReadDeadline(t)
}

func (cproxy *ConnProxy) SetWriteDeadline(t time.Time) error {
	return cproxy.conn.SetWriteDeadline(t)
}

func NewHttpClient(netTraffic *NetTraffic, timeout time.Duration) http.Client {
	return http.Client{Transport: NewTransport(netTraffic, timeout)}
}

func NewTransport(netTraffic *NetTraffic, timeout time.Duration) *http.Transport {
	tr := &http.Transport{DisableKeepAlives: true}
	tr.Dial = NewDial(netTraffic, timeout)
	return tr
}

func NewDial(netTraffic *NetTraffic, timeout time.Duration) func(net2 string, addr string) (net.Conn, error) {
	if netTraffic == nil {
		panic("netTraffic is nil")
	}
	return func(net2 string, addr string) (net.Conn, error) {
		ipport := ddns.FetchDns(addr)
		if ipport != "" {
			addr = ipport
		}
		_conn, err := net.DialTimeout(net2, addr, timeout)
		if err != nil {
			return nil, err
		}
		//return _conn, nil // disable netTraffic :  uncomment this line and commnet next line 
		return &ConnProxy{_conn, netTraffic}, nil
	}
}

type ProxyServer struct {
	Transport *http.Transport
}

func (p *ProxyServer) ServeHTTP(rw http.ResponseWriter, source *http.Request) {
	director := func(req *http.Request) {
		req.URL.Scheme = "http"
		req.URL.Host = source.Host
		req.URL.Path = source.URL.Path
		req.URL.RawQuery = source.URL.RawQuery
	}
	rp := &httputil.ReverseProxy{Director: director}
	rp.Transport = p.Transport
	rp.ServeHTTP(rw, source)
}