diff options
| -rw-r--r-- | pktline/pktline.go | 61 | ||||
| -rw-r--r-- | pktline/pktline_test.go | 21 |
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) } } |