From a26534c79c3518d9238af128dc74c88078547afe Mon Sep 17 00:00:00 2001
From: Shelikhoo <xiaokangwang@outlook.com>
Date: Thu, 18 Jun 2020 11:19:05 +0800
Subject: [PATCH] Allow client to submit a context in v2ray.New

---
 v2ray.go | 47 ++++++++++++++++++++++++++++++++++-------------
 1 file changed, 34 insertions(+), 13 deletions(-)

diff --git a/v2ray.go b/v2ray.go
index 9b47b368..b9e868d9 100755
--- a/v2ray.go
+++ b/v2ray.go
@@ -92,6 +92,8 @@ type Instance struct {
 	features           []features.Feature
 	featureResolutions []resolution
 	running            bool
+
+	ctx context.Context
 }
 
 func AddInboundHandler(server *Instance, config *InboundHandlerConfig) error {
@@ -104,7 +106,7 @@ func AddInboundHandler(server *Instance, config *InboundHandlerConfig) error {
 	if !ok {
 		return newError("not an InboundHandler")
 	}
-	if err := inboundManager.AddHandler(context.Background(), handler); err != nil {
+	if err := inboundManager.AddHandler(server.ctx, handler); err != nil {
 		return err
 	}
 	return nil
@@ -130,7 +132,7 @@ func AddOutboundHandler(server *Instance, config *OutboundHandlerConfig) error {
 	if !ok {
 		return newError("not an OutboundHandler")
 	}
-	if err := outboundManager.AddHandler(context.Background(), handler); err != nil {
+	if err := outboundManager.AddHandler(server.ctx, handler); err != nil {
 		return err
 	}
 	return nil
@@ -157,27 +159,47 @@ func RequireFeatures(ctx context.Context, callback interface{}) error {
 // 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.
 func New(config *Config) (*Instance, error) {
-	var server = &Instance{}
+	var server = &Instance{ctx: context.Background()}
+
+	err, done := initInstanceWithConfig(config, server)
+	if done {
+		return nil, err
+	}
+
+	return server, nil
+}
+
+func NewWithContext(config *Config, ctx context.Context) (*Instance, error) {
+	var server = &Instance{ctx: ctx}
 
+	err, done := initInstanceWithConfig(config, server)
+	if done {
+		return nil, err
+	}
+
+	return server, nil
+}
+
+func initInstanceWithConfig(config *Config, server *Instance) (error, bool) {
 	if config.Transport != nil {
 		features.PrintDeprecatedFeatureWarning("global transport settings")
 	}
 	if err := config.Transport.Apply(); err != nil {
-		return nil, err
+		return err, true
 	}
 
 	for _, appSettings := range config.App {
 		settings, err := appSettings.GetInstance()
 		if err != nil {
-			return nil, err
+			return err, true
 		}
 		obj, err := CreateObject(server, settings)
 		if err != nil {
-			return nil, err
+			return err, true
 		}
 		if feature, ok := obj.(features.Feature); ok {
 			if err := server.AddFeature(feature); err != nil {
-				return nil, err
+				return err, true
 			}
 		}
 	}
@@ -195,24 +217,23 @@ func New(config *Config) (*Instance, error) {
 	for _, f := range essentialFeatures {
 		if server.GetFeature(f.Type) == nil {
 			if err := server.AddFeature(f.Instance); err != nil {
-				return nil, err
+				return err, true
 			}
 		}
 	}
 
 	if server.featureResolutions != nil {
-		return nil, newError("not all dependency are resolved.")
+		return newError("not all dependency are resolved."), true
 	}
 
 	if err := addInboundHandlers(server, config.Inbound); err != nil {
-		return nil, err
+		return err, true
 	}
 
 	if err := addOutboundHandlers(server, config.Outbound); err != nil {
-		return nil, err
+		return err, true
 	}
-
-	return server, nil
+	return nil, false
 }
 
 // Type implements common.HasType.