From b8c5a78f2e8c23ee47fd0705102146e744b8adab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E7=A5=A5=E8=B6=85?= <iwind.liu@gmail.com>
Date: Wed, 3 Aug 2022 16:09:47 +0800
Subject: [PATCH] =?UTF-8?q?TCP/TLS/UDP=E7=AC=AC=E4=B8=80=E6=AC=A1=E8=BF=9E?=
 =?UTF-8?q?=E6=8E=A5=E6=BA=90=E7=AB=99=E5=A4=B1=E8=B4=A5=E5=90=8E=EF=BC=8C?=
 =?UTF-8?q?=E8=87=AA=E5=8A=A8=E5=B0=9D=E8=AF=95=E4=B8=8B=E4=B8=80=E4=B8=AA?=
 =?UTF-8?q?=E6=BA=90=E7=AB=99?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/nodes/listener_tcp.go | 18 ++++++++++++++++--
 internal/nodes/listener_udp.go | 30 ++++++++++++++++++++++++++++--
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/internal/nodes/listener_tcp.go b/internal/nodes/listener_tcp.go
index fbfbb54..0bd93ec 100644
--- a/internal/nodes/listener_tcp.go
+++ b/internal/nodes/listener_tcp.go
@@ -231,8 +231,17 @@ func (this *TCPListener) connectOrigin(serverId int64, requestHost string, rever
 
 	var retries = 3
 	var addr string
+
+	var failedOriginIds []int64
+
 	for i := 0; i < retries; i++ {
-		var origin = reverseProxy.NextOrigin(requestCall)
+		var origin *serverconfigs.OriginConfig
+		if len(failedOriginIds) > 0 {
+			origin = reverseProxy.AnyOrigin(requestCall, failedOriginIds)
+		}
+		if origin == nil {
+			origin = reverseProxy.NextOrigin(requestCall)
+		}
 		if origin == nil {
 			continue
 		}
@@ -246,6 +255,8 @@ func (this *TCPListener) connectOrigin(serverId int64, requestHost string, rever
 
 		conn, addr, err = OriginConnect(origin, this.port, remoteAddr, requestHost)
 		if err != nil {
+			failedOriginIds = append(failedOriginIds, origin.Id)
+
 			remotelogs.ServerError(serverId, "TCP_LISTENER", "unable to connect origin server: "+addr+": "+err.Error(), "", nil)
 
 			SharedOriginStateManager.Fail(origin, requestHost, reverseProxy, func() {
@@ -263,7 +274,10 @@ func (this *TCPListener) connectOrigin(serverId int64, requestHost string, rever
 			return
 		}
 	}
-	err = errors.New("server '" + types.String(serverId) + "': no available origin server can be used")
+
+	if err == nil {
+		err = errors.New("server '" + types.String(serverId) + "': no available origin server can be used")
+	}
 	return
 }
 
diff --git a/internal/nodes/listener_udp.go b/internal/nodes/listener_udp.go
index 491c73d..aedf025 100644
--- a/internal/nodes/listener_udp.go
+++ b/internal/nodes/listener_udp.go
@@ -139,16 +139,39 @@ func (this *UDPListener) connectOrigin(serverId int64, reverseProxy *serverconfi
 
 	var retries = 3
 	var addr string
+
+	var failedOriginIds []int64
+
 	for i := 0; i < retries; i++ {
-		var origin = reverseProxy.NextOrigin(nil)
+		var origin *serverconfigs.OriginConfig
+		if len(failedOriginIds) > 0 {
+			origin = reverseProxy.AnyOrigin(nil, failedOriginIds)
+		}
+		if origin == nil {
+			origin = reverseProxy.NextOrigin(nil)
+		}
 		if origin == nil {
 			continue
 		}
+
 		conn, addr, err = OriginConnect(origin, this.port, remoteAddr.String(), "")
 		if err != nil {
+			failedOriginIds = append(failedOriginIds, origin.Id)
+
 			remotelogs.ServerError(serverId, "UDP_LISTENER", "unable to connect origin server: "+addr+": "+err.Error(), "", nil)
+
+			SharedOriginStateManager.Fail(origin, "", reverseProxy, func() {
+				reverseProxy.ResetScheduling()
+			})
+
 			continue
 		} else {
+			if !origin.IsOk {
+				SharedOriginStateManager.Success(origin, func() {
+					reverseProxy.ResetScheduling()
+				})
+			}
+
 			// PROXY Protocol
 			if reverseProxy != nil &&
 				reverseProxy.ProxyProtocol != nil &&
@@ -175,7 +198,10 @@ func (this *UDPListener) connectOrigin(serverId int64, reverseProxy *serverconfi
 			return
 		}
 	}
-	err = errors.New("server '" + types.String(serverId) + "': no available origin server can be used")
+
+	if err == nil {
+		err = errors.New("server '" + types.String(serverId) + "': no available origin server can be used")
+	}
 	return
 }