dns: optimize code

This commit is contained in:
nadoo 2020-05-06 20:10:18 +08:00
parent e3c57ba369
commit ab6a448d00
6 changed files with 43 additions and 41 deletions

View File

@ -31,8 +31,8 @@ type Client struct {
proxy proxy.Proxy proxy proxy.Proxy
cache *Cache cache *Cache
config *Config config *Config
upStream *UpStream upStream *UPStream
upStreamMap map[string]*UpStream upStreamMap map[string]*UPStream
handlers []HandleFunc handlers []HandleFunc
} }
@ -42,8 +42,8 @@ func NewClient(proxy proxy.Proxy, config *Config) (*Client, error) {
proxy: proxy, proxy: proxy,
cache: NewCache(), cache: NewCache(),
config: config, config: config,
upStream: NewUpStream(config.Servers), upStream: NewUPStream(config.Servers),
upStreamMap: make(map[string]*UpStream), upStreamMap: make(map[string]*UPStream),
} }
// custom records // custom records
@ -233,11 +233,11 @@ func (c *Client) exchangeUDP(rc net.Conn, reqBytes []byte) ([]byte, error) {
// SetServers sets upstream dns servers for the given domain. // SetServers sets upstream dns servers for the given domain.
func (c *Client) SetServers(domain string, servers []string) { func (c *Client) SetServers(domain string, servers []string) {
c.upStreamMap[domain] = NewUpStream(servers) c.upStreamMap[domain] = NewUPStream(servers)
} }
// UpStream returns upstream dns server for the given domain. // UpStream returns upstream dns server for the given domain.
func (c *Client) UpStream(domain string) *UpStream { func (c *Client) UpStream(domain string) *UPStream {
domainParts := strings.Split(domain, ".") domainParts := strings.Split(domain, ".")
length := len(domainParts) length := len(domainParts)
for i := length - 1; i >= 0; i-- { for i := length - 1; i >= 0; i-- {

View File

@ -20,13 +20,13 @@ const UDPMaxLen = 510
// HeaderLen is the length of dns msg header. // HeaderLen is the length of dns msg header.
const HeaderLen = 12 const HeaderLen = 12
// Message types // Message types.
const ( const (
Query = 0 Query = 0
Response = 1 Response = 1
) )
// Query types // Query types.
const ( const (
QTypeA uint16 = 1 //ipv4 QTypeA uint16 = 1 //ipv4
QTypeAAAA uint16 = 28 ///ipv6 QTypeAAAA uint16 = 28 ///ipv6
@ -35,7 +35,7 @@ const (
// ClassINET . // ClassINET .
const ClassINET uint16 = 1 const ClassINET uint16 = 1
// Message format // Message format:
// https://tools.ietf.org/html/rfc1035#section-4.1 // https://tools.ietf.org/html/rfc1035#section-4.1
// All communications inside of the domain protocol are carried in a single // All communications inside of the domain protocol are carried in a single
// format called a message. The top level format of message is divided // format called a message. The top level format of message is divided
@ -143,7 +143,7 @@ func UnmarshalMessage(b []byte) (*Message, error) {
return m, nil return m, nil
} }
// Header format // Header format:
// https://tools.ietf.org/html/rfc1035#section-4.1.1 // https://tools.ietf.org/html/rfc1035#section-4.1.1
// The header contains the following fields: // The header contains the following fields:
// //
@ -222,7 +222,7 @@ func UnmarshalHeader(b []byte, h *Header) error {
return nil return nil
} }
// Question format // Question format:
// https://tools.ietf.org/html/rfc1035#section-4.1.2 // https://tools.ietf.org/html/rfc1035#section-4.1.2
// The question section is used to carry the "question" in most queries, // The question section is used to carry the "question" in most queries,
// i.e., the parameters that define what is being asked. The section // i.e., the parameters that define what is being asked. The section
@ -287,7 +287,7 @@ func (m *Message) UnmarshalQuestion(b []byte, q *Question) (n int, err error) {
return idx + 3 + 1, nil return idx + 3 + 1, nil
} }
// RR format // RR format:
// https://tools.ietf.org/html/rfc1035#section-3.2.1 // https://tools.ietf.org/html/rfc1035#section-3.2.1
// https://tools.ietf.org/html/rfc1035#section-4.1.3 // https://tools.ietf.org/html/rfc1035#section-4.1.3
// The answer, authority, and additional sections all share the same // The answer, authority, and additional sections all share the same
@ -328,8 +328,7 @@ type RR struct {
// NewRR returns a new dns rr. // NewRR returns a new dns rr.
func NewRR() *RR { func NewRR() *RR {
rr := &RR{} return &RR{}
return rr
} }
// MarshalTo marshals RR struct to []byte and write to w. // MarshalTo marshals RR struct to []byte and write to w.
@ -403,12 +402,8 @@ func MarshalDomainTo(w io.Writer, domain string) (int, error) {
// UnmarshalDomain gets domain from bytes. // UnmarshalDomain gets domain from bytes.
func (m *Message) UnmarshalDomain(b []byte) (string, int, error) { func (m *Message) UnmarshalDomain(b []byte) (string, int, error) {
if len(b) < 2 {
return "", 0, errors.New("UnmarshalDomain: not enough size")
}
var idx, size int var idx, size int
var labels = []string{} var labels []string
for { for {
// https://tools.ietf.org/html/rfc1035#section-4.1.4 // https://tools.ietf.org/html/rfc1035#section-4.1.4
@ -416,7 +411,10 @@ func (m *Message) UnmarshalDomain(b []byte) (string, int, error) {
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ // +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// | 1 1| OFFSET | // | 1 1| OFFSET |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+ // +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
if b[idx]&0xC0 == 0xC0 { if len(b[idx:]) == 0 {
break
} else if b[idx]&0xC0 == 0xC0 {
offset := binary.BigEndian.Uint16(b[idx : idx+2]) offset := binary.BigEndian.Uint16(b[idx : idx+2])
label, err := m.UnmarshalDomainPoint(int(offset & 0x3FFF)) label, err := m.UnmarshalDomainPoint(int(offset & 0x3FFF))
if err != nil { if err != nil {
@ -426,10 +424,13 @@ func (m *Message) UnmarshalDomain(b []byte) (string, int, error) {
labels = append(labels, label) labels = append(labels, label)
idx += 2 idx += 2
break break
} else { } else {
size = int(b[idx]) size = int(b[idx])
idx++
// root domain name
if size == 0 { if size == 0 {
idx++
break break
} }
@ -437,13 +438,14 @@ func (m *Message) UnmarshalDomain(b []byte) (string, int, error) {
return "", 0, errors.New("UnmarshalDomain: label size larger than 63") return "", 0, errors.New("UnmarshalDomain: label size larger than 63")
} }
if idx+size+1 > len(b) { if idx+size > len(b) {
return "", 0, errors.New("UnmarshalDomain: label size larger than msg length") return "", 0, errors.New("UnmarshalDomain: label size larger than msg length")
} }
labels = append(labels, string(b[idx+1:idx+size+1])) labels = append(labels, string(b[idx:idx+size]))
idx += (size + 1) idx += size
} }
} }
domain := strings.Join(labels, ".") domain := strings.Join(labels, ".")

View File

@ -2,29 +2,29 @@ package dns
import "sync/atomic" import "sync/atomic"
// UpStream is a dns upstream. // UPStream is a dns upstream.
type UpStream struct { type UPStream struct {
index uint32 index uint32
servers []string servers []string
} }
// NewUpStream returns a new UpStream. // NewUPStream returns a new UpStream.
func NewUpStream(servers []string) *UpStream { func NewUPStream(servers []string) *UPStream {
return &UpStream{servers: servers} return &UPStream{servers: servers}
} }
// Server returns a dns server. // Server returns a dns server.
func (u *UpStream) Server() string { func (u *UPStream) Server() string {
return u.servers[atomic.LoadUint32(&u.index)%uint32(len(u.servers))] return u.servers[atomic.LoadUint32(&u.index)%uint32(len(u.servers))]
} }
// Switch switches to the next dns server. // Switch switches to the next dns server.
func (u *UpStream) Switch() string { func (u *UPStream) Switch() string {
return u.servers[atomic.AddUint32(&u.index, 1)%uint32(len(u.servers))] return u.servers[atomic.AddUint32(&u.index, 1)%uint32(len(u.servers))]
} }
// SwitchIf switches to the next dns server if needed. // SwitchIf switches to the next dns server if needed.
func (u *UpStream) SwitchIf(server string) string { func (u *UPStream) SwitchIf(server string) string {
if u.Server() == server { if u.Server() == server {
return u.Switch() return u.Switch()
} }
@ -32,6 +32,6 @@ func (u *UpStream) SwitchIf(server string) string {
} }
// Len returns the number of dns servers. // Len returns the number of dns servers.
func (u *UpStream) Len() int { func (u *UPStream) Len() int {
return len(u.servers) return len(u.servers)
} }

5
go.mod
View File

@ -3,8 +3,7 @@ module github.com/nadoo/glider
go 1.14 go 1.14
require ( require (
github.com/klauspost/cpuid v1.2.4 // indirect github.com/klauspost/reedsolomon v1.9.6 // indirect
github.com/klauspost/reedsolomon v1.9.4 // indirect
github.com/mzz2017/shadowsocksR v0.0.0-20200126130347-721f53a7b15a github.com/mzz2017/shadowsocksR v0.0.0-20200126130347-721f53a7b15a
github.com/nadoo/conflag v0.2.3 github.com/nadoo/conflag v0.2.3
github.com/nadoo/go-shadowsocks2 v0.1.2 github.com/nadoo/go-shadowsocks2 v0.1.2
@ -12,7 +11,7 @@ require (
github.com/tjfoc/gmsm v1.3.0 // indirect github.com/tjfoc/gmsm v1.3.0 // indirect
github.com/xtaci/kcp-go/v5 v5.5.12 github.com/xtaci/kcp-go/v5 v5.5.12
golang.org/x/crypto v0.0.0-20200429183012-4b2356b1ed79 golang.org/x/crypto v0.0.0-20200429183012-4b2356b1ed79
golang.org/x/net v0.0.0-20200501053045-e0ff5e5a1de5 // indirect golang.org/x/net v0.0.0-20200505041828-1ed23360d12c // indirect
golang.org/x/sys v0.0.0-20200501145240-bc7a7d42d5c3 // indirect golang.org/x/sys v0.0.0-20200501145240-bc7a7d42d5c3 // indirect
) )

8
go.sum
View File

@ -17,8 +17,8 @@ github.com/klauspost/cpuid v1.2.4 h1:EBfaK0SWSwk+fgk6efYFWdzl8MwRWoOO1gkmiaTXPW4
github.com/klauspost/cpuid v1.2.4/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek= github.com/klauspost/cpuid v1.2.4/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
github.com/klauspost/reedsolomon v1.9.3 h1:N/VzgeMfHmLc+KHMD1UL/tNkfXAt8FnUqlgXGIduwAY= github.com/klauspost/reedsolomon v1.9.3 h1:N/VzgeMfHmLc+KHMD1UL/tNkfXAt8FnUqlgXGIduwAY=
github.com/klauspost/reedsolomon v1.9.3/go.mod h1:CwCi+NUr9pqSVktrkN+Ondf06rkhYZ/pcNv7fu+8Un4= github.com/klauspost/reedsolomon v1.9.3/go.mod h1:CwCi+NUr9pqSVktrkN+Ondf06rkhYZ/pcNv7fu+8Un4=
github.com/klauspost/reedsolomon v1.9.4 h1:FB9jDBGqUNyhUg4Gszz384ulFqVSc61Pdap+HRPgnSo= github.com/klauspost/reedsolomon v1.9.6 h1:sXZANEgYACIcmbk90z6MV4XL29d0Lm6AFleWRPZJxi8=
github.com/klauspost/reedsolomon v1.9.4/go.mod h1:CwCi+NUr9pqSVktrkN+Ondf06rkhYZ/pcNv7fu+8Un4= github.com/klauspost/reedsolomon v1.9.6/go.mod h1:+8WD025Xpby8/kG5h/HDPIFhiiuGEtZOKw+5Y4drAD8=
github.com/mzz2017/shadowsocksR v0.0.0-20200126130347-721f53a7b15a h1:lM2cQgk1E5mNMUBw5cdE5922waAyan6jlMpbjRxfEas= github.com/mzz2017/shadowsocksR v0.0.0-20200126130347-721f53a7b15a h1:lM2cQgk1E5mNMUBw5cdE5922waAyan6jlMpbjRxfEas=
github.com/mzz2017/shadowsocksR v0.0.0-20200126130347-721f53a7b15a/go.mod h1:1SJEvxD2Y+N7SK2NpCC4wSatvfGGTUo2rhPdthUFsCU= github.com/mzz2017/shadowsocksR v0.0.0-20200126130347-721f53a7b15a/go.mod h1:1SJEvxD2Y+N7SK2NpCC4wSatvfGGTUo2rhPdthUFsCU=
github.com/nadoo/conflag v0.2.2 h1:xywuyaevdBnA3+4g9S11ng+Nby725WN1LXargWnAXpM= github.com/nadoo/conflag v0.2.2 h1:xywuyaevdBnA3+4g9S11ng+Nby725WN1LXargWnAXpM=
@ -65,8 +65,8 @@ golang.org/x/crypto v0.0.0-20200429183012-4b2356b1ed79/go.mod h1:LzIPMQfyMNhhGPh
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
golang.org/x/net v0.0.0-20191014212845-da9a3fd4c582/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20191014212845-da9a3fd4c582/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
golang.org/x/net v0.0.0-20191209160850-c0dbc17a3553/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20191209160850-c0dbc17a3553/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
golang.org/x/net v0.0.0-20200501053045-e0ff5e5a1de5 h1:WQ8q63x+f/zpC8Ac1s9wLElVoHhm32p6tudrU72n1QA= golang.org/x/net v0.0.0-20200505041828-1ed23360d12c h1:zJ0mtu4jCalhKg6Oaukv6iIkb+cOvDrajDH9DH46Q4M=
golang.org/x/net v0.0.0-20200501053045-e0ff5e5a1de5/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= golang.org/x/net v0.0.0-20200505041828-1ed23360d12c/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A=
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20190902133755-9109b7679e13/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190902133755-9109b7679e13/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

View File

@ -13,6 +13,7 @@ import (
"github.com/nadoo/glider/proxy" "github.com/nadoo/glider/proxy"
) )
// SSH is a base ssh struct.
type SSH struct { type SSH struct {
dialer proxy.Dialer dialer proxy.Dialer
proxy proxy.Proxy proxy proxy.Proxy
@ -24,7 +25,7 @@ func init() {
proxy.RegisterDialer("ssh", NewSSHDialer) proxy.RegisterDialer("ssh", NewSSHDialer)
} }
// NewSS returns a ssh proxy. // NewSSH returns a ssh proxy.
func NewSSH(s string, d proxy.Dialer, p proxy.Proxy) (*SSH, error) { func NewSSH(s string, d proxy.Dialer, p proxy.Proxy) (*SSH, error) {
u, err := url.Parse(s) u, err := url.Parse(s)
if err != nil { if err != nil {