pktline: Make Reader/Writer structs

2 files changed, 49 insertions, 33 deletions

diff --git a/pktline/pktline.go b/pktline/pktline.go
index b8e4ccee..de5360a6 100644
--- a/pktline/pktline.go
+++ b/pktline/pktline.go
@@ -27,14 +27,21 @@ const (
 	StatusResponseEnd
 )
 
-// ReadLine reads a single pkt-line from r into buf.
+// Reader reads pkt-line data from an io.Reader.
+type Reader struct {
+	r io.Reader
+}
+
+// NewReader returns a Reader that reads from r.
+func NewReader(r io.Reader) *Reader {
+	return &Reader{r: r}
+}
+
+// ReadLine reads a single pkt-line from the underlying reader into buf.
 // It returns the payload slice, number of payload bytes, and a status.
-func ReadLine(r io.Reader, buf []byte) ([]byte, int, Status, error) {
-	if r == nil {
-		return nil, 0, StatusEOF, ErrInvalidHeader
-	}
+func (pr *Reader) ReadLine(buf []byte) ([]byte, int, Status, error) {
 	var header [4]byte
-	if _, err := io.ReadFull(r, header[:]); err != nil {
+	if _, err := io.ReadFull(pr.r, header[:]); err != nil {
 		if errors.Is(err, io.EOF) {
 			return nil, 0, StatusEOF, io.EOF
 		}
@@ -66,7 +73,7 @@ func ReadLine(r io.Reader, buf []byte) ([]byte, int, Status, error) {
 	if n > len(buf) {
 		return nil, 0, StatusEOF, ErrBufferTooSmall
 	}
-	if _, err := io.ReadFull(r, buf[:n]); err != nil {
+	if _, err := io.ReadFull(pr.r, buf[:n]); err != nil {
 		if errors.Is(err, io.ErrUnexpectedEOF) {
 			return nil, 0, StatusEOF, io.ErrUnexpectedEOF
 		}
@@ -75,46 +82,48 @@ func ReadLine(r io.Reader, buf []byte) ([]byte, int, Status, error) {
 	return buf[:n], n, StatusData, nil
 }
 
+// Writer writes pkt-line data to an io.Writer.
+type Writer struct {
+	w io.Writer
+}
+
+// NewWriter returns a Writer that writes to w.
+func NewWriter(w io.Writer) *Writer {
+	return &Writer{w: w}
+}
+
 // WriteLine writes a single pkt-line with data as its payload.
-func WriteLine(w io.Writer, data []byte) error {
-	if w == nil {
-		return ErrInvalidHeader
-	}
+func (pw *Writer) WriteLine(data []byte) error {
 	if len(data) > maxPacketDataLen {
 		return ErrPacketTooLarge
 	}
 	var header [4]byte
 	setHeader(header[:], len(data)+4)
-	if _, err := w.Write(header[:]); err != nil {
+	if _, err := pw.w.Write(header[:]); err != nil {
 		return err
 	}
 	if len(data) == 0 {
 		return nil
 	}
-	_, err := w.Write(data)
+	_, err := pw.w.Write(data)
 	return err
 }
 
 // Flush writes a flush-pkt ("0000").
-func Flush(w io.Writer) error {
-	return writeLiteral(w, "0000")
+func (pw *Writer) Flush() error {
+	_, err := io.WriteString(pw.w, "0000")
+	return err
 }
 
 // Delim writes a delim-pkt ("0001").
-func Delim(w io.Writer) error {
-	return writeLiteral(w, "0001")
+func (pw *Writer) Delim() error {
+	_, err := io.WriteString(pw.w, "0001")
+	return err
 }
 
 // ResponseEnd writes a response-end pkt ("0002").
-func ResponseEnd(w io.Writer) error {
-	return writeLiteral(w, "0002")
-}
-
-func writeLiteral(w io.Writer, s string) error {
-	if w == nil {
-		return ErrInvalidHeader
-	}
-	_, err := io.WriteString(w, s)
+func (pw *Writer) ResponseEnd() error {
+	_, err := io.WriteString(pw.w, "0002")
 	return err
 }
 
diff --git a/pktline/pktline_test.go b/pktline/pktline_test.go
index 4dae708b..185b0d98 100644
--- a/pktline/pktline_test.go
+++ b/pktline/pktline_test.go
@@ -10,12 +10,14 @@ import (
 func TestWriteReadLineRoundtrip(t *testing.T) {
 	var buf bytes.Buffer
 	payload := []byte("hello\n")
-	if err := WriteLine(&buf, payload); err != nil {
+	pw := NewWriter(&buf)
+	if err := pw.WriteLine(payload); err != nil {
 		t.Fatalf("WriteLine: %v", err)
 	}
 
 	dst := make([]byte, 64)
-	line, n, status, err := ReadLine(&buf, dst)
+	pr := NewReader(&buf)
+	line, n, status, err := pr.ReadLine(dst)
 	if err != nil {
 		t.Fatalf("ReadLine: %v", err)
 	}
@@ -43,8 +45,9 @@ func TestReadLineSpecialPackets(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			r := bytes.NewBufferString(tt.input)
+			pr := NewReader(r)
 			dst := make([]byte, 16)
-			line, n, status, err := ReadLine(r, dst)
+			line, n, status, err := pr.ReadLine(dst)
 			if err != nil {
 				t.Fatalf("ReadLine: %v", err)
 			}
@@ -60,8 +63,9 @@ func TestReadLineSpecialPackets(t *testing.T) {
 
 func TestReadLineInvalidHeader(t *testing.T) {
 	r := bytes.NewBufferString("zzzz")
+	pr := NewReader(r)
 	dst := make([]byte, 16)
-	_, _, _, err := ReadLine(r, dst)
+	_, _, _, err := pr.ReadLine(dst)
 	if !errors.Is(err, ErrInvalidHeader) {
 		t.Fatalf("expected ErrInvalidHeader, got %v", err)
 	}
@@ -70,11 +74,13 @@ func TestReadLineInvalidHeader(t *testing.T) {
 func TestReadLineBufferTooSmall(t *testing.T) {
 	var buf bytes.Buffer
 	payload := []byte("abcd")
-	if err := WriteLine(&buf, payload); err != nil {
+	pw := NewWriter(&buf)
+	if err := pw.WriteLine(payload); err != nil {
 		t.Fatalf("WriteLine: %v", err)
 	}
 	dst := make([]byte, 2)
-	_, _, _, err := ReadLine(&buf, dst)
+	pr := NewReader(&buf)
+	_, _, _, err := pr.ReadLine(dst)
 	if !errors.Is(err, ErrBufferTooSmall) {
 		t.Fatalf("expected ErrBufferTooSmall, got %v", err)
 	}
@@ -82,7 +88,8 @@ func TestReadLineBufferTooSmall(t *testing.T) {
 
 func TestWriteLineTooLarge(t *testing.T) {
 	payload := make([]byte, maxPacketDataLen+1)
-	if err := WriteLine(io.Discard, payload); !errors.Is(err, ErrPacketTooLarge) {
+	pw := NewWriter(io.Discard)
+	if err := pw.WriteLine(payload); !errors.Is(err, ErrPacketTooLarge) {
 		t.Fatalf("expected ErrPacketTooLarge, got %v", err)
 	}
 }