Skip to content

Instantly share code, notes, and snippets.

@wendal
Created January 22, 2015 08:18
Show Gist options
  • Save wendal/240e7ecc4def0f3eba67 to your computer and use it in GitHub Desktop.
Save wendal/240e7ecc4def0f3eba67 to your computer and use it in GitHub Desktop.
golang的dns缓存, 含流量统计
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
}
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)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment