Skip to content

Commit

Permalink
Enhance Options Buffer for >16 Message Unmarshaling
Browse files Browse the repository at this point in the history
Previously, the connection was being closed due to unmarshaling failures when
more than 16 options were included in the message. This fix addresses the issue
by increasing the buffer size to accommodate a higher number of options.
  • Loading branch information
jkralik committed Jan 12, 2024
1 parent 39a46aa commit c61c648
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 1 deletion.
16 changes: 15 additions & 1 deletion message/pool/message.go
Original file line number Diff line number Diff line change
Expand Up @@ -523,14 +523,28 @@ func (r *Message) MarshalWithEncoder(encoder Encoder) ([]byte, error) {
return r.bufferMarshal, nil
}

func (r *Message) decode(decoder Decoder) (int, error) {
var n int
var err error
for {
n, err = decoder.Decode(r.bufferUnmarshal, &r.msg)
if errors.Is(err, message.ErrOptionsTooSmall) {
// increase buffer size and try again
r.msg.Options = make(message.Options, 0, len(r.msg.Options)*2)
continue
}
return n, err
}
}

func (r *Message) UnmarshalWithDecoder(decoder Decoder, data []byte) (int, error) {
if len(r.bufferUnmarshal) < len(data) {
r.bufferUnmarshal = append(r.bufferUnmarshal, make([]byte, len(data)-len(r.bufferUnmarshal))...)
}
copy(r.bufferUnmarshal, data)
r.body = nil
r.bufferUnmarshal = r.bufferUnmarshal[:len(data)]
n, err := decoder.Decode(r.bufferUnmarshal, &r.msg)
n, err := r.decode(decoder)
if err != nil {
return n, err
}
Expand Down
31 changes: 31 additions & 0 deletions message/pool/message_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,13 @@ import (
"bytes"
"context"
"errors"
"fmt"
"strings"
"testing"

"github.com/plgd-dev/go-coap/v3/message"
"github.com/plgd-dev/go-coap/v3/message/pool"
"github.com/plgd-dev/go-coap/v3/tcp/coder"
"github.com/plgd-dev/go-coap/v3/test/net"
"github.com/stretchr/testify/require"
)
Expand Down Expand Up @@ -360,3 +362,32 @@ func TestMessageClone(t *testing.T) {
require.Error(t, err)
require.Contains(t, err.Error(), "read error")
}

func TestUnmarshalMessageWithMultipleOptions(t *testing.T) {
tests := []struct {
numOptions int
}{
{numOptions: 0},
{numOptions: 8},
{numOptions: 16},
{numOptions: 32},
{numOptions: 64},
{numOptions: 1023},
}

for _, tt := range tests {
t.Run(fmt.Sprintf("num-options-%v", tt.numOptions), func(t *testing.T) {
req := pool.NewMessage(context.Background())
for i := 0; i < 64; i++ {
req.AddOptionUint32(message.URIQuery, uint32(i))
}
data, err := req.MarshalWithEncoder(coder.DefaultCoder)
require.NoError(t, err)
msg := pool.NewMessage(context.Background())
n, err := msg.UnmarshalWithDecoder(coder.DefaultCoder, data)
require.NoError(t, err)
require.Equal(t, n, len(data))
require.Equal(t, req.Options(), msg.Options())
})
}
}

0 comments on commit c61c648

Please sign in to comment.