diff --git a/core/context.go b/core/context.go index 645ed277..30be55ed 100644 --- a/core/context.go +++ b/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 +} diff --git a/core/functions.go b/core/functions.go index 11282462..0188ada7 100644 --- a/core/functions.go +++ b/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")