From e166dd6557bad63dbf43ccc001a190dcbd1c86d2 Mon Sep 17 00:00:00 2001 From: omrzk <193513049+omrzk@users.noreply.github.com> Date: Sun, 25 Oct 2026 09:13:10 -0500 Subject: [PATCH 1/2] Return short-write errors when writing frame headers and payloads --- write.go | 10 +++++++-- write_test.go | 61 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 2 deletions(-) diff --git a/write.go b/write.go index 9e2fc07..60f2c60 100644 --- a/write.go +++ b/write.go @@ -93,7 +93,10 @@ func WriteHeaderBuffer(w io.Writer, h Header, bts []byte) error { n += copy(bts[n:], h.Mask[:]) } - _, err := w.Write(bts[:n]) + nn, err := w.Write(bts[:n]) + if nn < n && err == nil { + err = io.ErrShortWrite + } return err } @@ -104,7 +107,10 @@ func WriteFrame(w io.Writer, f Frame) error { if err != nil { return err } - _, err = w.Write(f.Payload) + n, err := w.Write(f.Payload) + if n < len(f.Payload) && err == nil { + err = io.ErrShortWrite + } return err } diff --git a/write_test.go b/write_test.go index e24f3c0..f2f780e 100644 --- a/write_test.go +++ b/write_test.go @@ -3,10 +3,71 @@ package ws import ( "bytes" "fmt" + "io" "io/ioutil" "testing" ) +type writeFunc func([]byte) (int, error) + +func (fn writeFunc) Write(p []byte) (int, error) { return fn(p) } + +func TestWriteHeaderWriteErrors(t *testing.T) { + failure := fmt.Errorf("writer failed") + for _, test := range []struct { + name string + err error + want error + }{ + {"short", nil, io.ErrShortWrite}, + {"error", failure, failure}, + } { + t.Run(test.name, func(t *testing.T) { + w := writeFunc(func(p []byte) (int, error) { return len(p) - 1, test.err }) + for _, write := range []func(io.Writer, Header) error{ + WriteHeader, + func(w io.Writer, h Header) error { return WriteHeaderBuffer(w, h, make([]byte, MaxHeaderSize)) }, + } { + if err := write(w, Header{Fin: true, OpCode: OpText, Length: 5}); err != test.want { + t.Fatalf("got %v, want %v", err, test.want) + } + } + }) + } +} + +func TestWriteFrameWriteErrors(t *testing.T) { + failure := fmt.Errorf("writer failed") + for _, test := range []struct { + name string + call int + err error + want error + }{ + {"header_short", 1, nil, io.ErrShortWrite}, + {"header_error", 1, failure, failure}, + {"payload_short", 2, nil, io.ErrShortWrite}, + {"payload_error", 2, failure, failure}, + } { + t.Run(test.name, func(t *testing.T) { + calls := 0 + w := writeFunc(func(p []byte) (int, error) { + calls++ + if calls == test.call { + return len(p) - 1, test.err + } + return len(p), nil + }) + if err := WriteFrame(w, NewTextFrame([]byte("hello"))); err != test.want { + t.Fatalf("got %v, want %v", err, test.want) + } + if calls != test.call { + t.Fatalf("got %d writes, want %d", calls, test.call) + } + }) + } +} + func TestWriteHeader(t *testing.T) { for i, test := range RWTestCases { t.Run(fmt.Sprintf("#%d", i), func(t *testing.T) { From c4c3f256a515dbd87243d697fae568d42d44fa8f Mon Sep 17 00:00:00 2001 From: omrzk <193513049+omrzk@users.noreply.github.com> Date: Sun, 25 Oct 2026 12:32:58 -0500 Subject: [PATCH 2/2] Refine short-write regression tests after review --- write_test.go | 77 ++++++++++++++++++++++++++++++++++++--------------- 1 file changed, 54 insertions(+), 23 deletions(-) diff --git a/write_test.go b/write_test.go index f2f780e..bde89ec 100644 --- a/write_test.go +++ b/write_test.go @@ -2,33 +2,46 @@ package ws import ( "bytes" + "errors" "fmt" "io" "io/ioutil" "testing" ) -type writeFunc func([]byte) (int, error) +type writerFunc func([]byte) (int, error) -func (fn writeFunc) Write(p []byte) (int, error) { return fn(p) } +func (fn writerFunc) Write(p []byte) (int, error) { return fn(p) } func TestWriteHeaderWriteErrors(t *testing.T) { failure := fmt.Errorf("writer failed") for _, test := range []struct { - name string - err error - want error + name string + failWith error + want error }{ - {"short", nil, io.ErrShortWrite}, - {"error", failure, failure}, + { + name: "short", + want: io.ErrShortWrite, + }, + { + name: "error", + failWith: failure, + want: failure, + }, } { t.Run(test.name, func(t *testing.T) { - w := writeFunc(func(p []byte) (int, error) { return len(p) - 1, test.err }) + w := writerFunc(func(p []byte) (int, error) { return len(p) - 1, test.failWith }) for _, write := range []func(io.Writer, Header) error{ WriteHeader, func(w io.Writer, h Header) error { return WriteHeaderBuffer(w, h, make([]byte, MaxHeaderSize)) }, } { - if err := write(w, Header{Fin: true, OpCode: OpText, Length: 5}); err != test.want { + h := Header{ + Fin: true, + OpCode: OpText, + Length: 5, + } + if err := write(w, h); !errors.Is(err, test.want) { t.Fatalf("got %v, want %v", err, test.want) } } @@ -39,30 +52,48 @@ func TestWriteHeaderWriteErrors(t *testing.T) { func TestWriteFrameWriteErrors(t *testing.T) { failure := fmt.Errorf("writer failed") for _, test := range []struct { - name string - call int - err error - want error + name string + retErrAt int + failWith error + want error }{ - {"header_short", 1, nil, io.ErrShortWrite}, - {"header_error", 1, failure, failure}, - {"payload_short", 2, nil, io.ErrShortWrite}, - {"payload_error", 2, failure, failure}, + { + name: "header_short", + retErrAt: 1, + want: io.ErrShortWrite, + }, + { + name: "header_error", + retErrAt: 1, + failWith: failure, + want: failure, + }, + { + name: "payload_short", + retErrAt: 2, + want: io.ErrShortWrite, + }, + { + name: "payload_error", + retErrAt: 2, + failWith: failure, + want: failure, + }, } { t.Run(test.name, func(t *testing.T) { calls := 0 - w := writeFunc(func(p []byte) (int, error) { + w := writerFunc(func(p []byte) (int, error) { calls++ - if calls == test.call { - return len(p) - 1, test.err + if calls == test.retErrAt { + return len(p) - 1, test.failWith } return len(p), nil }) - if err := WriteFrame(w, NewTextFrame([]byte("hello"))); err != test.want { + if err := WriteFrame(w, NewTextFrame([]byte("hello"))); !errors.Is(err, test.want) { t.Fatalf("got %v, want %v", err, test.want) } - if calls != test.call { - t.Fatalf("got %d writes, want %d", calls, test.call) + if calls != test.retErrAt { + t.Fatalf("got %d writes, want %d", calls, test.retErrAt) } }) }