Fix bugs in the Grow feature

[ZekeLu]

Thanks for your excellent work!

https://github.com/chromedp/chromedp depends on this package, and it fails to send large requests due to the lack of fragmentation support in Chrome (reported here: ChromeDevTools/devtools-protocol#175). Now that Grow() is added to this package, I want to address this issue on the client side. It works like a charm! Thank you very much!

I found some issues though. So here is the pull request.

BTW, DisableFlush() is supposed to be used to enable auto growing of the buffer, right? Is it better to call it EnableAutoGrow()?

[gobwas]

Hi @ZekeLu, can you please explain which exactly issues you are fixing here?
For example, the reasons of "set the reserved bytes for the header to the max value" are not clear to me.

[ZekeLu]

Hi @gobwas, sorry that I don't make it clear. The commits are organized in this way: first, it adds a test to reproduce the issue, then a commit fixing that issue is followed. Maybe it's a little easier to review the commits one by one. I will submit this comment first so that you can get my reply asap. I will explain more in the next comment.

[ZekeLu]

Here is the first issue:

The writer is created with w := NewWriterSize(&dest, 0, 0, 16), in this case:

w
    .raw: len:18 cap:18
    .buf: len:16 cap:16
    .n: 0
    .Available(): 16

p: len:20 // message size

in Grow():
offset: 2
n: len(p) - w.Available() = 4 // <--- it should be len(p)
size: ceilPowerOfTwo(offset + w.n + n) = ceilPowerOfTwo(2 + 0 + 4) = 8

The size 8 is less than len(w.raw), so results in panic.

This issue is reproduced in d562c93 and fixed in b0474c8.

ws/wsutil/writer.go

Lines 286 to 296 in 2effe5e

	func (w *Writer) Write(p []byte) (n int, err error) {
	// Even empty p may make a sense.
	w.dirty = true
	var nn int
	for len(p) > w.Available() && w.err == nil {
	if w.noFlush {
	w.Grow(len(p) - w.Available())
	continue
	}
	if w.Buffered() == 0 {

ws/wsutil/writer.go

Lines 271 to 273 in 2effe5e

	func (w *Writer) Available() int {
	return len(w.buf) - w.n
	}

ws/wsutil/writer.go

Lines 339 to 346 in 2effe5e

	func (w *Writer) Grow(n int) {
	var (
	offset = len(w.raw) - len(w.buf)
	size = ceilPowerOfTwo(offset + w.n + n)
	)
	if size <= len(w.raw) {
	panic("wsutil: buffer grow leads to its reduce")
	}

[ZekeLu]

Here is the second issue:

Again, the writer is created with w := NewWriterSize(&dest, 0, 0, 16), in this case:

w
    .raw: len:18 cap:18
    .buf: len:16 cap:16

When the buffer grows, the offset is kept (see line 350 below):

ws/wsutil/writer.go

Lines 339 to 351 in 2effe5e

	func (w *Writer) Grow(n int) {
	var (
	offset = len(w.raw) - len(w.buf)
	size = ceilPowerOfTwo(offset + w.n + n)
	)
	if size <= len(w.raw) {
	panic("wsutil: buffer grow leads to its reduce")
	}
	p := make([]byte, size)
	copy(p, w.raw[:offset+w.n])
	w.raw = p
	w.buf = w.raw[offset:]
	}

This time, we set the message size to wsutil.len16 + 10, which will cause the header size to increase from 2 to 10:

ws/write.go

Lines 31 to 46 in 2effe5e

	func HeaderSize(h Header) (n int) {
	switch {
	case h.Length < 126:
	n = 2
	case h.Length <= len16:
	n = 4
	case h.Length <= len64:
	n = 10
	default:
	return -1
	}
	if h.Masked {
	n += len(h.Mask)
	}
	return n
	}

And later, when flushFragment is called, the computed skip value is incorrect (see line 495 below):

offset: 2
skip: offset - ws.HeaderSize(header) = 2 - 10 = -8

panic with runtime error: slice bounds out of range [-8:].

ws/wsutil/writer.go

Lines 472 to 500 in 2effe5e

	func (w *Writer) flushFragment(fin bool) (err error) {
	var (
	payload = w.buf[:w.n]
	header = ws.Header{
	OpCode: w.opCode(),
	Fin: fin,
	Length: int64(len(payload)),
	}
	)
	for _, ext := range w.extensions {
	header, err = ext.SetBits(header)
	if err != nil {
	return err
	}
	}
	if w.state.ClientSide() {
	header.Masked = true
	header.Mask = ws.NewMask()
	ws.Cipher(payload, header.Mask, 0)
	}
	// Write header to the header segment of the raw buffer.
	var (
	offset = len(w.raw) - len(w.buf)
	skip = offset - ws.HeaderSize(header)
	)
	buf := bytesWriter{
	buf: w.raw[skip:offset],
	}
	if err := ws.WriteHeader(&buf, header); err != nil {

This issue is reproduced in 5112688 and fixed in 458fafb.

[ZekeLu]

I'm sorry that the explanation looks like a mess. But it's all about calculation, and I don't know how to explain the issues more clearly.

And the final commit 59a2d64 is to fix bugs introduced by its previous commit. I intentionally not to meld it into the previous commit because the bugs are just covered by the autobahn test suite. Maybe we want to add unit tests to cover them.

That's all. Thank you!

[gobwas]

Hey @ZekeLu,

Thanks for explanation! I think it is a good idea to add unit tests for each (if multiple) mentioned issues. It would be easier to reproduce the bugs and understand your fixes. Can you add it please?

Anyway I will delve into this on weekend, so thanks for reporting and fixing.

upd. Ah, I see one unit test already, so maybe only one unit test is missing (the one you mentioned that is covered by autobahn)?

[ZekeLu]

Ah, I see one unit test already, so maybe only one unit test is missing (the one you mentioned that is covered by autobahn)?

Yes, I'm going to add unit tests for that. Thanks for your time!

[ZekeLu]

@gobwas unit tests have been added now. I put the commit with the unit tests (
22e5241) before the commit that fixes the issues (0322f0d), so that you can check out the commit and see the failed tests.

[gobwas]

Hey @ZekeLu, that's what I was thinking of! Great, thanks!

[ZekeLu]

Kindly ping @gobwas. Any updates on this PR? BTW, do you have a plan to release v1.1.0?

[gobwas]

Hey @ZekeLu!

Sorry for being silent, had pretty busy days. Anyway, I just finished delving into your work and want to thank you again for finding out this bugs! I did cherry-pick most of your commits. However, I didn't pick the 458fafb not because I didn't like the idea of simplification, but because it seemed to me that it takes too many places to change (and moreover produced some issues which you fixed in 6th comit?). So I fixed it in other way, please check the master branch if it works for you.

Regarding release, yeah, I was just waiting if anyone will find some inconvinience and then your bug fixes arrived, so I wanted to include them as well into v1.1.0. So if master works for your use cases, lets do release-candidate today and true release later on this week? 🎉

[ZekeLu]

Hey @gobwas, thank you for your review! I'm very appreciated for your excellent work!

I agree that 458fafb touches too many places, and I don't have the confidence that it doesn't bring any new bugs. I'm appreciated that you find a better way to fix it. I feel relaxed now.

The schedule is good to me. I'm going to test my use cases, and will feedback later.

Thank you very much!

[ZekeLu]

@gobwas I'm sorry but the changes failed my test. I guess here is the problem (see the comments below):

func (w *Writer) Grow(n int) {
	var (
		size     = len(w.raw)
		offset   = len(w.raw) - len(w.buf)
		buffered = w.Buffered()
	)
	for cap := size - offset - buffered; cap < n; {
		size = ceilPowerOfTwo(offset + buffered + n)
		// the size will be greater than the payload size most of the time,
		// and I think we should use the payload size to calculate the header size.
		offset = headerSize(w.state, size)
		cap = size - offset - buffered
	}
	if size <= len(w.raw) {
		panic("wsutil: buffer grow leads to its reduce")
	}
	p := make([]byte, size)
	// if the header size grows, the payload should be shifted,
	// but the copy below does not shift the payload.
	copy(p, w.raw[:offset+buffered])
	w.raw = p
	w.buf = w.raw[offset:]
}

Inspired by your changes, I think we can simply do it like this (diff based on cfd0d39):

diff --git a/wsutil/writer.go b/wsutil/writer.go
index 6941dc0..41a6016 100644
--- a/wsutil/writer.go
+++ b/wsutil/writer.go
@@ -336,18 +336,22 @@ func ceilPowerOfTwo(n int) int {
    return n
 }

 func (w *Writer) Grow(n int) {
    var (
-       offset = len(w.raw) - len(w.buf)
-       size   = ceilPowerOfTwo(offset + w.n + n)
+       oldOffset = len(w.raw) - len(w.buf)
+       newOffset = headerSize(w.state, w.n+n)
+       size      = ceilPowerOfTwo(newOffset + w.n + n)
    )
    if size <= len(w.raw) {
        panic("wsutil: buffer grow leads to its reduce")
    }
    p := make([]byte, size)
-   copy(p, w.raw[:offset+w.n])
+   copy(p[newOffset-oldOffset:], w.raw[:oldOffset+w.n])
    w.raw = p
-   w.buf = w.raw[offset:]
+   w.buf = w.raw[newOffset:]
 }

This change passed my tests. I will send another PR so that it will run the autobahn test suite to verify the change.

I know that you had been super busy recently, sorry to occupy your spare time! ☕

---

Update: I just realized that I can not send a PR based on cfd0d39 (Github required me to pick a branch or tag as the base). So the PR #138 is based on the current master branch. I keep the commit based on cfd0d39 here: ZekeLu@4e6089e. Maybe it will make it easy for you to review the change.

[gobwas]

Hey @ZekeLu, yeah, there was a bug with that old and new offsets, nice catch, thank you! Which tests did it fail? Can you commit some?

[ZekeLu]

It's a test in another repository which depends on your excellent package:

https://github.com/chromedp/chromedp/blob/5b511f121cb2fc1e868b8667049d3cff32a6e0b1/chromedp_test.go#L583

Now you ask, I will create a simple test and commit it later.

[gobwas]

Thanks! Will get back to it tomorrow Moscow Time.

[ZekeLu]

Then I will create the test tomorrow morning Beijing Time. It's very late here now (~ 3 AM).

[gobwas]

Sure! Night!

[ZekeLu]

Good morning @gobwas, I just forced pushed to #138. As usual, test first, fix follows.

I have changed the test code a little so that it covers more cases:

1. split into two writes, so that it will cover the copy() part of the Grow() method;
2. check the payload to make sure the result is correct.

---

The section below is off-topic, I'm not sure is it okay to discuss it here.

Since we are handling revered space for header, I thought that we should use reserve() in Grow(). Then I found out that the parameter n passed to reserve() means the length of the whole raw buffer, while that passed to headerSize() means the length of the payload. And it seems that reserve() is incorrectly used in TestControlWriter() (I think here it should use headerSize()). Should I send another PR to document it and update TestControlWriter()?

update: here is the doc I wanted to add: 8fa6aa1

[gobwas]

Hey @ZekeLu, hope you are well!
Thank you for your work! I've cherry-picked your test commit and did fix in a bit different way (but main fix is the same as in yours), hope you don't mind. Also, I've switched from headerSize() to reserve() as you pointed out, and I agree with you, nice catch again 👍🏻

Regarding the doc, seems that I missed you update; however I added it also, but fill free to open separate PR to clarify things better, if you feel we need it!

Waiting for your feedback regarding fix!
Have a great day! ☕

upd.

It's fun to see that comments to reserve() are almost identical :D

[ZekeLu]

Hey @gobwas, hope you are well too!

It feels like pair programming, and it's good! I took a carefully look into your changes, and found that I made some mistakes before:

1. I didn't realize that the reservation for the header will shrink the capacity and it should be handled dynamically. Nice job for the for loop!
2. I thought that buffered + wantSpace (n) was the final size of the payload. But that's not true, it's very possible that a client will write to the free space of the buffer later. So your approach to use reserve() in Grow() is the right way to go (It's not explicitly mentioned in my previous comment, but I thought that we should use headerSize() in Grow(). Now I know that it's wrong).

Thank you for your fix! I think everything is clear and we could close this PR now.

I really learn a lot from you! Thank you again! ☕ 🤝

[gobwas]

Hey @ZekeLu! Great to hear that! Yeah, it was nice to work together on this!
Closing this for now. See you! 👋🏻