From f35c485a6317feaeee25da33888b2b5456465e32 Mon Sep 17 00:00:00 2001
From: e1732a364fed <75717694+e1732a364fed@users.noreply.github.com>
Date: Sat, 1 Jan 2000 00:00:00 +0000
Subject: [PATCH] =?UTF-8?q?=E5=AE=9A=E4=B9=89=E5=A5=BD=E6=96=B0=E7=9A=84Re?=
 =?UTF-8?q?layMsg=E5=87=BD=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 netLayer/relay_udp.go | 34 +++++++++++++++++++++++++++++++---
 1 file changed, 31 insertions(+), 3 deletions(-)

diff --git a/netLayer/relay_udp.go b/netLayer/relay_udp.go
index 86d1c0c..b238545 100644
--- a/netLayer/relay_udp.go
+++ b/netLayer/relay_udp.go
@@ -308,18 +308,46 @@ relay udp  有两种针对不同通道的技术
 我们重新定义一个MsgProducer 和 MsgConsumer，就方便很多. 这是针对多来源的转发。
 
 如此，一个 UDPConn就相当于一个 MsgProducer, 它的to 可以由 tproxy 或者 msg内部的数据提取出来
+
+而且 这个模型也可以实现 单来源，所以更实用
 */
-func RelayMsgP2C(p MsgProducer, c MsgConsumer) {
+func RelayMsg(rc, lc MsgHub, downloadByteCount, uploadByteCount *uint64) uint64 {
+	go CopyMsgFromP2C(lc, rc, uploadByteCount)
+
+	var dbc uint64
+	CopyMsgFromP2C(rc, lc, &dbc)
+
+	if downloadByteCount != nil {
+		atomic.AddUint64(downloadByteCount, dbc)
+	}
+	return dbc
+
+}
+
+func CopyMsgFromP2C(p MsgProducer, c MsgConsumer, countPtr *uint64) {
+	var bc uint64
+
 	for {
 		msg, from, to, err := p.ProduceMsg()
 		if err != nil {
-			return
+			break
 		}
 		err = c.ConsumeMsg(msg, from, to)
 		if err != nil {
-			return
+			break
 		}
+		bc += uint64(len(msg))
 	}
+
+	if countPtr != nil {
+		atomic.AddUint64(countPtr, bc)
+	}
+
+}
+
+type MsgHub interface {
+	MsgProducer
+	MsgConsumer
 }
 
 type MsgProducer interface {