optimize: memory consumption

internal/algo/crypto.go

@@ -5,58 +5,51 @@ import (
 	"crypto/rand"
 	"encoding/binary"
 	"errors"
-
-	"github.com/fumiama/orbyte/pbuf"
 )
 
 var (
 	ErrCipherTextTooShort = errors.New("ciphertext too short")
 )
 
-func EncodeAEAD(aead cipher.AEAD, additional uint16, b []byte) pbuf.Bytes {
+func EncodeAEAD(aead cipher.AEAD, additional uint16, b []byte) []byte {
 	nsz := aead.NonceSize()
 	// Accocate capacity for all the stuffs.
-	buf := pbuf.NewBytes(2 + nsz + len(b) + aead.Overhead())
+	buf := make([]byte, 2+nsz+len(b)+aead.Overhead())
 	n := 0
-	buf.V(func(buf []byte) {
-		binary.LittleEndian.PutUint16(buf[:2], additional)
-		nonce := buf[2 : 2+nsz]
-		// Select a random nonce
-		_, err := rand.Read(nonce)
-		if err != nil {
-			panic(err)
-		}
-		// Encrypt the message and append the ciphertext to the nonce.
-		eb := aead.Seal(nonce[nsz:nsz], nonce, b, buf[:2])
-		n = len(eb)
-	})
-	return buf.Slice(2, 2+nsz+n)
+	binary.LittleEndian.PutUint16(buf[:2], additional)
+	nonce := buf[2 : 2+nsz]
+	// Select a random nonce
+	_, err := rand.Read(nonce)
+	if err != nil {
+		panic(err)
+	}
+	// Encrypt the message and append the ciphertext to the nonce.
+	eb := aead.Seal(nonce[nsz:nsz], nonce, b, buf[:2])
+	n = len(eb)
+	return buf[2 : 2+nsz+n]
 }
 
-func DecodeAEAD(aead cipher.AEAD, additional uint16, b []byte) (data pbuf.Bytes, err error) {
+func DecodeAEAD(aead cipher.AEAD, additional uint16, b []byte) (data []byte, err error) {
 	nsz := aead.NonceSize()
 	if len(b) < nsz {
-		return pbuf.Bytes{}, ErrCipherTextTooShort
+		return nil, ErrCipherTextTooShort
 	}
 	// Split nonce and ciphertext.
 	nonce, ciphertext := b[:nsz], b[nsz:]
 	if len(ciphertext) == 0 {
-		return pbuf.Bytes{}, nil
+		return nil, nil
 	}
 	// Decrypt the message and check it wasn't tampered with.
 	var buf [2]byte
 	binary.LittleEndian.PutUint16(buf[:], additional)
-	data = pbuf.NewBytes(len(ciphertext))
+	data = make([]byte, len(ciphertext))
 	n := 0
-	data.V(func(b []byte) {
-		var d []byte
-		d, err = aead.Open(b[:0], nonce, ciphertext, buf[:])
-		n = len(d)
-	})
+	d, err := aead.Open(data[:0], nonce, ciphertext, buf[:])
+	n = len(d)
 	if err != nil {
 		return
 	}
-	return data.SliceTo(n), nil
+	return data[:n], nil
 }
 
 func EncodeXORLen(datalen int) int {
@@ -65,29 +58,27 @@ func EncodeXORLen(datalen int) int {
 }
 
 // EncodeXOR 按 8 字节, 以初始 mask 循环异或编码 data
-func EncodeXOR(data []byte, mask uint64, seq uint32) pbuf.Bytes {
+func EncodeXOR(data []byte, mask uint64, seq uint32) []byte {
 	batchsz := len(data) / 8
 	remain := len(data) % 8
 	sum := mask
-	newdat := pbuf.NewBytes(EncodeXORLen(len(data)))
-	newdat.V(func(buf []byte) {
-		binary.LittleEndian.PutUint32(buf[:4], seq)
-		_, _ = rand.Read(buf[4:8])                 // seqrand
-		sum ^= binary.LittleEndian.Uint64(buf[:8]) // init from seqrand
-		binary.LittleEndian.PutUint64(buf[:8], sum)
-		for i := 0; i < batchsz; i++ { // range on batch data
-			a := i * 8
-			b := (i + 1) * 8
-			sum ^= binary.LittleEndian.Uint64(data[a:b])
-			binary.LittleEndian.PutUint64(buf[a+8:b+8], sum)
-		}
-		p := batchsz * 8
-		copy(buf[8+p:], data[p:])
-		buf[newdat.Len()-1] = byte(remain)
-		sum ^= binary.LittleEndian.Uint64(buf[8+p:])
-		binary.LittleEndian.PutUint64(buf[8+p:], sum)
-	})
-	return newdat
+	buf := make([]byte, EncodeXORLen(len(data)))
+	binary.LittleEndian.PutUint32(buf[:4], seq)
+	_, _ = rand.Read(buf[4:8])                 // seqrand
+	sum ^= binary.LittleEndian.Uint64(buf[:8]) // init from seqrand
+	binary.LittleEndian.PutUint64(buf[:8], sum)
+	for i := 0; i < batchsz; i++ { // range on batch data
+		a := i * 8
+		b := (i + 1) * 8
+		sum ^= binary.LittleEndian.Uint64(data[a:b])
+		binary.LittleEndian.PutUint64(buf[a+8:b+8], sum)
+	}
+	p := batchsz * 8
+	copy(buf[8+p:], data[p:])
+	buf[len(buf)-1] = byte(remain)
+	sum ^= binary.LittleEndian.Uint64(buf[8+p:])
+	binary.LittleEndian.PutUint64(buf[8+p:], sum)
+	return buf
 }
 
 // DecodeXOR 按 8 字节, 以初始 mask 循环异或解码 data,

internal/algo/crypto_test.go

@@ -25,7 +25,7 @@ func TestXOR(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		seq, dec := DecodeXOR(EncodeXOR(r1.Bytes(), mask, uint32(i)).Trans(), mask)
+		seq, dec := DecodeXOR(EncodeXOR(r1.Bytes(), mask, uint32(i)), mask)
 		if !bytes.Equal(dec, r2.Bytes()) {
 			t.Fatal("unexpected xor at", i, "except", hex.EncodeToString(r2.Bytes()), "got", hex.EncodeToString(dec))
 		}
@@ -51,11 +51,11 @@ func TestXChacha20(t *testing.T) {
 		t.Fatal(err)
 	}
 	for i := 0; i < 4096; i++ {
-		db, err := DecodeAEAD(aead, uint16(i), EncodeAEAD(aead, uint16(i), data[:i]).Trans())
+		db, err := DecodeAEAD(aead, uint16(i), EncodeAEAD(aead, uint16(i), data[:i]))
 		if err != nil {
 			t.Fatal(err)
 		}
-		if !bytes.Equal(db.Trans(), data[:i]) {
+		if !bytes.Equal(db, data[:i]) {
 			t.Fatal("unexpected preshared at idx(len)", i, "addt", uint16(i))
 		}
 	}

internal/algo/zstd.go

@@ -9,7 +9,7 @@ import (
 	"github.com/klauspost/compress/zstd"
 )
 
-func EncodeZstd(data []byte) pbuf.Bytes {
+func EncodeZstd(data []byte) []byte {
 	return bin.SelectWriter().P(func(w *pbuf.Buffer) {
 		enc, err := zstd.NewWriter(w, zstd.WithEncoderLevel(zstd.SpeedFastest))
 		if err != nil {
@@ -23,19 +23,19 @@ func EncodeZstd(data []byte) pbuf.Bytes {
 		if err != nil {
 			panic(err)
 		}
-	}).ToBytes()
+	}).ToBytes().Copy().Ignore().Trans()
 }
 
-func DecodeZstd(data []byte) (b pbuf.Bytes, err error) {
+func DecodeZstd(data []byte) (b []byte, err error) {
 	dec, err := zstd.NewReader(bytes.NewReader(data))
 	if err != nil {
-		return pbuf.Bytes{}, err
+		return
 	}
 
 	b = bin.SelectWriter().P(func(w *pbuf.Buffer) {
 		_, err = io.Copy(w, dec)
 		dec.Close()
-	}).ToBytes()
+	}).ToBytes().Copy().Ignore().Trans()
 
 	return
 }