From 3da6cefbcacaca3fc4c77dce5ae9952506b2e5f3 Mon Sep 17 00:00:00 2001
From: yuhan6665 <1588741+yuhan6665@users.noreply.github.com>
Date: Sat, 13 Apr 2024 16:04:58 -0400
Subject: [PATCH] Add more traffic info for future use

---
 proxy/proxy.go | 70 ++++++++++++++++++++++++++++++++------------------
 1 file changed, 45 insertions(+), 25 deletions(-)

diff --git a/proxy/proxy.go b/proxy/proxy.go
index 1e4c69f5..7281a205 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -101,6 +101,11 @@ type GetOutbound interface {
 // It is used by XTLS to determine if switch to raw copy mode, It is used by Vision to calculate padding
 type TrafficState struct {
 	UserUUID               []byte
+	StartTime              time.Time
+	ByteSent               int64
+	ByteReceived           int64      
+	NumberOfPacketSent     int
+	NumberOfPacketReceived int
 	NumberOfPacketToFilter int
 	EnableXtls             bool
 	IsTLS12orAbove         bool
@@ -124,6 +129,11 @@ type TrafficState struct {
 func NewTrafficState(userUUID []byte) *TrafficState {
 	return &TrafficState{
 		UserUUID:                 userUUID,
+		StartTime:                time.Time{},
+		ByteSent:                 0,
+		ByteReceived:             0,
+		NumberOfPacketSent:       0,
+		NumberOfPacketReceived:   0,
 		NumberOfPacketToFilter:   8,
 		EnableXtls:               false,
 		IsTLS12orAbove:           false,
@@ -160,6 +170,10 @@ func NewVisionReader(reader buf.Reader, state *TrafficState, context context.Con
 func (w *VisionReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
 	buffer, err := w.Reader.ReadMultiBuffer()
 	if !buffer.IsEmpty() {
+		if w.trafficState.StartTime.IsZero() {
+			w.trafficState.StartTime = time.Now()
+		}
+		w.trafficState.ByteReceived += int64(buffer.Len())
 		if w.trafficState.WithinPaddingBuffers || w.trafficState.NumberOfPacketToFilter > 0 {
 			mb2 := make(buf.MultiBuffer, 0, len(buffer))
 			for _, b := range buffer {
@@ -180,6 +194,7 @@ func (w *VisionReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
 				errors.LogInfo(w.ctx, "XtlsRead unknown command ", w.trafficState.CurrentCommand, buffer.Len())
 			}
 		}
+		w.trafficState.NumberOfPacketReceived += len(buffer)
 		if w.trafficState.NumberOfPacketToFilter > 0 {
 			XtlsFilterTls(buffer, w.trafficState, w.ctx)
 		}
@@ -208,47 +223,52 @@ func NewVisionWriter(writer buf.Writer, state *TrafficState, context context.Con
 }
 
 func (w *VisionWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
+	w.trafficState.NumberOfPacketSent += len(mb)
 	if w.trafficState.NumberOfPacketToFilter > 0 {
 		XtlsFilterTls(mb, w.trafficState, w.ctx)
 	}
 	if w.trafficState.IsPadding {
 		if len(mb) == 1 && mb[0] == nil {
 			mb[0] = XtlsPadding(nil, CommandPaddingContinue, &w.writeOnceUserUUID, true, w.ctx) // we do a long padding to hide vless header
-			return w.Writer.WriteMultiBuffer(mb)
-		}
-		mb = ReshapeMultiBuffer(w.ctx, mb)
-		longPadding := w.trafficState.IsTLS
-		for i, b := range mb {
-			if w.trafficState.IsTLS && b.Len() >= 6 && bytes.Equal(TlsApplicationDataStart, b.BytesTo(3)) {
-				if w.trafficState.EnableXtls {
-					w.trafficState.WriterSwitchToDirectCopy = true
+		} else {
+			mb = ReshapeMultiBuffer(w.ctx, mb)
+			longPadding := w.trafficState.IsTLS
+			for i, b := range mb {
+				if w.trafficState.IsTLS && b.Len() >= 6 && bytes.Equal(TlsApplicationDataStart, b.BytesTo(3)) {
+					if w.trafficState.EnableXtls {
+						w.trafficState.WriterSwitchToDirectCopy = true
+					}
+					var command byte = CommandPaddingContinue
+					if i == len(mb) - 1 {
+						command = CommandPaddingEnd
+						if w.trafficState.EnableXtls {
+							command = CommandPaddingDirect
+						}
+					}
+					mb[i] = XtlsPadding(b, command, &w.writeOnceUserUUID, true, w.ctx)
+					w.trafficState.IsPadding = false // padding going to end
+					longPadding = false
+					continue
+				} else if !w.trafficState.IsTLS12orAbove && w.trafficState.NumberOfPacketToFilter <= 1 { // For compatibility with earlier vision receiver, we finish padding 1 packet early
+					w.trafficState.IsPadding = false
+					mb[i] = XtlsPadding(b, CommandPaddingEnd, &w.writeOnceUserUUID, longPadding, w.ctx)
+					break
 				}
 				var command byte = CommandPaddingContinue
-				if i == len(mb)-1 {
+				if i == len(mb) - 1 && !w.trafficState.IsPadding {
 					command = CommandPaddingEnd
 					if w.trafficState.EnableXtls {
 						command = CommandPaddingDirect
 					}
 				}
-				mb[i] = XtlsPadding(b, command, &w.writeOnceUserUUID, true, w.ctx)
-				w.trafficState.IsPadding = false // padding going to end
-				longPadding = false
-				continue
-			} else if !w.trafficState.IsTLS12orAbove && w.trafficState.NumberOfPacketToFilter <= 1 { // For compatibility with earlier vision receiver, we finish padding 1 packet early
-				w.trafficState.IsPadding = false
-				mb[i] = XtlsPadding(b, CommandPaddingEnd, &w.writeOnceUserUUID, longPadding, w.ctx)
-				break
+				mb[i] = XtlsPadding(b, command, &w.writeOnceUserUUID, longPadding, w.ctx)
 			}
-			var command byte = CommandPaddingContinue
-			if i == len(mb)-1 && !w.trafficState.IsPadding {
-				command = CommandPaddingEnd
-				if w.trafficState.EnableXtls {
-					command = CommandPaddingDirect
-				}
-			}
-			mb[i] = XtlsPadding(b, command, &w.writeOnceUserUUID, longPadding, w.ctx)
 		}
 	}
+	w.trafficState.ByteSent += int64(mb.Len())
+	if w.trafficState.StartTime.IsZero() {
+		w.trafficState.StartTime = time.Now()
+	}
 	return w.Writer.WriteMultiBuffer(mb)
 }