From a8559a1b465909ef4d7ce222361c316bcaaa7f4f Mon Sep 17 00:00:00 2001
From: RPRX <63339210+RPRX@users.noreply.github.com>
Date: Sat, 28 Dec 2024 04:06:49 +0000
Subject: [PATCH] Dokodemo TPROXY: Interrupt UDP download if upload timeouts

https://github.com/XTLS/Xray-core/issues/4194#issuecomment-2564156340

Fixes https://github.com/XTLS/Xray-core/issues/4194
---
 proxy/dokodemo/dokodemo.go | 27 ++++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/proxy/dokodemo/dokodemo.go b/proxy/dokodemo/dokodemo.go
index bbeab18f..6522f071 100644
--- a/proxy/dokodemo/dokodemo.go
+++ b/proxy/dokodemo/dokodemo.go
@@ -2,6 +2,7 @@ package dokodemo
 
 import (
 	"context"
+	"runtime"
 	"sync/atomic"
 
 	"github.com/xtls/xray-core/common"
@@ -147,10 +148,6 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn st
 		return nil
 	}
 
-	tproxyRequest := func() error {
-		return nil
-	}
-
 	var writer buf.Writer
 	if network == net.Network_TCP {
 		writer = buf.NewWriter(conn)
@@ -180,7 +177,12 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn st
 				return err
 			}
 			writer = NewPacketWriter(pConn, &dest, mark, back)
-			defer writer.(*PacketWriter).Close()
+			defer func() {
+				runtime.Gosched()
+				common.Interrupt(link.Reader) // maybe duplicated
+				runtime.Gosched()
+				writer.(*PacketWriter).Close() // close fake UDP conns
+			}()
 			/*
 				sockopt := &internet.SocketConfig{
 					Tproxy: internet.SocketConfig_TProxy,
@@ -219,17 +221,24 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn st
 	responseDone := func() error {
 		defer timer.SetTimeout(plcy.Timeouts.UplinkOnly)
 
+		if network == net.Network_UDP && destinationOverridden {
+			buf.Copy(link.Reader, writer) // respect upload's timeout
+			return nil
+		}
+
 		if err := buf.Copy(link.Reader, writer, buf.UpdateActivity(timer)); err != nil {
 			return errors.New("failed to transport response").Base(err)
 		}
 		return nil
 	}
 
-	if err := task.Run(ctx, task.OnSuccess(func() error {
-		return task.Run(ctx, requestDone, tproxyRequest)
-	}, task.Close(link.Writer)), responseDone); err != nil {
-		common.Interrupt(link.Reader)
+	if err := task.Run(ctx,
+		task.OnSuccess(func() error { return task.Run(ctx, requestDone) }, task.Close(link.Writer)),
+		responseDone); err != nil {
+		runtime.Gosched()
 		common.Interrupt(link.Writer)
+		runtime.Gosched()
+		common.Interrupt(link.Reader)
 		return errors.New("connection ends").Base(err)
 	}