From 53a0c7148dd51a98f6ff8b42061d5b0792b96ae4 Mon Sep 17 00:00:00 2001 From: Grivn <56754250+Grivn@users.noreply.github.com> Date: Fri, 16 Oct 2026 19:38:30 +0800 Subject: [PATCH] fix(read): reject truncated messages Transport EOF before the final frame is complete can silently return a partial message. Report unexpected EOF and finish the remaining framing after a DEFLATE BFINAL block so valid fragmented messages still work. --- read.go | 19 ++++++- read_test.go | 136 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 154 insertions(+), 1 deletion(-) create mode 100644 read_test.go diff --git a/read.go b/read.go index c19f41de..0954e945 100644 --- a/read.go +++ b/read.go @@ -444,7 +444,18 @@ func (mr *msgReader) Read(p []byte) (n int, err error) { p = p[:n] mr.dict.write(p) } - if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) && mr.fin && mr.flate { + if mr.flate && errors.Is(err, io.EOF) { + // A DEFLATE block with BFINAL set can end before the final frame. + // Discard its padding, but still require the complete message framing. + mr.flateBufio.Reset(mr.readFunc) + if !mr.fin || mr.payloadLength != 0 { + _, err = io.Copy(io.Discard, mr.readFunc) + if err == nil { + err = io.EOF + } + } + } + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) && mr.fin && mr.flate && mr.payloadLength == 0 { mr.putFlateReader() return n, io.EOF } @@ -466,6 +477,9 @@ func (mr *msgReader) read(p []byte) (int, error) { h, err := mr.c.readLoop(mr.ctx) if err != nil { + if errors.Is(err, io.EOF) { + return 0, io.ErrUnexpectedEOF + } return 0, err } if h.opcode != opContinuation { @@ -484,6 +498,9 @@ func (mr *msgReader) read(p []byte) (int, error) { n, err := mr.c.readFramePayload(mr.ctx, p) if err != nil { + if errors.Is(err, io.EOF) { + return n, io.ErrUnexpectedEOF + } return n, err } diff --git a/read_test.go b/read_test.go new file mode 100644 index 00000000..e1362fac --- /dev/null +++ b/read_test.go @@ -0,0 +1,136 @@ +//go:build !js + +package websocket + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "net" + "testing" + "time" +) + +func TestReadMessageCompletion(t *testing.T) { + t.Parallel() + + for _, encoding := range []struct { + name string + payload []byte + flate bool + }{ + {"plain", []byte("Hello"), false}, + // RFC 7692 sections 7.2.3.1 and 7.2.3.4 encode the same message. + {"deflate", []byte{0xf2, 0x48, 0xcd, 0xc9, 0xc9, 0x07, 0x00}, true}, + {"deflateBFINAL", []byte{0xf3, 0x48, 0xcd, 0xc9, 0xc9, 0x07, 0x00, 0x00}, true}, + } { + for _, client := range []bool{false, true} { + for _, reader := range []bool{false, true} { + name := fmt.Sprintf("%s/client=%t/reader=%t", encoding.name, client, reader) + t.Run(name, func(t *testing.T) { + payload := encoding.payload + first := byte(opText) + if encoding.flate { + first |= 0x40 + } + frame := func(op byte, length int, data []byte) []byte { + b := []byte{op, byte(length)} + if !client { + b[1] |= 0x80 + b = append(b, 0, 0, 0, 0) + } + return append(b, data...) + } + split := len(payload) - 1 + fragment := func(last []byte) []byte { + return append(frame(first, split, payload[:split]), last...) + } + + for _, tc := range []struct { + name string + wire []byte + complete bool + }{ + {"missingContinuation", frame(first, len(payload), payload), false}, + {"partialContinuationHeader", fragment([]byte{0x80}), false}, + {"missingPayload", frame(first|0x80, len(payload), nil), false}, + {"partialPayload", frame(first|0x80, len(payload), payload[:split]), false}, + {"missingContinuationPayload", fragment(frame(0x80, 1, nil)), false}, + {"partialContinuationPayload", fragment(frame(0x80, 2, payload[split:])), false}, + {"complete", frame(first|0x80, len(payload), payload), true}, + {"completeFragmented", fragment(frame(0x80, 1, payload[split:])), true}, + {"emptyFinalFragment", append(frame(first, len(payload), payload), frame(0x80, 0, nil)...), true}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + local, peer := net.Pipe() + c := newConn(connConfig{ + rwc: local, client: client, + br: bufio.NewReader(local), bw: bufio.NewWriter(local), + }) + if encoding.flate { + c.copts = CompressionNoContextTakeover.opts() + } + defer c.CloseNow() + defer peer.Close() + + wire := tc.wire + if tc.complete { + wire = append(wire, frame(0x80|byte(opBinary), 4, []byte("next"))...) + wire = append(wire, frame(first|0x80, len(payload), payload)...) + } + written := make(chan error, 1) + go func() { + _, err := peer.Write(wire) + peer.Close() + written <- err + }() + + read := func() (MessageType, []byte, error) { + if !reader { + return c.Read(ctx) + } + typ, r, err := c.Reader(ctx) + if err != nil { + return typ, nil, err + } + b, err := io.ReadAll(r) + return typ, b, err + } + typ, b, err := read() + if !tc.complete { + if !errors.Is(err, io.ErrUnexpectedEOF) { + t.Fatalf("read truncated message %q: got %v, want unexpected EOF", b, err) + } + } else { + if err != nil || typ != MessageText || string(b) != "Hello" { + t.Fatalf("read complete message: got %v %q %v", typ, b, err) + } + typ, b, err = read() + if err != nil || typ != MessageBinary || string(b) != "next" { + t.Fatalf("read following message: got %v %q %v", typ, b, err) + } + typ, b, err = read() + if err != nil || typ != MessageText || string(b) != "Hello" { + t.Fatalf("read repeated message: got %v %q %v", typ, b, err) + } + } + select { + case err := <-written: + if err != nil { + t.Fatal(err) + } + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + }) + } + }) + } + } + } +}