Export instance context

core/context.go

@@ -25,3 +25,10 @@ func MustFromContext(ctx context.Context) *Instance {
 	}
 	return x
 }
+
+func WithContext(ctx context.Context, v *Instance) context.Context {
+	if FromContext(ctx) != v {
+		ctx = context.WithValue(ctx, xrayKey, v)
+	}
+	return ctx
+}

core/functions.go

@@ -15,7 +15,7 @@ import (
 func CreateObject(v *Instance, config interface{}) (interface{}, error) {
 	ctx := v.ctx
 	if v != nil {
-		ctx = context.WithValue(ctx, xrayKey, v)
+		ctx = WithContext(v.ctx, v)
 	}
 	return common.CreateObject(ctx, config)
 }
@@ -46,6 +46,8 @@ func StartInstance(configFormat string, configBytes []byte) (*Instance, error) {
 //
 // xray:api:stable
 func Dial(ctx context.Context, v *Instance, dest net.Destination) (net.Conn, error) {
+	ctx = WithContext(ctx, v)
+
 	dispatcher := v.GetFeature(routing.DispatcherType())
 	if dispatcher == nil {
 		return nil, newError("routing.Dispatcher is not registered in Xray core")
@@ -70,6 +72,8 @@ func Dial(ctx context.Context, v *Instance, dest net.Destination) (net.Conn, err
 //
 // xray:api:beta
 func DialUDP(ctx context.Context, v *Instance) (net.PacketConn, error) {
+	ctx = WithContext(ctx, v)
+
 	dispatcher := v.GetFeature(routing.DispatcherType())
 	if dispatcher == nil {
 		return nil, newError("routing.Dispatcher is not registered in Xray core")