From d730637239e5fc5d49dcc655aa89aa3e8d5ab1ec Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Fri, 12 Oct 2018 23:57:56 +0200
Subject: [PATCH] Require Type() for Feature

---
 app/commander/commander.go                |  4 +--
 app/dispatcher/default.go                 |  8 ++++--
 app/dns/server.go                         |  6 +++-
 app/log/log.go                            |  2 +-
 app/policy/manager.go                     |  7 ++++-
 app/proxyman/command/command.go           |  2 +-
 app/proxyman/inbound/inbound.go           |  7 ++++-
 app/proxyman/mux/mux.go                   |  4 +++
 app/proxyman/outbound/handler.go          |  2 +-
 app/proxyman/outbound/handler_test.go     |  2 +-
 app/proxyman/outbound/outbound.go         | 16 +++++++----
 app/router/router.go                      |  7 ++++-
 app/stats/stats.go                        |  6 +++-
 common/interfaces.go                      |  1 +
 dns.go                                    |  4 +++
 features/dns/client.go                    |  4 +++
 features/feature.go                       |  1 +
 features/inbound/inbound.go               |  8 ++++--
 features/outbound/outbound.go             | 12 +++++---
 features/policy/policy.go                 |  4 +++
 features/routing/dispatcher.go            |  4 +++
 features/routing/router.go                |  4 +++
 features/stats/stats.go                   |  4 +++
 network.go                                | 34 ++++++++++++++---------
 policy.go                                 |  4 +++
 router.go                                 |  8 ++++++
 stats.go                                  |  4 +++
 testing/scenarios/socks_test.go           |  4 +--
 transport/internet/udp/dispatcher_test.go |  5 ++++
 v2ray.go                                  | 33 ++++++++++++++++------
 30 files changed, 162 insertions(+), 49 deletions(-)

diff --git a/app/commander/commander.go b/app/commander/commander.go
index 0048128f8..87e9bc9b5 100644
--- a/app/commander/commander.go
+++ b/app/commander/commander.go
@@ -21,7 +21,7 @@ type Commander struct {
 	server *grpc.Server
 	config Config
 	v      *core.Instance
-	ohm    outbound.HandlerManager
+	ohm    outbound.Manager
 }
 
 // NewCommander creates a new Commander based on the given config.
@@ -32,7 +32,7 @@ func NewCommander(ctx context.Context, config *Config) (*Commander, error) {
 		ohm:    v.OutboundHandlerManager(),
 		v:      v,
 	}
-	if err := v.RegisterFeature((*Commander)(nil), c); err != nil {
+	if err := v.RegisterFeature(c); err != nil {
 		return nil, err
 	}
 	return c, nil
diff --git a/app/dispatcher/default.go b/app/dispatcher/default.go
index 845f74b5a..8517d82eb 100644
--- a/app/dispatcher/default.go
+++ b/app/dispatcher/default.go
@@ -84,7 +84,7 @@ func (r *cachedReader) CloseError() {
 
 // DefaultDispatcher is a default implementation of Dispatcher.
 type DefaultDispatcher struct {
-	ohm    outbound.HandlerManager
+	ohm    outbound.Manager
 	router routing.Router
 	policy policy.Manager
 	stats  feature_stats.Manager
@@ -100,12 +100,16 @@ func NewDefaultDispatcher(ctx context.Context, config *Config) (*DefaultDispatch
 		stats:  v.Stats(),
 	}
 
-	if err := v.RegisterFeature((*routing.Dispatcher)(nil), d); err != nil {
+	if err := v.RegisterFeature(d); err != nil {
 		return nil, newError("unable to register Dispatcher").Base(err)
 	}
 	return d, nil
 }
 
+func (*DefaultDispatcher) Type() interface{} {
+	return routing.DispatcherType()
+}
+
 // Start implements common.Runnable.
 func (*DefaultDispatcher) Start() error {
 	return nil
diff --git a/app/dns/server.go b/app/dns/server.go
index d93b7af37..ce5afe34a 100644
--- a/app/dns/server.go
+++ b/app/dns/server.go
@@ -41,7 +41,7 @@ func New(ctx context.Context, config *Config) (*Server, error) {
 	server.hosts = hosts
 
 	v := core.MustFromContext(ctx)
-	if err := v.RegisterFeature((*dns.Client)(nil), server); err != nil {
+	if err := v.RegisterFeature(server); err != nil {
 		return nil, newError("unable to register DNSClient.").Base(err)
 	}
 
@@ -97,6 +97,10 @@ func New(ctx context.Context, config *Config) (*Server, error) {
 	return server, nil
 }
 
+func (*Server) Type() interface{} {
+	return dns.ClientType()
+}
+
 // Start implements common.Runnable.
 func (s *Server) Start() error {
 	return nil
diff --git a/app/log/log.go b/app/log/log.go
index 82e95aeab..5cb0bfef9 100644
--- a/app/log/log.go
+++ b/app/log/log.go
@@ -30,7 +30,7 @@ func New(ctx context.Context, config *Config) (*Instance, error) {
 
 	v := core.FromContext(ctx)
 	if v != nil {
-		common.Must(v.RegisterFeature((*log.Handler)(nil), g))
+		common.Must(v.RegisterFeature(g))
 	}
 
 	return g, nil
diff --git a/app/policy/manager.go b/app/policy/manager.go
index 5da6ce28c..d23ddf0eb 100644
--- a/app/policy/manager.go
+++ b/app/policy/manager.go
@@ -30,7 +30,7 @@ func New(ctx context.Context, config *Config) (*Instance, error) {
 
 	v := core.FromContext(ctx)
 	if v != nil {
-		if err := v.RegisterFeature((*policy.Manager)(nil), m); err != nil {
+		if err := v.RegisterFeature(m); err != nil {
 			return nil, newError("unable to register PolicyManager in core").Base(err).AtError()
 		}
 	}
@@ -38,6 +38,11 @@ func New(ctx context.Context, config *Config) (*Instance, error) {
 	return m, nil
 }
 
+// Type implements common.HasType.
+func (*Instance) Type() interface{} {
+	return policy.ManagerType()
+}
+
 // ForLevel implements policy.Manager.
 func (m *Instance) ForLevel(level uint32) policy.Session {
 	if p, ok := m.levels[level]; ok {
diff --git a/app/proxyman/command/command.go b/app/proxyman/command/command.go
index e9cc8aabf..2421a5b1b 100755
--- a/app/proxyman/command/command.go
+++ b/app/proxyman/command/command.go
@@ -64,7 +64,7 @@ func (op *RemoveUserOperation) ApplyInbound(ctx context.Context, handler inbound
 type handlerServer struct {
 	s   *core.Instance
 	ihm inbound.Manager
-	ohm outbound.HandlerManager
+	ohm outbound.Manager
 }
 
 func (s *handlerServer) AddInbound(ctx context.Context, request *AddInboundRequest) (*AddInboundResponse, error) {
diff --git a/app/proxyman/inbound/inbound.go b/app/proxyman/inbound/inbound.go
index 8caf34227..98071269c 100644
--- a/app/proxyman/inbound/inbound.go
+++ b/app/proxyman/inbound/inbound.go
@@ -28,12 +28,17 @@ func New(ctx context.Context, config *proxyman.InboundConfig) (*Manager, error)
 		taggedHandlers: make(map[string]inbound.Handler),
 	}
 	v := core.MustFromContext(ctx)
-	if err := v.RegisterFeature((*inbound.Manager)(nil), m); err != nil {
+	if err := v.RegisterFeature(m); err != nil {
 		return nil, newError("unable to register InboundHandlerManager").Base(err)
 	}
 	return m, nil
 }
 
+// Type implements common.HasType.
+func (*Manager) Type() interface{} {
+	return inbound.ManagerType()
+}
+
 // AddHandler implements inbound.Manager.
 func (m *Manager) AddHandler(ctx context.Context, handler inbound.Handler) error {
 	m.access.Lock()
diff --git a/app/proxyman/mux/mux.go b/app/proxyman/mux/mux.go
index 71284872e..73e4b2f50 100644
--- a/app/proxyman/mux/mux.go
+++ b/app/proxyman/mux/mux.go
@@ -310,6 +310,10 @@ func NewServer(ctx context.Context) *Server {
 	return s
 }
 
+func (s *Server) Type() interface{} {
+	return s.dispatcher.Type()
+}
+
 func (s *Server) Dispatch(ctx context.Context, dest net.Destination) (*vio.Link, error) {
 	if dest.Address != muxCoolAddress {
 		return s.dispatcher.Dispatch(ctx, dest)
diff --git a/app/proxyman/outbound/handler.go b/app/proxyman/outbound/handler.go
index fbab94399..70655d7a6 100644
--- a/app/proxyman/outbound/handler.go
+++ b/app/proxyman/outbound/handler.go
@@ -21,7 +21,7 @@ type Handler struct {
 	senderSettings  *proxyman.SenderConfig
 	streamSettings  *internet.MemoryStreamConfig
 	proxy           proxy.Outbound
-	outboundManager outbound.HandlerManager
+	outboundManager outbound.Manager
 	mux             *mux.ClientManager
 }
 
diff --git a/app/proxyman/outbound/handler_test.go b/app/proxyman/outbound/handler_test.go
index 3423ad8d1..d0edde503 100644
--- a/app/proxyman/outbound/handler_test.go
+++ b/app/proxyman/outbound/handler_test.go
@@ -12,5 +12,5 @@ func TestInterfaces(t *testing.T) {
 	assert := With(t)
 
 	assert((*Handler)(nil), Implements, (*outbound.Handler)(nil))
-	assert((*Manager)(nil), Implements, (*outbound.HandlerManager)(nil))
+	assert((*Manager)(nil), Implements, (*outbound.Manager)(nil))
 }
diff --git a/app/proxyman/outbound/outbound.go b/app/proxyman/outbound/outbound.go
index adae1f713..4a072a1f6 100644
--- a/app/proxyman/outbound/outbound.go
+++ b/app/proxyman/outbound/outbound.go
@@ -27,12 +27,16 @@ func New(ctx context.Context, config *proxyman.OutboundConfig) (*Manager, error)
 		taggedHandler: make(map[string]outbound.Handler),
 	}
 	v := core.MustFromContext(ctx)
-	if err := v.RegisterFeature((*outbound.HandlerManager)(nil), m); err != nil {
-		return nil, newError("unable to register OutboundHandlerManager").Base(err)
+	if err := v.RegisterFeature(m); err != nil {
+		return nil, newError("unable to register outbound.Manager").Base(err)
 	}
 	return m, nil
 }
 
+func (m *Manager) Type() interface{} {
+	return outbound.ManagerType()
+}
+
 // Start implements core.Feature
 func (m *Manager) Start() error {
 	m.access.Lock()
@@ -73,7 +77,7 @@ func (m *Manager) Close() error {
 	return nil
 }
 
-// GetDefaultHandler implements outbound.HandlerManager.
+// GetDefaultHandler implements outbound.Manager.
 func (m *Manager) GetDefaultHandler() outbound.Handler {
 	m.access.RLock()
 	defer m.access.RUnlock()
@@ -84,7 +88,7 @@ func (m *Manager) GetDefaultHandler() outbound.Handler {
 	return m.defaultHandler
 }
 
-// GetHandler implements outbound.HandlerManager.
+// GetHandler implements outbound.Manager.
 func (m *Manager) GetHandler(tag string) outbound.Handler {
 	m.access.RLock()
 	defer m.access.RUnlock()
@@ -94,7 +98,7 @@ func (m *Manager) GetHandler(tag string) outbound.Handler {
 	return nil
 }
 
-// AddHandler implements outbound.HandlerManager.
+// AddHandler implements outbound.Manager.
 func (m *Manager) AddHandler(ctx context.Context, handler outbound.Handler) error {
 	m.access.Lock()
 	defer m.access.Unlock()
@@ -117,7 +121,7 @@ func (m *Manager) AddHandler(ctx context.Context, handler outbound.Handler) erro
 	return nil
 }
 
-// RemoveHandler implements outbound.HandlerManager.
+// RemoveHandler implements outbound.Manager.
 func (m *Manager) RemoveHandler(ctx context.Context, tag string) error {
 	if len(tag) == 0 {
 		return common.ErrNoClue
diff --git a/app/router/router.go b/app/router/router.go
index baec3a453..cfcbbdefd 100644
--- a/app/router/router.go
+++ b/app/router/router.go
@@ -39,7 +39,7 @@ func NewRouter(ctx context.Context, config *Config) (*Router, error) {
 		r.rules[idx].Condition = cond
 	}
 
-	if err := v.RegisterFeature((*routing.Router)(nil), r); err != nil {
+	if err := v.RegisterFeature(r); err != nil {
 		return nil, newError("unable to register Router").Base(err)
 	}
 	return r, nil
@@ -124,6 +124,11 @@ func (*Router) Close() error {
 	return nil
 }
 
+// Type implement common.HasType.
+func (*Router) Type() interface{} {
+	return routing.RouterType()
+}
+
 func init() {
 	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
 		return NewRouter(ctx, config.(*Config))
diff --git a/app/stats/stats.go b/app/stats/stats.go
index a39432ed8..ae6f184b2 100644
--- a/app/stats/stats.go
+++ b/app/stats/stats.go
@@ -44,7 +44,7 @@ func NewManager(ctx context.Context, config *Config) (*Manager, error) {
 
 	v := core.FromContext(ctx)
 	if v != nil {
-		if err := v.RegisterFeature((*stats.Manager)(nil), m); err != nil {
+		if err := v.RegisterFeature(m); err != nil {
 			return nil, newError("failed to register StatManager").Base(err)
 		}
 	}
@@ -52,6 +52,10 @@ func NewManager(ctx context.Context, config *Config) (*Manager, error) {
 	return m, nil
 }
 
+func (*Manager) Type() interface{} {
+	return stats.ManagerType()
+}
+
 func (m *Manager) RegisterCounter(name string) (stats.Counter, error) {
 	m.access.Lock()
 	defer m.access.Unlock()
diff --git a/common/interfaces.go b/common/interfaces.go
index 962141784..c313f9e4f 100644
--- a/common/interfaces.go
+++ b/common/interfaces.go
@@ -25,6 +25,7 @@ type Runnable interface {
 // HasType is the interface for objects that knows its type.
 type HasType interface {
 	// Type returns the type of the object.
+	// Usually it returns (*Type)(nil) of the object.
 	Type() interface{}
 }
 
diff --git a/dns.go b/dns.go
index 34fbe1361..a7f8d7f31 100644
--- a/dns.go
+++ b/dns.go
@@ -13,6 +13,10 @@ type syncDNSClient struct {
 	dns.Client
 }
 
+func (d *syncDNSClient) Type() interface{} {
+	return dns.ClientType()
+}
+
 func (d *syncDNSClient) LookupIP(host string) ([]net.IP, error) {
 	d.RLock()
 	defer d.RUnlock()
diff --git a/features/dns/client.go b/features/dns/client.go
index 70adc064b..822f102a6 100644
--- a/features/dns/client.go
+++ b/features/dns/client.go
@@ -10,3 +10,7 @@ type Client interface {
 	features.Feature
 	LookupIP(host string) ([]net.IP, error)
 }
+
+func ClientType() interface{} {
+	return (*Client)(nil)
+}
diff --git a/features/feature.go b/features/feature.go
index cc52c18d3..2ee1c0c73 100644
--- a/features/feature.go
+++ b/features/feature.go
@@ -5,5 +5,6 @@ import "v2ray.com/core/common"
 // Feature is the interface for V2Ray features. All features must implement this interface.
 // All existing features have an implementation in app directory. These features can be replaced by third-party ones.
 type Feature interface {
+	common.HasType
 	common.Runnable
 }
diff --git a/features/inbound/inbound.go b/features/inbound/inbound.go
index fc7845474..8c9f0e6e4 100644
--- a/features/inbound/inbound.go
+++ b/features/inbound/inbound.go
@@ -23,9 +23,13 @@ type Manager interface {
 	features.Feature
 	// GetHandlers returns an InboundHandler for the given tag.
 	GetHandler(ctx context.Context, tag string) (Handler, error)
-	// AddHandler adds the given handler into this InboundHandlerManager.
+	// AddHandler adds the given handler into this Manager.
 	AddHandler(ctx context.Context, handler Handler) error
 
-	// RemoveHandler removes a handler from InboundHandlerManager.
+	// RemoveHandler removes a handler from Manager.
 	RemoveHandler(ctx context.Context, tag string) error
 }
+
+func ManagerType() interface{} {
+	return (*Manager)(nil)
+}
diff --git a/features/outbound/outbound.go b/features/outbound/outbound.go
index d9f5bb2f5..b79fcf4f3 100644
--- a/features/outbound/outbound.go
+++ b/features/outbound/outbound.go
@@ -15,16 +15,20 @@ type Handler interface {
 	Dispatch(ctx context.Context, link *vio.Link)
 }
 
-// HandlerManager is a feature that manages outbound.Handlers.
-type HandlerManager interface {
+// Manager is a feature that manages outbound.Handlers.
+type Manager interface {
 	features.Feature
 	// GetHandler returns an outbound.Handler for the given tag.
 	GetHandler(tag string) Handler
 	// GetDefaultHandler returns the default outbound.Handler. It is usually the first outbound.Handler specified in the configuration.
 	GetDefaultHandler() Handler
-	// AddHandler adds a handler into this outbound.HandlerManager.
+	// AddHandler adds a handler into this outbound.Manager.
 	AddHandler(ctx context.Context, handler Handler) error
 
-	// RemoveHandler removes a handler from outbound.HandlerManager.
+	// RemoveHandler removes a handler from outbound.Manager.
 	RemoveHandler(ctx context.Context, tag string) error
 }
+
+func ManagerType() interface{} {
+	return (*Manager)(nil)
+}
diff --git a/features/policy/policy.go b/features/policy/policy.go
index 650de93a0..65908815c 100644
--- a/features/policy/policy.go
+++ b/features/policy/policy.go
@@ -67,6 +67,10 @@ type Manager interface {
 	ForSystem() System
 }
 
+func ManagerType() interface{} {
+	return (*Manager)(nil)
+}
+
 var defaultBufferSize int32
 
 func init() {
diff --git a/features/routing/dispatcher.go b/features/routing/dispatcher.go
index ac3177a09..133043ea7 100644
--- a/features/routing/dispatcher.go
+++ b/features/routing/dispatcher.go
@@ -16,3 +16,7 @@ type Dispatcher interface {
 	// Dispatch returns a Ray for transporting data for the given request.
 	Dispatch(ctx context.Context, dest net.Destination) (*vio.Link, error)
 }
+
+func DispatcherType() interface{} {
+	return (*Dispatcher)(nil)
+}
diff --git a/features/routing/router.go b/features/routing/router.go
index ba70a37c8..748b60241 100644
--- a/features/routing/router.go
+++ b/features/routing/router.go
@@ -13,3 +13,7 @@ type Router interface {
 	// PickRoute returns a tag of an OutboundHandler based on the given context.
 	PickRoute(ctx context.Context) (string, error)
 }
+
+func RouterType() interface{} {
+	return (*Router)(nil)
+}
diff --git a/features/stats/stats.go b/features/stats/stats.go
index a5cdf82a7..e1e00891d 100644
--- a/features/stats/stats.go
+++ b/features/stats/stats.go
@@ -24,3 +24,7 @@ func GetOrRegisterCounter(m Manager, name string) (Counter, error) {
 
 	return m.RegisterCounter(name)
 }
+
+func ManagerType() interface{} {
+	return (*Manager)(nil)
+}
diff --git a/network.go b/network.go
index 3e9119568..af0f2ddac 100644
--- a/network.go
+++ b/network.go
@@ -14,6 +14,10 @@ type syncInboundHandlerManager struct {
 	inbound.Manager
 }
 
+func (*syncInboundHandlerManager) Type() interface{} {
+	return inbound.ManagerType()
+}
+
 func (m *syncInboundHandlerManager) GetHandler(ctx context.Context, tag string) (inbound.Handler, error) {
 	m.RLock()
 	defer m.RUnlock()
@@ -68,61 +72,65 @@ func (m *syncInboundHandlerManager) Set(manager inbound.Manager) {
 
 type syncOutboundHandlerManager struct {
 	sync.RWMutex
-	outbound.HandlerManager
+	outbound.Manager
+}
+
+func (*syncOutboundHandlerManager) Type() interface{} {
+	return outbound.ManagerType()
 }
 
 func (m *syncOutboundHandlerManager) GetHandler(tag string) outbound.Handler {
 	m.RLock()
 	defer m.RUnlock()
 
-	if m.HandlerManager == nil {
+	if m.Manager == nil {
 		return nil
 	}
 
-	return m.HandlerManager.GetHandler(tag)
+	return m.Manager.GetHandler(tag)
 }
 
 func (m *syncOutboundHandlerManager) GetDefaultHandler() outbound.Handler {
 	m.RLock()
 	defer m.RUnlock()
 
-	if m.HandlerManager == nil {
+	if m.Manager == nil {
 		return nil
 	}
 
-	return m.HandlerManager.GetDefaultHandler()
+	return m.Manager.GetDefaultHandler()
 }
 
 func (m *syncOutboundHandlerManager) AddHandler(ctx context.Context, handler outbound.Handler) error {
 	m.RLock()
 	defer m.RUnlock()
 
-	if m.HandlerManager == nil {
+	if m.Manager == nil {
 		return newError("OutboundHandlerManager not set.").AtError()
 	}
 
-	return m.HandlerManager.AddHandler(ctx, handler)
+	return m.Manager.AddHandler(ctx, handler)
 }
 
 func (m *syncOutboundHandlerManager) Start() error {
 	m.RLock()
 	defer m.RUnlock()
 
-	if m.HandlerManager == nil {
+	if m.Manager == nil {
 		return newError("OutboundHandlerManager not set.").AtError()
 	}
 
-	return m.HandlerManager.Start()
+	return m.Manager.Start()
 }
 
 func (m *syncOutboundHandlerManager) Close() error {
 	m.RLock()
 	defer m.RUnlock()
 
-	return common.Close(m.HandlerManager)
+	return common.Close(m.Manager)
 }
 
-func (m *syncOutboundHandlerManager) Set(manager outbound.HandlerManager) {
+func (m *syncOutboundHandlerManager) Set(manager outbound.Manager) {
 	if manager == nil {
 		return
 	}
@@ -130,6 +138,6 @@ func (m *syncOutboundHandlerManager) Set(manager outbound.HandlerManager) {
 	m.Lock()
 	defer m.Unlock()
 
-	common.Close(m.HandlerManager) // nolint: errcheck
-	m.HandlerManager = manager
+	common.Close(m.Manager) // nolint: errcheck
+	m.Manager = manager
 }
diff --git a/policy.go b/policy.go
index 24e4aaa77..e5d80d2df 100644
--- a/policy.go
+++ b/policy.go
@@ -13,6 +13,10 @@ type syncPolicyManager struct {
 	policy.Manager
 }
 
+func (*syncPolicyManager) Type() interface{} {
+	return policy.ManagerType()
+}
+
 func (m *syncPolicyManager) ForLevel(level uint32) policy.Session {
 	m.RLock()
 	defer m.RUnlock()
diff --git a/router.go b/router.go
index ba05b3687..5bbc0f8b5 100644
--- a/router.go
+++ b/router.go
@@ -15,6 +15,10 @@ type syncDispatcher struct {
 	routing.Dispatcher
 }
 
+func (*syncDispatcher) Type() interface{} {
+	return routing.DispatcherType()
+}
+
 func (d *syncDispatcher) Dispatch(ctx context.Context, dest net.Destination) (*vio.Link, error) {
 	d.RLock()
 	defer d.RUnlock()
@@ -61,6 +65,10 @@ type syncRouter struct {
 	routing.Router
 }
 
+func (*syncRouter) Type() interface{} {
+	return routing.RouterType()
+}
+
 func (r *syncRouter) PickRoute(ctx context.Context) (string, error) {
 	r.RLock()
 	defer r.RUnlock()
diff --git a/stats.go b/stats.go
index 7665c6d6f..871408ff5 100644
--- a/stats.go
+++ b/stats.go
@@ -11,6 +11,10 @@ type syncStatManager struct {
 	stats.Manager
 }
 
+func (*syncStatManager) Type() interface{} {
+	return stats.ManagerType()
+}
+
 func (s *syncStatManager) Start() error {
 	s.RLock()
 	defer s.RUnlock()
diff --git a/testing/scenarios/socks_test.go b/testing/scenarios/socks_test.go
index 830c04b47..a87a43653 100644
--- a/testing/scenarios/socks_test.go
+++ b/testing/scenarios/socks_test.go
@@ -3,12 +3,12 @@ package scenarios
 import (
 	"testing"
 
-	"v2ray.com/core/app/router"
-
 	xproxy "golang.org/x/net/proxy"
 	socks4 "h12.io/socks"
+
 	"v2ray.com/core"
 	"v2ray.com/core/app/proxyman"
+	"v2ray.com/core/app/router"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/common/protocol"
 	"v2ray.com/core/common/serial"
diff --git a/transport/internet/udp/dispatcher_test.go b/transport/internet/udp/dispatcher_test.go
index 06ba01dee..8d43f4698 100644
--- a/transport/internet/udp/dispatcher_test.go
+++ b/transport/internet/udp/dispatcher_test.go
@@ -9,6 +9,7 @@ import (
 	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/common/vio"
+	"v2ray.com/core/features/routing"
 	. "v2ray.com/core/transport/internet/udp"
 	"v2ray.com/core/transport/pipe"
 	. "v2ray.com/ext/assert"
@@ -30,6 +31,10 @@ func (d *TestDispatcher) Close() error {
 	return nil
 }
 
+func (*TestDispatcher) Type() interface{} {
+	return routing.DispatcherType()
+}
+
 func TestSameDestinationDispatching(t *testing.T) {
 	assert := With(t)
 
diff --git a/v2ray.go b/v2ray.go
index 289836ba0..22b224869 100755
--- a/v2ray.go
+++ b/v2ray.go
@@ -140,10 +140,10 @@ func (s *Instance) Start() error {
 // RegisterFeature registers the given feature into V2Ray.
 // If feature is one of the following types, the corresponding feature in this Instance
 // will be replaced: DNSClient, PolicyManager, Router, Dispatcher, InboundHandlerManager, OutboundHandlerManager.
-func (s *Instance) RegisterFeature(feature interface{}, instance features.Feature) error {
+func (s *Instance) RegisterFeature(instance features.Feature) error {
 	running := false
 
-	switch feature.(type) {
+	switch instance.Type().(type) {
 	case dns.Client, *dns.Client:
 		s.dnsClient.Set(instance.(dns.Client))
 	case policy.Manager, *policy.Manager:
@@ -154,8 +154,8 @@ func (s *Instance) RegisterFeature(feature interface{}, instance features.Featur
 		s.dispatcher.Set(instance.(routing.Dispatcher))
 	case inbound.Manager, *inbound.Manager:
 		s.ihm.Set(instance.(inbound.Manager))
-	case outbound.HandlerManager, *outbound.HandlerManager:
-		s.ohm.Set(instance.(outbound.HandlerManager))
+	case outbound.Manager, *outbound.Manager:
+		s.ohm.Set(instance.(outbound.Manager))
 	case stats.Manager, *stats.Manager:
 		s.stats.Set(instance.(stats.Manager))
 	default:
@@ -178,14 +178,29 @@ func (s *Instance) allFeatures() []features.Feature {
 // GetFeature returns a feature that was registered in this Instance. Nil if not found.
 // The returned Feature must implement common.HasType and whose type equals to the given feature type.
 func (s *Instance) GetFeature(featureType interface{}) features.Feature {
-	for _, f := range s.features {
-		if hasType, ok := f.(common.HasType); ok {
-			if hasType.Type() == featureType {
+	switch featureType.(type) {
+	case dns.Client, *dns.Client:
+		return s.DNSClient()
+	case policy.Manager, *policy.Manager:
+		return s.PolicyManager()
+	case routing.Router, *routing.Router:
+		return s.Router()
+	case routing.Dispatcher, *routing.Dispatcher:
+		return s.Dispatcher()
+	case inbound.Manager, *inbound.Manager:
+		return s.InboundHandlerManager()
+	case outbound.Manager, *outbound.Manager:
+		return s.OutboundHandlerManager()
+	case stats.Manager, *stats.Manager:
+		return s.Stats()
+	default:
+		for _, f := range s.features {
+			if f.Type() == featureType {
 				return f
 			}
 		}
+		return nil
 	}
-	return nil
 }
 
 // DNSClient returns the dns.Client used by this Instance. The returned dns.Client is always functional.
@@ -214,7 +229,7 @@ func (s *Instance) InboundHandlerManager() inbound.Manager {
 }
 
 // OutboundHandlerManager returns the OutboundHandlerManager used by this Instance. If OutboundHandlerManager was not registered before, the returned value doesn't work.
-func (s *Instance) OutboundHandlerManager() outbound.HandlerManager {
+func (s *Instance) OutboundHandlerManager() outbound.Manager {
 	return &(s.ohm)
 }