Correctly manage partial writes

mtglib/internal/faketls/record/pools.go

@@ -1,12 +1,22 @@
 package record
 
-import "sync"
+import (
+	"bytes"
+	"sync"
+)
 
-var recordPool = sync.Pool{
-	New: func() interface{} {
-		return &Record{}
-	},
-}
+var (
+	recordPool = sync.Pool{
+		New: func() interface{} {
+			return &Record{}
+		},
+	}
+	bytesBufferPool = sync.Pool{
+		New: func() interface{} {
+			return &bytes.Buffer{}
+		},
+	}
+)
 
 func AcquireRecord() *Record {
 	return recordPool.Get().(*Record)
@@ -16,3 +26,12 @@ func ReleaseRecord(r *Record) {
 	r.Reset()
 	recordPool.Put(r)
 }
+
+func acquireBytesBuffer() *bytes.Buffer {
+	return bytesBufferPool.Get().(*bytes.Buffer)
+}
+
+func releaseBytesBuffer(buf *bytes.Buffer) {
+	buf.Reset()
+	bytesBufferPool.Put(buf)
+}

mtglib/internal/faketls/record/record.go

@@ -61,26 +61,22 @@ func (r *Record) Read(reader io.Reader) error {
 }
 
 func (r *Record) Dump(writer io.Writer) error {
-	buf := [2]byte{byte(r.Type), 0}
+	buf := acquireBytesBuffer()
+	defer releaseBytesBuffer(buf)
 
-	if _, err := writer.Write(buf[:1]); err != nil {
-		return fmt.Errorf("cannot dump type: %w", err)
-	}
+	bufSlice := [2]byte{byte(r.Type), 0}
+	buf.Write(bufSlice[:1])
 
-	binary.BigEndian.PutUint16(buf[:], uint16(r.Version))
+	binary.BigEndian.PutUint16(bufSlice[:], uint16(r.Version))
+	buf.Write(bufSlice[:])
 
-	if _, err := writer.Write(buf[:]); err != nil {
-		return fmt.Errorf("cannot dump version: %w", err)
-	}
+	binary.BigEndian.PutUint16(bufSlice[:], uint16(r.Payload.Len()))
+	buf.Write(bufSlice[:])
 
-	binary.BigEndian.PutUint16(buf[:], uint16(r.Payload.Len()))
+	buf.Write(r.Payload.Bytes())
 
-	if _, err := writer.Write(buf[:]); err != nil {
-		return fmt.Errorf("cannot dump payload length: %w", err)
-	}
-
-	if _, err := writer.Write(r.Payload.Bytes()); err != nil {
-		return fmt.Errorf("cannot dump payload: %w", err)
+	if _, err := buf.WriteTo(writer); err != nil {
+		return fmt.Errorf("cannot dump record: %w", err)
 	}
 
 	return nil