From 307aac26b32cdf58cd705704ee4a349692ffc27f Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Mon, 22 Oct 2018 11:26:22 +0200
Subject: [PATCH] simplify dependency resolution

---
 app/app.go                         |  1 +
 app/commander/commander.go         |  6 ++--
 app/dispatcher/default.go          | 24 ++++++-------
 app/dns/server.go                  | 13 ++++---
 app/proxyman/command/command.go    |  7 ++--
 app/proxyman/mux/mux.go            |  7 ++--
 app/router/router.go               |  7 ++--
 app/stats/command/command.go       | 14 +++++---
 features/stats/errors.generated.go |  6 ++++
 features/stats/stats.go            | 26 ++++++++++++++
 proxy/freedom/freedom.go           |  9 ++---
 v2ray.go                           | 56 ++++++++++++++++++++++++++----
 v2ray_test.go                      | 15 ++++++++
 13 files changed, 135 insertions(+), 56 deletions(-)
 create mode 100644 features/stats/errors.generated.go

diff --git a/app/app.go b/app/app.go
index 4879f7a4..d42e4cf4 100644
--- a/app/app.go
+++ b/app/app.go
@@ -1 +1,2 @@
+// Package app contains feature implementations of V2Ray. The features may be enabled during runtime.
 package app
diff --git a/app/commander/commander.go b/app/commander/commander.go
index e753448b..41075bb0 100644
--- a/app/commander/commander.go
+++ b/app/commander/commander.go
@@ -12,7 +12,6 @@ import (
 	"v2ray.com/core"
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/signal/done"
-	"v2ray.com/core/features"
 	"v2ray.com/core/features/outbound"
 )
 
@@ -31,9 +30,8 @@ func NewCommander(ctx context.Context, config *Config) (*Commander, error) {
 		tag: config.Tag,
 	}
 
-	v := core.MustFromContext(ctx)
-	v.RequireFeatures([]interface{}{outbound.ManagerType()}, func(fs []features.Feature) {
-		c.ohm = fs[0].(outbound.Manager)
+	core.RequireFeatures(ctx, func(om outbound.Manager) {
+		c.ohm = om
 	})
 
 	for _, rawConfig := range config.Service {
diff --git a/app/dispatcher/default.go b/app/dispatcher/default.go
index 0ba2f4a7..517a4559 100644
--- a/app/dispatcher/default.go
+++ b/app/dispatcher/default.go
@@ -16,7 +16,6 @@ import (
 	"v2ray.com/core/common/protocol"
 	"v2ray.com/core/common/session"
 	"v2ray.com/core/common/vio"
-	"v2ray.com/core/features"
 	"v2ray.com/core/features/outbound"
 	"v2ray.com/core/features/policy"
 	"v2ray.com/core/features/routing"
@@ -94,23 +93,20 @@ type DefaultDispatcher struct {
 func NewDefaultDispatcher(ctx context.Context, config *Config) (*DefaultDispatcher, error) {
 	d := &DefaultDispatcher{}
 
-	v := core.MustFromContext(ctx)
-	v.RequireFeatures([]interface{}{outbound.ManagerType(), routing.RouterType(), policy.ManagerType()}, func(fs []features.Feature) {
-		d.ohm = fs[0].(outbound.Manager)
-		d.router = fs[1].(routing.Router)
-		d.policy = fs[2].(policy.Manager)
-	})
-	v.RequireFeatures([]interface{}{core.ServerType()}, func([]features.Feature) {
-		f := v.GetFeature(stats.ManagerType())
-		if f == nil {
-			return
-		}
-		d.stats = f.(stats.Manager)
-	})
+	core.RequireFeatures(ctx, d.Init)
 
 	return d, nil
 }
 
+// Init initializes DefaultDispatcher.
+// This method is visible for testing purpose.
+func (d *DefaultDispatcher) Init(om outbound.Manager, router routing.Router, pm policy.Manager, sm stats.Manager) {
+	d.ohm = om
+	d.router = router
+	d.policy = pm
+	d.stats = sm
+}
+
 // Type implements common.HasType.
 func (*DefaultDispatcher) Type() interface{} {
 	return routing.DispatcherType()
diff --git a/app/dns/server.go b/app/dns/server.go
index 82902864..68a7c390 100644
--- a/app/dns/server.go
+++ b/app/dns/server.go
@@ -7,14 +7,13 @@ import (
 	"sync"
 	"time"
 
-	"v2ray.com/core/features/routing"
-
 	"v2ray.com/core"
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/common/strmatcher"
 	"v2ray.com/core/features"
 	"v2ray.com/core/features/dns"
+	"v2ray.com/core/features/routing"
 )
 
 type Server struct {
@@ -43,8 +42,6 @@ func New(ctx context.Context, config *Config) (*Server, error) {
 	}
 	server.hosts = hosts
 
-	v := core.MustFromContext(ctx)
-
 	addNameServer := func(endpoint *net.Endpoint) int {
 		address := endpoint.Address.AsAddress()
 		if address.Family().IsDomain() && address.Domain() == "localhost" {
@@ -57,9 +54,9 @@ func New(ctx context.Context, config *Config) (*Server, error) {
 			if dest.Network == net.Network_UDP {
 				idx := len(server.servers)
 				server.servers = append(server.servers, nil)
-				v.RequireFeatures([]interface{}{routing.DispatcherType()}, func(fs []features.Feature) {
-					dispatcher := fs[0].(routing.Dispatcher)
-					server.servers[idx] = NewClassicNameServer(dest, dispatcher, server.clientIP)
+
+				core.RequireFeatures(ctx, func(d routing.Dispatcher) {
+					server.servers[idx] = NewClassicNameServer(dest, d, server.clientIP)
 				})
 			}
 		}
@@ -102,6 +99,7 @@ func New(ctx context.Context, config *Config) (*Server, error) {
 	return server, nil
 }
 
+// Type implements common.HasType.
 func (*Server) Type() interface{} {
 	return dns.ClientType()
 }
@@ -123,6 +121,7 @@ func (s *Server) queryIPTimeout(server NameServerInterface, domain string) ([]ne
 	return ips, err
 }
 
+// LookupIP implements dns.Client.
 func (s *Server) LookupIP(domain string) ([]net.IP, error) {
 	if ip := s.hosts.LookupIP(domain); len(ip) > 0 {
 		return ip, nil
diff --git a/app/proxyman/command/command.go b/app/proxyman/command/command.go
index 14df7fcc..aef371ae 100755
--- a/app/proxyman/command/command.go
+++ b/app/proxyman/command/command.go
@@ -7,7 +7,6 @@ import (
 
 	"v2ray.com/core"
 	"v2ray.com/core/common"
-	"v2ray.com/core/features"
 	"v2ray.com/core/features/inbound"
 	"v2ray.com/core/features/outbound"
 	"v2ray.com/core/proxy"
@@ -132,9 +131,9 @@ func (s *service) Register(server *grpc.Server) {
 	hs := &handlerServer{
 		s: s.v,
 	}
-	s.v.RequireFeatures([]interface{}{inbound.ManagerType(), outbound.ManagerType()}, func(fs []features.Feature) {
-		hs.ihm = fs[0].(inbound.Manager)
-		hs.ohm = fs[1].(outbound.Manager)
+	s.v.RequireFeatures(func(im inbound.Manager, om outbound.Manager) {
+		hs.ihm = im
+		hs.ohm = om
 	})
 	RegisterHandlerServiceServer(server, hs)
 }
diff --git a/app/proxyman/mux/mux.go b/app/proxyman/mux/mux.go
index 57497b0d..c35984d2 100644
--- a/app/proxyman/mux/mux.go
+++ b/app/proxyman/mux/mux.go
@@ -8,8 +8,6 @@ import (
 	"sync"
 	"time"
 
-	"v2ray.com/core/features"
-
 	"v2ray.com/core"
 	"v2ray.com/core/app/proxyman"
 	"v2ray.com/core/common"
@@ -307,9 +305,8 @@ type Server struct {
 // NewServer creates a new mux.Server.
 func NewServer(ctx context.Context) *Server {
 	s := &Server{}
-	v := core.MustFromContext(ctx)
-	v.RequireFeatures([]interface{}{routing.DispatcherType()}, func(fs []features.Feature) {
-		s.dispatcher = fs[0].(routing.Dispatcher)
+	core.RequireFeatures(ctx, func(d routing.Dispatcher) {
+		s.dispatcher = d
 	})
 	return s
 }
diff --git a/app/router/router.go b/app/router/router.go
index e800d8cf..51f8cd61 100644
--- a/app/router/router.go
+++ b/app/router/router.go
@@ -9,7 +9,6 @@ import (
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/common/session"
-	"v2ray.com/core/features"
 	"v2ray.com/core/features/dns"
 	"v2ray.com/core/features/routing"
 	"v2ray.com/core/proxy"
@@ -38,10 +37,10 @@ func NewRouter(ctx context.Context, config *Config) (*Router, error) {
 		r.rules[idx].Condition = cond
 	}
 
-	v := core.MustFromContext(ctx)
-	v.RequireFeatures([]interface{}{dns.ClientType()}, func(fs []features.Feature) {
-		r.dns = fs[0].(dns.Client)
+	core.RequireFeatures(ctx, func(d dns.Client) {
+		r.dns = d
 	})
+
 	return r, nil
 }
 
diff --git a/app/stats/command/command.go b/app/stats/command/command.go
index 24bea2d7..31c9301e 100644
--- a/app/stats/command/command.go
+++ b/app/stats/command/command.go
@@ -75,17 +75,21 @@ func (s *statsServer) QueryStats(ctx context.Context, request *QueryStatsRequest
 }
 
 type service struct {
-	v *core.Instance
+	statsManager feature_stats.Manager
 }
 
 func (s *service) Register(server *grpc.Server) {
-	f := s.v.GetFeature(feature_stats.ManagerType())
-	RegisterStatsServiceServer(server, NewStatsServer(f.(feature_stats.Manager)))
+	RegisterStatsServiceServer(server, NewStatsServer(s.statsManager))
 }
 
 func init() {
 	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, cfg interface{}) (interface{}, error) {
-		s := core.MustFromContext(ctx)
-		return &service{v: s}, nil
+		s := new(service)
+
+		core.RequireFeatures(ctx, func(sm feature_stats.Manager) {
+			s.statsManager = sm
+		})
+
+		return s, nil
 	}))
 }
diff --git a/features/stats/errors.generated.go b/features/stats/errors.generated.go
new file mode 100644
index 00000000..79ec4cf1
--- /dev/null
+++ b/features/stats/errors.generated.go
@@ -0,0 +1,6 @@
+package stats
+
+import "v2ray.com/core/common/errors"
+
+type errPathObjHolder struct {}
+func newError(values ...interface{}) *errors.Error { return errors.New(values...).WithPathObj(errPathObjHolder{}) }
diff --git a/features/stats/stats.go b/features/stats/stats.go
index 67bf022d..8b966c82 100644
--- a/features/stats/stats.go
+++ b/features/stats/stats.go
@@ -1,5 +1,7 @@
 package stats
 
+//go:generate errorgen
+
 import "v2ray.com/core/features"
 
 // Counter is the interface for stats counters.
@@ -36,3 +38,27 @@ func GetOrRegisterCounter(m Manager, name string) (Counter, error) {
 func ManagerType() interface{} {
 	return (*Manager)(nil)
 }
+
+// NoopManager is an implementation of Manager, which doesn't has actual functionalities.
+type NoopManager struct{}
+
+// Type implements common.HasType.
+func (NoopManager) Type() interface{} {
+	return ManagerType()
+}
+
+// RegisterCounter implements Manager.
+func (NoopManager) RegisterCounter(string) (Counter, error) {
+	return nil, newError("not implemented")
+}
+
+// GetCounter implements Manager.
+func (NoopManager) GetCounter(string) Counter {
+	return nil
+}
+
+// Start implements common.Runnable.
+func (NoopManager) Start() error { return nil }
+
+// Close implements common.Closable.
+func (NoopManager) Close() error { return nil }
diff --git a/proxy/freedom/freedom.go b/proxy/freedom/freedom.go
index 4baaee82..4742c6f6 100644
--- a/proxy/freedom/freedom.go
+++ b/proxy/freedom/freedom.go
@@ -6,8 +6,6 @@ import (
 	"context"
 	"time"
 
-	"v2ray.com/core/features"
-
 	"v2ray.com/core"
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/buf"
@@ -37,10 +35,9 @@ func New(ctx context.Context, config *Config) (*Handler, error) {
 		config: *config,
 	}
 
-	v := core.MustFromContext(ctx)
-	v.RequireFeatures([]interface{}{policy.ManagerType(), dns.ClientType()}, func(fs []features.Feature) {
-		f.policyManager = fs[0].(policy.Manager)
-		f.dns = fs[1].(dns.Client)
+	core.RequireFeatures(ctx, func(pm policy.Manager, d dns.Client) {
+		f.policyManager = pm
+		f.dns = d
 	})
 
 	return f, nil
diff --git a/v2ray.go b/v2ray.go
index e52bd238..1b98f7f6 100755
--- a/v2ray.go
+++ b/v2ray.go
@@ -2,6 +2,7 @@ package core
 
 import (
 	"context"
+	"reflect"
 	"sync"
 
 	"v2ray.com/core/common"
@@ -12,6 +13,7 @@ import (
 	"v2ray.com/core/features/outbound"
 	"v2ray.com/core/features/policy"
 	"v2ray.com/core/features/routing"
+	"v2ray.com/core/features/stats"
 )
 
 // Server is an instance of V2Ray. At any time, there must be at most one Server instance running.
@@ -25,13 +27,13 @@ func ServerType() interface{} {
 }
 
 type resolution struct {
-	deps     []interface{}
-	callback func([]features.Feature)
+	deps     []reflect.Type
+	callback interface{}
 }
 
-func getFeature(allFeatures []features.Feature, t interface{}) features.Feature {
+func getFeature(allFeatures []features.Feature, t reflect.Type) features.Feature {
 	for _, f := range allFeatures {
-		if f.Type() == t {
+		if reflect.TypeOf(f.Type()) == t {
 			return f
 		}
 	}
@@ -48,7 +50,25 @@ func (r *resolution) resolve(allFeatures []features.Feature) bool {
 		fs = append(fs, f)
 	}
 
-	r.callback(fs)
+	callback := reflect.ValueOf(r.callback)
+	var input []reflect.Value
+	callbackType := callback.Type()
+	for i := 0; i < callbackType.NumIn(); i++ {
+		pt := callbackType.In(i)
+		for _, f := range fs {
+			if reflect.TypeOf(f).AssignableTo(pt) {
+				input = append(input, reflect.ValueOf(f))
+				break
+			}
+		}
+	}
+
+	if len(input) != callbackType.NumIn() {
+		panic("Can't get all input parameters")
+	}
+
+	callback.Call(input)
+
 	return true
 }
 
@@ -112,6 +132,13 @@ func addOutboundHandlers(server *Instance, configs []*OutboundHandlerConfig) err
 	return nil
 }
 
+// RequireFeatures is a helper function to require features from Instance in context.
+// See Instance.RequireFeatures for more information.
+func RequireFeatures(ctx context.Context, callback interface{}) {
+	v := MustFromContext(ctx)
+	v.RequireFeatures(callback)
+}
+
 // New returns a new V2Ray instance based on given configuration.
 // The instance is not started at this point.
 // To ensure V2Ray instance works properly, the config must contain one Dispatcher, one InboundHandlerManager and one OutboundHandlerManager. Other features are optional.
@@ -151,6 +178,10 @@ func New(config *Config) (*Instance, error) {
 		server.AddFeature(routing.DefaultRouter{})
 	}
 
+	if server.GetFeature(stats.ManagerType()) == nil {
+		server.AddFeature(stats.NoopManager{})
+	}
+
 	// Add an empty instance at the end, for optional feature requirement.
 	server.AddFeature(&Instance{})
 
@@ -195,7 +226,18 @@ func (s *Instance) Close() error {
 }
 
 // RequireFeatures registers a callback, which will be called when all dependent features are registered.
-func (s *Instance) RequireFeatures(featureTypes []interface{}, callback func([]features.Feature)) {
+// The callback must be a func(). All its parameters must be features.Feature.
+func (s *Instance) RequireFeatures(callback interface{}) {
+	callbackType := reflect.TypeOf(callback)
+	if callbackType.Kind() != reflect.Func {
+		panic("not a function")
+	}
+
+	var featureTypes []reflect.Type
+	for i := 0; i < callbackType.NumIn(); i++ {
+		featureTypes = append(featureTypes, reflect.PtrTo(callbackType.In(i)))
+	}
+
 	r := resolution{
 		deps:     featureTypes,
 		callback: callback,
@@ -236,7 +278,7 @@ func (s *Instance) AddFeature(feature features.Feature) {
 
 // GetFeature returns a feature of the given type, or nil if such feature is not registered.
 func (s *Instance) GetFeature(featureType interface{}) features.Feature {
-	return getFeature(s.features, featureType)
+	return getFeature(s.features, reflect.TypeOf(featureType))
 }
 
 // Start starts the V2Ray instance, including all registered features. When Start returns error, the state of the instance is unknown.
diff --git a/v2ray_test.go b/v2ray_test.go
index 934491af..ef6577c5 100644
--- a/v2ray_test.go
+++ b/v2ray_test.go
@@ -13,12 +13,27 @@ import (
 	"v2ray.com/core/common/protocol"
 	"v2ray.com/core/common/serial"
 	"v2ray.com/core/common/uuid"
+	"v2ray.com/core/features/dns"
 	_ "v2ray.com/core/main/distro/all"
 	"v2ray.com/core/proxy/dokodemo"
 	"v2ray.com/core/proxy/vmess"
 	"v2ray.com/core/proxy/vmess/outbound"
 )
 
+func TestV2RayDependency(t *testing.T) {
+	instance := new(Instance)
+
+	wait := make(chan bool, 1)
+	instance.RequireFeatures(func(d dns.Client) {
+		if d == nil {
+			t.Error("expected dns client fulfilled, but actually nil")
+		}
+		wait <- true
+	})
+	instance.AddFeature(dns.LocalClient{})
+	<-wait
+}
+
 func TestV2RayClose(t *testing.T) {
 	port := net.Port(dice.RollUint16())
 	userId := uuid.New()