Prechádzať zdrojové kódy

fix tcp testing server

Darien Raymond 7 rokov pred
rodič
commit
c4491d06d3
1 zmenil súbory, kde vykonal 33 pridanie a 15 odobranie
  1. 33 15
      testing/servers/tcp/tcp.go

+ 33 - 15
testing/servers/tcp/tcp.go

@@ -4,7 +4,10 @@ import (
 	"fmt"
 	"io"
 
+	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/net"
+	"v2ray.com/core/common/task"
+	"v2ray.com/core/transport/pipe"
 )
 
 type Server struct {
@@ -52,25 +55,40 @@ func (server *Server) handleConnection(conn net.Conn) {
 	if len(server.SendFirst) > 0 {
 		conn.Write(server.SendFirst)
 	}
-	request := make([]byte, 4096)
-	for {
-		nBytes, err := conn.Read(request)
-		if err != nil {
-			if err != io.EOF {
-				fmt.Println("Failed to read request:", err)
+
+	pReader, pWriter := pipe.New(pipe.WithoutSizeLimit())
+	err := task.Run(task.Parallel(func() error {
+		for {
+			b := buf.New()
+			if err := b.AppendSupplier(buf.ReadFrom(conn)); err != nil {
+				if err == io.EOF {
+					return nil
+				}
+				return err
+			}
+			copy(b.Bytes(), server.MsgProcessor(b.Bytes()))
+			if err := pWriter.WriteMultiBuffer(buf.NewMultiBufferValue(b)); err != nil {
+				return err
 			}
-			break
-		}
-		response := server.MsgProcessor(request[:nBytes])
-		if _, err := conn.Write(response); err != nil {
-			fmt.Println("Failed to write response:", err)
-			break
 		}
-		if server.ShouldClose {
-			break
+	}, func() error {
+		w := buf.NewWriter(conn)
+		for {
+			mb, err := pReader.ReadMultiBuffer()
+			if err != nil {
+				return err
+			}
+			if err := w.WriteMultiBuffer(mb); err != nil {
+				return err
+			}
 		}
+	}))()
+
+	if err != nil {
+		fmt.Println("failed to transfer data: ", err.Error())
 	}
-	conn.Close()
+
+	conn.Close() // nolint: errcheck
 }
 
 func (server *Server) Close() error {