From 496b2c02c5431fb34cda77b7a4920838d6ea5d83 Mon Sep 17 00:00:00 2001
From: yuhan6665 <1588741+yuhan6665@users.noreply.github.com>
Date: Sat, 19 Feb 2022 22:45:41 -0500
Subject: [PATCH] Fix context api and document (#727)

* Fix Dial function crash instance when there is no instance context in the ctx

* check ctx to fix. (#841)

* Feat: core.ToContext(ctx, v) for ctx initialization (#852)

* remove exported API: toContext

* Remove unnecessary API

* rework document for API

* fix: make sure the ctx is propagated to connections by detached connection

Co-authored-by: Shelikhoo <xiaokangwang@outlook.com>
Co-authored-by: rurirei <72071920+rurirei@users.noreply.github.com>
---
 app/dns/nameserver_udp.go |  3 ++-
 core/context.go           | 25 +++++++++++++++++++++++++
 core/context_test.go      |  3 ++-
 core/functions.go         |  7 ++++++-
 4 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/app/dns/nameserver_udp.go b/app/dns/nameserver_udp.go
index d68ffdd0..36747016 100644
--- a/app/dns/nameserver_udp.go
+++ b/app/dns/nameserver_udp.go
@@ -17,6 +17,7 @@ import (
 	"github.com/xtls/xray-core/common/session"
 	"github.com/xtls/xray-core/common/signal/pubsub"
 	"github.com/xtls/xray-core/common/task"
+	"github.com/xtls/xray-core/core"
 	dns_feature "github.com/xtls/xray-core/features/dns"
 	"github.com/xtls/xray-core/features/routing"
 	"github.com/xtls/xray-core/transport/internet/udp"
@@ -195,7 +196,7 @@ func (s *ClassicNameServer) sendQuery(ctx context.Context, domain string, client
 	for _, req := range reqs {
 		s.addPendingRequest(req)
 		b, _ := dns.PackMessage(req.msg)
-		udpCtx := context.Background()
+		udpCtx := core.ToBackgroundDetachedContext(ctx)
 		if inbound := session.InboundFromContext(ctx); inbound != nil {
 			udpCtx = session.ContextWithInbound(udpCtx, inbound)
 		}
diff --git a/core/context.go b/core/context.go
index 645ed277..d28ac7e3 100644
--- a/core/context.go
+++ b/core/context.go
@@ -25,3 +25,28 @@ func MustFromContext(ctx context.Context) *Instance {
 	}
 	return x
 }
+
+/* toContext returns ctx from the given context, or creates an Instance if the context doesn't find that.
+
+It is unsupported to use this function to create a context that is suitable to invoke Xray's internal component
+in third party code, you shouldn't use //go:linkname to alias of this function into your own package and
+use this function in your third party code.
+
+For third party code, usage enabled by creating a context to interact with Xray's internal component is unsupported,
+and may break at any time.
+
+*/
+func toContext(ctx context.Context, v *Instance) context.Context {
+	if FromContext(ctx) != v {
+		ctx = context.WithValue(ctx, xrayKey, v)
+	}
+	return ctx
+}
+
+/*ToBackgroundDetachedContext create a detached context from another context
+Internal API
+*/
+func ToBackgroundDetachedContext(ctx context.Context) context.Context {
+	instance := MustFromContext(ctx)
+	return toContext(context.Background(), instance)
+}
diff --git a/core/context_test.go b/core/context_test.go
index a16f6120..423585e1 100644
--- a/core/context_test.go
+++ b/core/context_test.go
@@ -5,9 +5,10 @@ import (
 	"testing"
 
 	. "github.com/xtls/xray-core/core"
+	_ "unsafe"
 )
 
-func TestContextPanic(t *testing.T) {
+func TestFromContextPanic(t *testing.T) {
 	defer func() {
 		r := recover()
 		if r == nil {
diff --git a/core/functions.go b/core/functions.go
index 11282462..7885fd14 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 = toContext(v.ctx, v)
 	}
 	return common.CreateObject(ctx, config)
 }
@@ -46,10 +46,13 @@ 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 = toContext(ctx, v)
+
 	dispatcher := v.GetFeature(routing.DispatcherType())
 	if dispatcher == nil {
 		return nil, newError("routing.Dispatcher is not registered in Xray core")
 	}
+
 	r, err := dispatcher.(routing.Dispatcher).Dispatch(ctx, dest)
 	if err != nil {
 		return nil, err
@@ -70,6 +73,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 = toContext(ctx, v)
+
 	dispatcher := v.GetFeature(routing.DispatcherType())
 	if dispatcher == nil {
 		return nil, newError("routing.Dispatcher is not registered in Xray core")