bug fixes in sniffer

app/dispatcher/impl/default.go

@@ -12,6 +12,7 @@ import (
 	"v2ray.com/core/app/proxyman"
 	"v2ray.com/core/app/router"
 	"v2ray.com/core/common"
+	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/proxy"
 	"v2ray.com/core/transport/ray"
@@ -21,11 +22,13 @@ var (
 	errSniffingTimeout = newError("timeout on sniffing")
 )
 
+// DefaultDispatcher is a default implementation of Dispatcher.
 type DefaultDispatcher struct {
 	ohm    proxyman.OutboundHandlerManager
 	router *router.Router
 }
 
+// NewDefaultDispatcher create a new DefaultDispatcher.
 func NewDefaultDispatcher(ctx context.Context, config *dispatcher.Config) (*DefaultDispatcher, error) {
 	space := app.SpaceFromContext(ctx)
 	if space == nil {
@@ -43,21 +46,17 @@ func NewDefaultDispatcher(ctx context.Context, config *dispatcher.Config) (*Defa
 	return d, nil
 }
 
-func (DefaultDispatcher) Start() error {
+func (*DefaultDispatcher) Start() error {
 	return nil
 }
 
-func (DefaultDispatcher) Close() {}
+func (*DefaultDispatcher) Close() {}
 
-func (DefaultDispatcher) Interface() interface{} {
+func (*DefaultDispatcher) Interface() interface{} {
 	return (*dispatcher.Interface)(nil)
 }
 
-type domainOrError struct {
-	domain string
-	err    error
-}
-
+// Dispatch implements Dispatcher.Interface.
 func (d *DefaultDispatcher) Dispatch(ctx context.Context, destination net.Destination) (ray.InboundRay, error) {
 	if !destination.IsValid() {
 		panic("Dispatcher: Invalid destination.")
@@ -70,47 +69,36 @@ func (d *DefaultDispatcher) Dispatch(ctx context.Context, destination net.Destin
 		go d.routedDispatch(ctx, outbound, destination)
 	} else {
 		go func() {
-			done := make(chan domainOrError)
-			go snifer(ctx, sniferList, outbound, done)
-			de := <-done
-			if de.err != nil {
-				log.Trace(newError("failed to snif").Base(de.err))
-				return
+			domain, err := snifer(ctx, sniferList, outbound)
+			if err == nil {
+				log.Trace(newError("sniffed domain: ", domain))
+				destination.Address = net.ParseAddress(domain)
+				ctx = proxy.ContextWithTarget(ctx, destination)
 			}
-			log.Trace(newError("sniffed domain: ", de.domain))
-			destination.Address = net.ParseAddress(de.domain)
-			ctx = proxy.ContextWithTarget(ctx, destination)
 			d.routedDispatch(ctx, outbound, destination)
 		}()
 	}
 	return outbound, nil
 }
 
-func snifer(ctx context.Context, sniferList []proxyman.KnownProtocols, outbound ray.OutboundRay, done chan<- domainOrError) {
-	payload := make([]byte, 2048)
+func snifer(ctx context.Context, sniferList []proxyman.KnownProtocols, outbound ray.OutboundRay) (string, error) {
+	payload := buf.New()
+	defer payload.Release()
+
 	totalAttempt := 0
 	for {
 		select {
 		case <-ctx.Done():
-			done <- domainOrError{
-				domain: "",
-				err:    ctx.Err(),
-			}
-			return
+			return "", ctx.Err()
 		case <-time.After(time.Millisecond * 100):
 			totalAttempt++
 			if totalAttempt > 5 {
-				done <- domainOrError{
-					domain: "",
-					err:    errSniffingTimeout,
-				}
-				return
+				return "", errSniffingTimeout
 			}
-			mb := outbound.OutboundInput().Peek()
-			if mb.IsEmpty() {
+			outbound.OutboundInput().Peek(payload)
+			if payload.IsEmpty() {
 				continue
 			}
-			nBytes := mb.Copy(payload)
 			for _, protocol := range sniferList {
 				var f func([]byte) (string, error)
 				switch protocol {
@@ -122,21 +110,13 @@ func snifer(ctx context.Context, sniferList []proxyman.KnownProtocols, outbound
 					panic("Unsupported protocol")
 				}
 
-				domain, err := f(payload[:nBytes])
+				domain, err := f(payload.Bytes())
 				if err != ErrMoreData {
-					done <- domainOrError{
-						domain: domain,
-						err:    err,
-					}
-					return
+					return domain, err
 				}
 			}
-			if nBytes == 2048 {
-				done <- domainOrError{
-					domain: "",
-					err:    ErrInvalidData,
-				}
-				return
+			if payload.IsFull() {
+				return "", ErrInvalidData
 			}
 		}
 	}

app/dispatcher/impl/sniffer.go

@@ -44,7 +44,8 @@ func SniffHTTP(b []byte) (string, error) {
 		key := strings.ToLower(string(parts[0]))
 		value := strings.ToLower(string(bytes.Trim(parts[1], " ")))
 		if key == "host" {
-			return value, nil
+			domain := strings.Split(value, ":")
+			return strings.TrimSpace(domain[0]), nil
 		}
 	}
 	return "", ErrMoreData
@@ -60,11 +61,11 @@ func ReadClientHello(data []byte) (string, error) {
 	if len(data) < 42 {
 		return "", ErrMoreData
 	}
-	sessionIdLen := int(data[38])
-	if sessionIdLen > 32 || len(data) < 39+sessionIdLen {
+	sessionIDLen := int(data[38])
+	if sessionIDLen > 32 || len(data) < 39+sessionIDLen {
 		return "", ErrInvalidData
 	}
-	data = data[39+sessionIdLen:]
+	data = data[39+sessionIDLen:]
 	if len(data) < 2 {
 		return "", ErrMoreData
 	}

common/protocol/server_picker.go

@@ -13,32 +13,32 @@ func NewServerList() *ServerList {
 	return &ServerList{}
 }
 
-func (v *ServerList) AddServer(server *ServerSpec) {
-	v.Lock()
-	defer v.Unlock()
+func (sl *ServerList) AddServer(server *ServerSpec) {
+	sl.Lock()
+	defer sl.Unlock()
 
-	v.servers = append(v.servers, server)
+	sl.servers = append(sl.servers, server)
 }
 
-func (v *ServerList) Size() uint32 {
-	v.RLock()
-	defer v.RUnlock()
+func (sl *ServerList) Size() uint32 {
+	sl.RLock()
+	defer sl.RUnlock()
 
-	return uint32(len(v.servers))
+	return uint32(len(sl.servers))
 }
 
-func (v *ServerList) GetServer(idx uint32) *ServerSpec {
-	v.RLock()
-	defer v.RUnlock()
+func (sl *ServerList) GetServer(idx uint32) *ServerSpec {
+	sl.Lock()
+	defer sl.Unlock()
 
 	for {
-		if idx >= uint32(len(v.servers)) {
+		if idx >= uint32(len(sl.servers)) {
 			return nil
 		}
 
-		server := v.servers[idx]
+		server := sl.servers[idx]
 		if !server.IsValid() {
-			v.RemoveServer(idx)
+			sl.removeServer(idx)
 			continue
 		}
 
@@ -46,11 +46,10 @@ func (v *ServerList) GetServer(idx uint32) *ServerSpec {
 	}
 }
 
-// Private: Visible for testing.
-func (v *ServerList) RemoveServer(idx uint32) {
-	n := len(v.servers)
-	v.servers[idx] = v.servers[n-1]
-	v.servers = v.servers[:n-1]
+func (sl *ServerList) removeServer(idx uint32) {
+	n := len(sl.servers)
+	sl.servers[idx] = sl.servers[n-1]
+	sl.servers = sl.servers[:n-1]
 }
 
 type ServerPicker interface {
@@ -70,21 +69,21 @@ func NewRoundRobinServerPicker(serverlist *ServerList) *RoundRobinServerPicker {
 	}
 }
 
-func (v *RoundRobinServerPicker) PickServer() *ServerSpec {
-	v.Lock()
-	defer v.Unlock()
+func (p *RoundRobinServerPicker) PickServer() *ServerSpec {
+	p.Lock()
+	defer p.Unlock()
 
-	next := v.nextIndex
-	server := v.serverlist.GetServer(next)
+	next := p.nextIndex
+	server := p.serverlist.GetServer(next)
 	if server == nil {
 		next = 0
-		server = v.serverlist.GetServer(0)
+		server = p.serverlist.GetServer(0)
 	}
 	next++
-	if next >= v.serverlist.Size() {
+	if next >= p.serverlist.Size() {
 		next = 0
 	}
-	v.nextIndex = next
+	p.nextIndex = next
 
 	return server
 }

common/protocol/server_spec.go

@@ -13,34 +13,34 @@ type ValidationStrategy interface {
 	Invalidate()
 }
 
-type AlwaysValidStrategy struct{}
+type alwaysValidStrategy struct{}
 
 func AlwaysValid() ValidationStrategy {
-	return AlwaysValidStrategy{}
+	return alwaysValidStrategy{}
 }
 
-func (v AlwaysValidStrategy) IsValid() bool {
+func (alwaysValidStrategy) IsValid() bool {
 	return true
 }
 
-func (v AlwaysValidStrategy) Invalidate() {}
+func (alwaysValidStrategy) Invalidate() {}
 
-type TimeoutValidStrategy struct {
+type timeoutValidStrategy struct {
 	until time.Time
 }
 
 func BeforeTime(t time.Time) ValidationStrategy {
-	return &TimeoutValidStrategy{
+	return &timeoutValidStrategy{
 		until: t,
 	}
 }
 
-func (v *TimeoutValidStrategy) IsValid() bool {
-	return v.until.After(time.Now())
+func (s *timeoutValidStrategy) IsValid() bool {
+	return s.until.After(time.Now())
 }
 
-func (v *TimeoutValidStrategy) Invalidate() {
-	v.until = time.Time{}
+func (s *timeoutValidStrategy) Invalidate() {
+	s.until = time.Time{}
 }
 
 type ServerSpec struct {
@@ -63,19 +63,19 @@ func NewServerSpecFromPB(spec ServerEndpoint) *ServerSpec {
 	return NewServerSpec(dest, AlwaysValid(), spec.User...)
 }
 
-func (v *ServerSpec) Destination() net.Destination {
-	return v.dest
+func (s *ServerSpec) Destination() net.Destination {
+	return s.dest
 }
 
-func (v *ServerSpec) HasUser(user *User) bool {
-	v.RLock()
-	defer v.RUnlock()
+func (s *ServerSpec) HasUser(user *User) bool {
+	s.RLock()
+	defer s.RUnlock()
 
 	accountA, err := user.GetTypedAccount()
 	if err != nil {
 		return false
 	}
-	for _, u := range v.users {
+	for _, u := range s.users {
 		accountB, err := u.GetTypedAccount()
 		if err == nil && accountA.Equals(accountB) {
 			return true
@@ -84,26 +84,29 @@ func (v *ServerSpec) HasUser(user *User) bool {
 	return false
 }
 
-func (v *ServerSpec) AddUser(user *User) {
-	if v.HasUser(user) {
+func (s *ServerSpec) AddUser(user *User) {
+	if s.HasUser(user) {
 		return
 	}
 
-	v.Lock()
-	defer v.Unlock()
+	s.Lock()
+	defer s.Unlock()
 
-	v.users = append(v.users, user)
+	s.users = append(s.users, user)
 }
 
-func (v *ServerSpec) PickUser() *User {
-	userCount := len(v.users)
+func (s *ServerSpec) PickUser() *User {
+	s.RLock()
+	defer s.RUnlock()
+
+	userCount := len(s.users)
 	switch userCount {
 	case 0:
 		return nil
 	case 1:
-		return v.users[0]
+		return s.users[0]
 	default:
-		return v.users[dice.Roll(userCount)]
+		return s.users[dice.Roll(userCount)]
 	}
 }
 

proxy/vmess/encoding/server.go

@@ -45,14 +45,13 @@ func NewSessionHistory(ctx context.Context) *SessionHistory {
 func (h *SessionHistory) add(session sessionId) {
 	h.Lock()
 	h.cache[session] = time.Now().Add(time.Minute * 3)
+	h.Unlock()
 
 	select {
 	case <-h.token.Wait():
 		go h.run()
 	default:
 	}
-
-	h.Unlock()
 }
 
 func (h *SessionHistory) has(session sessionId) bool {

testing/scenarios/feature_test.go

@@ -7,6 +7,8 @@ import (
 	"testing"
 	"time"
 
+	"io/ioutil"
+
 	xproxy "golang.org/x/net/proxy"
 	"v2ray.com/core"
 	"v2ray.com/core/app/log"
@@ -738,6 +740,8 @@ func TestDomainSniffing(t *testing.T) {
 		resp, err := client.Get("https://www.github.com/")
 		assert.Error(err).IsNil()
 		assert.Int(resp.StatusCode).Equals(200)
+
+		assert.Error(resp.Write(ioutil.Discard)).IsNil()
 	}
 
 	CloseAllServers()

tools/conf/v2ray.go

@@ -30,13 +30,29 @@ var (
 	}, "protocol", "settings")
 )
 
+func toProtocolList(s []string) ([]proxyman.KnownProtocols, error) {
+	kp := make([]proxyman.KnownProtocols, 0, 8)
+	for _, p := range s {
+		switch strings.ToLower(p) {
+		case "http":
+			kp = append(kp, proxyman.KnownProtocols_HTTP)
+		case "https", "tls", "ssl":
+			kp = append(kp, proxyman.KnownProtocols_TLS)
+		default:
+			return nil, newError("Unknown protocol: ", p)
+		}
+	}
+	return kp, nil
+}
+
 type InboundConnectionConfig struct {
-	Port          uint16          `json:"port"`
-	Listen        *Address        `json:"listen"`
-	Protocol      string          `json:"protocol"`
-	StreamSetting *StreamConfig   `json:"streamSettings"`
-	Settings      json.RawMessage `json:"settings"`
-	Tag           string          `json:"tag"`
+	Port           uint16          `json:"port"`
+	Listen         *Address        `json:"listen"`
+	Protocol       string          `json:"protocol"`
+	StreamSetting  *StreamConfig   `json:"streamSettings"`
+	Settings       json.RawMessage `json:"settings"`
+	Tag            string          `json:"tag"`
+	DomainOverride *StringList     `json:"domainOverride"`
 }
 
 func (v *InboundConnectionConfig) Build() (*proxyman.InboundHandlerConfig, error) {
@@ -59,6 +75,13 @@ func (v *InboundConnectionConfig) Build() (*proxyman.InboundHandlerConfig, error
 		}
 		receiverConfig.StreamSettings = ts
 	}
+	if v.DomainOverride != nil {
+		kp, err := toProtocolList(*v.DomainOverride)
+		if err != nil {
+			return nil, newError("failed to parse inbound config").Base(err)
+		}
+		receiverConfig.DomainOverride = kp
+	}
 
 	jsonConfig, err := inboundConfigLoader.LoadWithID(v.Settings, v.Protocol)
 	if err != nil {
@@ -183,13 +206,14 @@ func (v *InboundDetourAllocationConfig) Build() (*proxyman.AllocationStrategy, e
 }
 
 type InboundDetourConfig struct {
-	Protocol      string                         `json:"protocol"`
-	PortRange     *PortRange                     `json:"port"`
-	ListenOn      *Address                       `json:"listen"`
-	Settings      json.RawMessage                `json:"settings"`
-	Tag           string                         `json:"tag"`
-	Allocation    *InboundDetourAllocationConfig `json:"allocate"`
-	StreamSetting *StreamConfig                  `json:"streamSettings"`
+	Protocol       string                         `json:"protocol"`
+	PortRange      *PortRange                     `json:"port"`
+	ListenOn       *Address                       `json:"listen"`
+	Settings       json.RawMessage                `json:"settings"`
+	Tag            string                         `json:"tag"`
+	Allocation     *InboundDetourAllocationConfig `json:"allocate"`
+	StreamSetting  *StreamConfig                  `json:"streamSettings"`
+	DomainOverride *StringList                    `json:"domainOverride"`
 }
 
 func (v *InboundDetourConfig) Build() (*proxyman.InboundHandlerConfig, error) {
@@ -220,6 +244,13 @@ func (v *InboundDetourConfig) Build() (*proxyman.InboundHandlerConfig, error) {
 		}
 		receiverSettings.StreamSettings = ss
 	}
+	if v.DomainOverride != nil {
+		kp, err := toProtocolList(*v.DomainOverride)
+		if err != nil {
+			return nil, newError("failed to parse inbound detour config").Base(err)
+		}
+		receiverSettings.DomainOverride = kp
+	}
 
 	rawConfig, err := inboundConfigLoader.LoadWithID(v.Settings, v.Protocol)
 	if err != nil {

transport/ray/direct.go

@@ -75,11 +75,13 @@ func (s *Stream) getData() (buf.MultiBuffer, error) {
 	return nil, nil
 }
 
-func (s *Stream) Peek() buf.MultiBuffer {
+func (s *Stream) Peek(b *buf.Buffer) {
 	s.access.RLock()
 	defer s.access.RUnlock()
 
-	return s.data
+	b.Reset(func(data []byte) (int, error) {
+		return s.data.Copy(data), nil
+	})
 }
 
 func (s *Stream) Read() (buf.MultiBuffer, error) {

transport/ray/ray.go

@@ -42,7 +42,7 @@ type InputStream interface {
 	buf.Reader
 	buf.TimeoutReader
 	RayStream
-	Peek() buf.MultiBuffer
+	Peek(*buf.Buffer)
 }
 
 type OutputStream interface {