From 76b27a37cb9d1f82baef901949d3fabc03fa63fa Mon Sep 17 00:00:00 2001
From: RPRX <63339210+RPRX@users.noreply.github.com>
Date: Fri, 7 Apr 2023 19:13:20 +0800
Subject: [PATCH] Update common/xudp/xudp.go and common/mux/server.go

---
 common/mux/server.go |  3 +++
 common/xudp/xudp.go  | 18 +++++++++---------
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/common/mux/server.go b/common/mux/server.go
index e64e038f..b3c22a61 100644
--- a/common/mux/server.go
+++ b/common/mux/server.go
@@ -177,6 +177,9 @@ func (w *ServerWorker) handleStatusNew(ctx context.Context, meta *FrameMetadata,
 			// Actually, it won't return an error in Xray-core's implementations.
 			link, err := w.dispatcher.Dispatch(ctx, meta.Target)
 			if err != nil {
+				XUDPManager.Lock()
+				delete(XUDPManager.Map, x.GlobalID)
+				XUDPManager.Unlock()
 				err = newError("failed to dispatch request to ", meta.Target).Base(err)
 				if xudp.Show {
 					fmt.Printf("XUDP new: %v err: %v\n", meta.GlobalID, err)
diff --git a/common/xudp/xudp.go b/common/xudp/xudp.go
index 65096d16..32dda614 100644
--- a/common/xudp/xudp.go
+++ b/common/xudp/xudp.go
@@ -25,7 +25,7 @@ var AddrParser = protocol.NewAddressParser(
 
 var (
 	Show    bool
-	BaseKey [32]byte
+	BaseKey []byte
 )
 
 const (
@@ -37,24 +37,24 @@ func init() {
 	if strings.ToLower(os.Getenv(EnvShow)) == "true" {
 		Show = true
 	}
-	if raw := os.Getenv(EnvBaseKey); raw != "" {
-		if key, _ := base64.RawURLEncoding.DecodeString(raw); len(key) == len(BaseKey) {
-			copy(BaseKey[:], key)
+	if raw, found := os.LookupEnv(EnvBaseKey); found {
+		if BaseKey, _ = base64.RawURLEncoding.DecodeString(raw); len(BaseKey) == 32 {
 			return
-		} else {
-			panic(EnvBaseKey + ": invalid value: " + raw)
 		}
+		panic(EnvBaseKey + ": invalid value: " + raw)
 	}
-	rand.Read(BaseKey[:])
+	rand.Read(BaseKey)
 }
 
 func GetGlobalID(ctx context.Context) (globalID [8]byte) {
 	if inbound := session.InboundFromContext(ctx); inbound != nil && inbound.Source.Network == net.Network_UDP &&
 		(inbound.Name == "dokodemo-door" || inbound.Name == "socks" || inbound.Name == "shadowsocks") {
-		h := blake3.New(8, BaseKey[:])
+		h := blake3.New(8, BaseKey)
 		h.Write([]byte(inbound.Source.String()))
 		copy(globalID[:], h.Sum(nil))
-		fmt.Printf("XUDP inbound.Source.String(): %v\tglobalID: %v\n", inbound.Source.String(), globalID)
+		if Show {
+			fmt.Printf("XUDP inbound.Source.String(): %v\tglobalID: %v\n", inbound.Source.String(), globalID)
+		}
 	}
 	return
 }