XHTTP client: Fix edge-case issue for packet-up mode (#5020)

https://github.com/XTLS/Xray-core/pull/4952#issuecomment-3184080580

common/buf/buffer.go

@@ -13,6 +13,8 @@ const (
 	Size = 8192
 )
 
+var ErrBufferFull = errors.New("buffer is full")
+
 var zero = [Size * 10]byte{0}
 
 var pool = bytespool.GetPool(Size)
@@ -266,13 +268,16 @@ func (b *Buffer) IsFull() bool {
 func (b *Buffer) Write(data []byte) (int, error) {
 	nBytes := copy(b.v[b.end:], data)
 	b.end += int32(nBytes)
+	if nBytes < len(data) {
+		return nBytes, ErrBufferFull
+	}
 	return nBytes, nil
 }
 
 // WriteByte writes a single byte into the buffer.
 func (b *Buffer) WriteByte(v byte) error {
 	if b.IsFull() {
-		return errors.New("buffer full")
+		return ErrBufferFull
 	}
 	b.v[b.end] = v
 	b.end++

transport/internet/splithttp/dialer.go

@@ -489,15 +489,16 @@ func (w uploadWriter) Write(b []byte) (int, error) {
 		}
 	*/
 
-	buffer := buf.New()
+	buffer := buf.MultiBufferContainer{}
-	n, err := buffer.Write(b)
+	common.Must2(buffer.Write(b))
-	if err != nil {
-		return 0, err
-	}
 
-	err = w.WriteMultiBuffer([]*buf.Buffer{buffer})
+	var writed int
+	for _, buff := range buffer.MultiBuffer {
+		err := w.WriteMultiBuffer(buf.MultiBuffer{buff})
 		if err != nil {
-		return 0, err
+			return writed, err
 		}
-	return n, nil
+		writed += int(buff.Len())
+	}
+	return writed, nil
 }

transport/internet/splithttp/splithttp_test.go

@@ -1,6 +1,7 @@
 package splithttp_test
 
 import (
+	"bytes"
 	"context"
 	"crypto/rand"
 	"fmt"
@@ -421,18 +422,12 @@ func Test_maxUpload(t *testing.T) {
 		},
 	}
 
-	var uploadSize int
+	uploadReceived := make([]byte, 10001)
 	listen, err := ListenXH(context.Background(), net.LocalHostIP, listenPort, streamSettings, func(conn stat.Connection) {
 		go func(c stat.Connection) {
 			defer c.Close()
-			var b [10240]byte
 			c.SetReadDeadline(time.Now().Add(2 * time.Second))
-			n, err := c.Read(b[:])
+			io.ReadFull(c, uploadReceived)
-			if err != nil {
-				return
-			}
-
-			uploadSize = n
 
 			common.Must2(c.Write([]byte("Response")))
 		}(conn)
@@ -441,10 +436,12 @@ func Test_maxUpload(t *testing.T) {
 	ctx := context.Background()
 
 	conn, err := Dial(ctx, net.TCPDestination(net.DomainAddress("localhost"), listenPort), streamSettings)
+	common.Must(err)
 
 	// send a slightly too large upload
-	var upload [10001]byte
+	upload := make([]byte, 10001)
-	_, err = conn.Write(upload[:])
+	rand.Read(upload)
+	_, err = conn.Write(upload)
 	common.Must(err)
 
 	var b [10240]byte
@@ -455,8 +452,8 @@ func Test_maxUpload(t *testing.T) {
 	}
 	common.Must(conn.Close())
 
-	if uploadSize > 10000 || uploadSize == 0 {
+	if !bytes.Equal(upload, uploadReceived) {
-		t.Error("incorrect upload size: ", uploadSize)
+		t.Error("incorrect upload", upload, uploadReceived)
 	}
 
 	common.Must(listen.Close())