From d5aeb6c545ada26415ecb72fea2aa5bf6c63ad79 Mon Sep 17 00:00:00 2001
From: maskedeken <52683904+maskedeken@users.noreply.github.com>
Date: Fri, 8 Jan 2021 11:55:25 +0800
Subject: [PATCH] Refine Trojan packet reader & writer (#142)

---
 proxy/trojan/protocol.go      | 37 +----------------------------------
 proxy/trojan/protocol_test.go | 15 +++++++-------
 proxy/trojan/server.go        | 19 +++++++++++-------
 3 files changed, 21 insertions(+), 50 deletions(-)

diff --git a/proxy/trojan/protocol.go b/proxy/trojan/protocol.go
index c48b9f28..1273b0a2 100644
--- a/proxy/trojan/protocol.go
+++ b/proxy/trojan/protocol.go
@@ -146,26 +146,6 @@ func (w *PacketWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
 	return nil
 }
 
-// WriteMultiBufferWithMetadata writes udp packet with destination specified
-func (w *PacketWriter) WriteMultiBufferWithMetadata(mb buf.MultiBuffer, dest net.Destination) error {
-	for {
-		mb2, b := buf.SplitFirst(mb)
-		mb = mb2
-		if b == nil {
-			break
-		}
-		source := &dest
-		if b.UDP != nil {
-			source = b.UDP
-		}
-		if _, err := w.writePacket(b.Bytes(), *source); err != nil {
-			buf.ReleaseMulti(mb)
-			return err
-		}
-	}
-	return nil
-}
-
 func (w *PacketWriter) writePacket(payload []byte, dest net.Destination) (int, error) {
 	buffer := buf.StackNew()
 	defer buffer.Release()
@@ -259,12 +239,6 @@ func (c *ConnReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
 	return buf.MultiBuffer{b}, err
 }
 
-// PacketPayload combines udp payload and destination
-type PacketPayload struct {
-	Target net.Destination
-	Buffer buf.MultiBuffer
-}
-
 // PacketReader is UDP Connection Reader Wrapper for trojan protocol
 type PacketReader struct {
 	io.Reader
@@ -272,15 +246,6 @@ type PacketReader struct {
 
 // ReadMultiBuffer implements buf.Reader
 func (r *PacketReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
-	p, err := r.ReadMultiBufferWithMetadata()
-	if p != nil {
-		return p.Buffer, err
-	}
-	return nil, err
-}
-
-// ReadMultiBufferWithMetadata reads udp packet with destination
-func (r *PacketReader) ReadMultiBufferWithMetadata() (*PacketPayload, error) {
 	addr, port, err := addrParser.ReadAddressPort(nil, r)
 	if err != nil {
 		return nil, newError("failed to read address and port").Base(err)
@@ -321,7 +286,7 @@ func (r *PacketReader) ReadMultiBufferWithMetadata() (*PacketPayload, error) {
 		remain -= int(n)
 	}
 
-	return &PacketPayload{Target: dest, Buffer: mb}, nil
+	return mb, nil
 }
 
 func ReadV(reader buf.Reader, writer buf.Writer, timer signal.ActivityUpdater, conn *xtls.Conn, rawConn syscall.RawConn, counter stats.Counter, sctx context.Context) error {
diff --git a/proxy/trojan/protocol_test.go b/proxy/trojan/protocol_test.go
index ea1cf7ac..038f45fd 100644
--- a/proxy/trojan/protocol_test.go
+++ b/proxy/trojan/protocol_test.go
@@ -71,21 +71,22 @@ func TestUDPRequest(t *testing.T) {
 	common.Must(connReader.ParseHeader())
 
 	packetReader := &PacketReader{Reader: connReader}
-	p, err := packetReader.ReadMultiBufferWithMetadata()
+	mb, err := packetReader.ReadMultiBuffer()
 	common.Must(err)
 
-	if p.Buffer.IsEmpty() {
+	if mb.IsEmpty() {
 		t.Error("no request data")
 	}
 
-	if r := cmp.Diff(p.Target, destination); r != "" {
+	mb2, b := buf.SplitFirst(mb)
+	defer buf.ReleaseMulti(mb2)
+
+	dest := *b.UDP
+	if r := cmp.Diff(dest, destination); r != "" {
 		t.Error("destination: ", r)
 	}
 
-	mb, decoded := buf.SplitFirst(p.Buffer)
-	buf.ReleaseMulti(mb)
-
-	if r := cmp.Diff(decoded.Bytes(), payload); r != "" {
+	if r := cmp.Diff(b.Bytes(), payload); r != "" {
 		t.Error("data: ", r)
 	}
 }
diff --git a/proxy/trojan/server.go b/proxy/trojan/server.go
index c73d4e2f..57e232b1 100644
--- a/proxy/trojan/server.go
+++ b/proxy/trojan/server.go
@@ -250,7 +250,9 @@ func (s *Server) Process(ctx context.Context, network net.Network, conn internet
 
 func (s *Server) handleUDPPayload(ctx context.Context, clientReader *PacketReader, clientWriter *PacketWriter, dispatcher routing.Dispatcher) error {
 	udpServer := udp.NewDispatcher(dispatcher, func(ctx context.Context, packet *udp_proto.Packet) {
-		common.Must(clientWriter.WriteMultiBufferWithMetadata(buf.MultiBuffer{packet.Payload}, packet.Source))
+		udpPayload := packet.Payload
+		udpPayload.UDP = &packet.Source
+		common.Must(clientWriter.WriteMultiBuffer(buf.MultiBuffer{udpPayload}))
 	})
 
 	inbound := session.InboundFromContext(ctx)
@@ -263,7 +265,7 @@ func (s *Server) handleUDPPayload(ctx context.Context, clientReader *PacketReade
 		case <-ctx.Done():
 			return nil
 		default:
-			p, err := clientReader.ReadMultiBufferWithMetadata()
+			mb, err := clientReader.ReadMultiBuffer()
 			if err != nil {
 				if errors.Cause(err) != io.EOF {
 					return newError("unexpected EOF").Base(err)
@@ -271,21 +273,24 @@ func (s *Server) handleUDPPayload(ctx context.Context, clientReader *PacketReade
 				return nil
 			}
 
+			mb2, b := buf.SplitFirst(mb)
+			destination := *b.UDP
 			ctx = log.ContextWithAccessMessage(ctx, &log.AccessMessage{
 				From:   inbound.Source,
-				To:     p.Target,
+				To:     destination,
 				Status: log.AccessAccepted,
 				Reason: "",
 				Email:  user.Email,
 			})
-			newError("tunnelling request to ", p.Target).WriteToLog(session.ExportIDToError(ctx))
+			newError("tunnelling request to ", destination).WriteToLog(session.ExportIDToError(ctx))
 
 			if !buf.Cone || dest == nil {
-				dest = &p.Target
+				dest = &destination
 			}
 
-			for _, b := range p.Buffer {
-				udpServer.Dispatch(ctx, *dest, b)
+			udpServer.Dispatch(ctx, *dest, b) // first packet
+			for _, payload := range mb2 {
+				udpServer.Dispatch(ctx, *dest, payload)
 			}
 		}
 	}