Return short-write errors from frame writers

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
 }
 

write_test.go

@@ -2,11 +2,103 @@ package ws
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"testing"
 )
 
+type writerFunc func([]byte) (int, error)
+
+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
+		failWith error
+		want     error
+	}{
+		{
+			name: "short",
+			want: io.ErrShortWrite,
+		},
+		{
+			name:     "error",
+			failWith: failure,
+			want:     failure,
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			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)) },
+			} {
+				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)
+				}
+			}
+		})
+	}
+}
+
+func TestWriteFrameWriteErrors(t *testing.T) {
+	failure := fmt.Errorf("writer failed")
+	for _, test := range []struct {
+		name     string
+		retErrAt int
+		failWith error
+		want     error
+	}{
+		{
+			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 := writerFunc(func(p []byte) (int, error) {
+				calls++
+				if calls == test.retErrAt {
+					return len(p) - 1, test.failWith
+				}
+				return len(p), nil
+			})
+			if err := WriteFrame(w, NewTextFrame([]byte("hello"))); !errors.Is(err, test.want) {
+				t.Fatalf("got %v, want %v", err, test.want)
+			}
+			if calls != test.retErrAt {
+				t.Fatalf("got %d writes, want %d", calls, test.retErrAt)
+			}
+		})
+	}
+}
+
 func TestWriteHeader(t *testing.T) {
 	for i, test := range RWTestCases {
 		t.Run(fmt.Sprintf("#%d", i), func(t *testing.T) {