Pārlūkot izejas kodu

remove static bytes array

Darien Raymond 7 gadi atpakaļ
vecāks
revīzija
009d58dd6c
1 mainītis faili ar 18 papildinājumiem un 17 dzēšanām
  1. 18 17
      transport/internet/kcp/io.go

+ 18 - 17
transport/internet/kcp/io.go

@@ -5,6 +5,8 @@ import (
 	"crypto/rand"
 	"io"
 
+	"v2ray.com/core/common"
+	"v2ray.com/core/common/buf"
 	"v2ray.com/core/transport/internet"
 )
 
@@ -50,8 +52,6 @@ type KCPPacketWriter struct {
 	Header   internet.PacketHeader
 	Security cipher.AEAD
 	Writer   io.Writer
-
-	buffer [2048]byte
 }
 
 func (w *KCPPacketWriter) Overhead() int {
@@ -66,27 +66,28 @@ func (w *KCPPacketWriter) Overhead() int {
 }
 
 func (w *KCPPacketWriter) Write(b []byte) (int, error) {
-	x := w.buffer[:]
-	size := 0
+	bb := buf.NewSize(int32(len(b) + w.Overhead()))
+	defer bb.Release()
+
 	if w.Header != nil {
-		nBytes, _ := w.Header.Write(x)
-		size += nBytes
-		x = x[nBytes:]
+		common.Must(bb.AppendSupplier(func(x []byte) (int, error) {
+			return w.Header.Write(x)
+		}))
 	}
 	if w.Security != nil {
 		nonceSize := w.Security.NonceSize()
-		var nonce []byte
-		if nonceSize > 0 {
-			nonce = x[:nonceSize]
-			rand.Read(nonce)
-			x = x[nonceSize:]
-		}
-		x = w.Security.Seal(x[:0], nonce, b, nil)
-		size += nonceSize + len(x)
+		common.Must(bb.AppendSupplier(func(x []byte) (int, error) {
+			return rand.Read(x[:nonceSize])
+		}))
+		nonce := bb.BytesFrom(int32(-nonceSize))
+		common.Must(bb.AppendSupplier(func(x []byte) (int, error) {
+			eb := w.Security.Seal(x[:0], nonce, b, nil)
+			return len(eb), nil
+		}))
 	} else {
-		size += copy(x, b)
+		bb.Append(b)
 	}
 
-	_, err := w.Writer.Write(w.buffer[:size])
+	_, err := w.Writer.Write(bb.Bytes())
 	return len(b), err
 }