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
cache *Cache
config *Config
upStream *UpStream
upStreamMap map[string]*UpStream
upStream *UPStream
upStreamMap map[string]*UPStream
handlers []HandleFunc
}
@ -42,8 +42,8 @@ func NewClient(proxy proxy.Proxy, config *Config) (*Client, error) {
proxy: proxy,
cache: NewCache(),
config: config,
upStream: NewUpStream(config.Servers),
upStreamMap: make(map[string]*UpStream),
upStream: NewUPStream(config.Servers),
upStreamMap: make(map[string]*UPStream),
}
// 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.
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.
func (c *Client) UpStream(domain string) *UpStream {
func (c *Client) UpStream(domain string) *UPStream {
domainParts := strings.Split(domain, ".")
length := len(domainParts)
for i := length - 1; i >= 0; i-- {

View File

@ -20,13 +20,13 @@ const UDPMaxLen = 510
// HeaderLen is the length of dns msg header.
const HeaderLen = 12
// Message types
// Message types.
const (
Query = 0
Response = 1
)
// Query types
// Query types.
const (
QTypeA uint16 = 1 //ipv4
QTypeAAAA uint16 = 28 ///ipv6
@ -35,7 +35,7 @@ const (
// ClassINET .
const ClassINET uint16 = 1
// Message format
// Message format:
// https://tools.ietf.org/html/rfc1035#section-4.1
// All communications inside of the domain protocol are carried in a single
// 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
}
// Header format
// Header format:
// https://tools.ietf.org/html/rfc1035#section-4.1.1
// The header contains the following fields:
//
@ -222,7 +222,7 @@ func UnmarshalHeader(b []byte, h *Header) error {
return nil
}
// Question format
// Question format:
// https://tools.ietf.org/html/rfc1035#section-4.1.2
// The question section is used to carry the "question" in most queries,
// 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
}
// RR format
// RR format:
// https://tools.ietf.org/html/rfc1035#section-3.2.1
// https://tools.ietf.org/html/rfc1035#section-4.1.3
// The answer, authority, and additional sections all share the same
@ -328,8 +328,7 @@ type RR struct {
// NewRR returns a new dns rr.
func NewRR() *RR {
rr := &RR{}
return rr
return &RR{}
}
// 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.
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 labels = []string{}
var labels []string
for {
// 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 |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
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])
label, err := m.UnmarshalDomainPoint(int(offset & 0x3FFF))
if err != nil {
@ -426,10 +424,13 @@ func (m *Message) UnmarshalDomain(b []byte) (string, int, error) {
labels = append(labels, label)
idx += 2
break
} else {
size = int(b[idx])
idx++
// root domain name
if size == 0 {
idx++
break
}
@ -437,13 +438,14 @@ func (m *Message) UnmarshalDomain(b []byte) (string, int, error) {
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")
}
labels = append(labels, string(b[idx+1:idx+size+1]))
idx += (size + 1)
labels = append(labels, string(b[idx:idx+size]))
idx += size
}
}
domain := strings.Join(labels, ".")

View File

@ -2,29 +2,29 @@ package dns
import "sync/atomic"
// UpStream is a dns upstream.
type UpStream struct {
// UPStream is a dns upstream.
type UPStream struct {
index uint32
servers []string
}
// NewUpStream returns a new UpStream.
func NewUpStream(servers []string) *UpStream {
return &UpStream{servers: servers}
// NewUPStream returns a new UpStream.
func NewUPStream(servers []string) *UPStream {
return &UPStream{servers: servers}
}
// 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))]
}
// 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))]
}
// 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 {
return u.Switch()
}
@ -32,6 +32,6 @@ func (u *UpStream) SwitchIf(server string) string {
}
// Len returns the number of dns servers.
func (u *UpStream) Len() int {
func (u *UPStream) Len() int {
return len(u.servers)
}

5
go.mod
View File

@ -3,8 +3,7 @@ module github.com/nadoo/glider
go 1.14
require (
github.com/klauspost/cpuid v1.2.4 // indirect
github.com/klauspost/reedsolomon v1.9.4 // indirect
github.com/klauspost/reedsolomon v1.9.6 // indirect
github.com/mzz2017/shadowsocksR v0.0.0-20200126130347-721f53a7b15a
github.com/nadoo/conflag v0.2.3
github.com/nadoo/go-shadowsocks2 v0.1.2
@ -12,7 +11,7 @@ require (
github.com/tjfoc/gmsm v1.3.0 // indirect
github.com/xtaci/kcp-go/v5 v5.5.12
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
)

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/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.4 h1:FB9jDBGqUNyhUg4Gszz384ulFqVSc61Pdap+HRPgnSo=
github.com/klauspost/reedsolomon v1.9.4/go.mod h1:CwCi+NUr9pqSVktrkN+Ondf06rkhYZ/pcNv7fu+8Un4=
github.com/klauspost/reedsolomon v1.9.6 h1:sXZANEgYACIcmbk90z6MV4XL29d0Lm6AFleWRPZJxi8=
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/go.mod h1:1SJEvxD2Y+N7SK2NpCC4wSatvfGGTUo2rhPdthUFsCU=
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-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-20200501053045-e0ff5e5a1de5 h1:WQ8q63x+f/zpC8Ac1s9wLElVoHhm32p6tudrU72n1QA=
golang.org/x/net v0.0.0-20200501053045-e0ff5e5a1de5/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A=
golang.org/x/net v0.0.0-20200505041828-1ed23360d12c h1:zJ0mtu4jCalhKg6Oaukv6iIkb+cOvDrajDH9DH46Q4M=
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-20190412213103-97732733099d/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"
)
// SSH is a base ssh struct.
type SSH struct {
dialer proxy.Dialer
proxy proxy.Proxy
@ -24,7 +25,7 @@ func init() {
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) {
u, err := url.Parse(s)
if err != nil {