Skip to content

Commit

Permalink
dns: optimize code
Browse files Browse the repository at this point in the history
  • Loading branch information
@nadoo
nadoo committed May 6, 2020
1 parent e3c57ba commit ab6a448
Show file tree
Hide file tree
Showing 6 changed files with 43 additions and 41 deletions.
12 changes: 6 additions & 6 deletions dns/client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand All @@ -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
Expand Down Expand Up @@ -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-- {
Expand Down
38 changes: 20 additions & 18 deletions dns/message.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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:
//
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -403,20 +402,19 @@ 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
// "Message compression",
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// | 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 {
Expand All @@ -426,24 +424,28 @@ 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
}

if size > 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")
}

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, ".")
Expand Down
18 changes: 9 additions & 9 deletions dns/upstream.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,36 +2,36 @@ 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()
}
return u.Server()
}

// Len returns the number of dns servers.
func (u *UpStream) Len() int {
func (u *UPStream) Len() int {
return len(u.servers)
}
5 changes: 2 additions & 3 deletions go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,16 +3,15 @@ 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
github.com/pkg/errors v0.9.1 // indirect
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
)

Expand Down
8 changes: 4 additions & 4 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down Expand Up @@ -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=
Expand Down
3 changes: 2 additions & 1 deletion proxy/ssh/ssh.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 {
Expand Down

0 comments on commit ab6a448

Please sign in to comment.