From 1816fa521c4754732127f7b9ddba0f2cdd9d9544 Mon Sep 17 00:00:00 2001
From: Daniel Adam <daniel.adam1922@protonmail.com>
Date: Wed, 8 Nov 2023 12:37:54 +0100
Subject: [PATCH] Update UnmarshalJSON in codes package

Try fix Issue 61477 reported by libFuzzer in OSS-Fuzz
---
 .github/workflows/cifuzz.yml |  4 +++-
 message/codes/codes.go       | 23 +++++++++++++++++++++-
 message/codes/codes_test.go  | 37 +++++++++++++++++++++++++++++++++---
 3 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/cifuzz.yml b/.github/workflows/cifuzz.yml
index 74ae5992..255ff0a3 100644
--- a/.github/workflows/cifuzz.yml
+++ b/.github/workflows/cifuzz.yml
@@ -1,5 +1,7 @@
 name: CIFuzz
-on: [pull_request]
+on:
+  pull_request:
+  workflow_dispatch:
 permissions: {}
 jobs:
  Fuzzing:
diff --git a/message/codes/codes.go b/message/codes/codes.go
index 62e3d8c3..5cda66de 100644
--- a/message/codes/codes.go
+++ b/message/codes/codes.go
@@ -56,6 +56,8 @@ const (
 
 const _maxCode = 255
 
+var _maxCodeLen int
+
 var strToCode = map[string]Code{
 	`"GET"`:                                GET,
 	`"POST"`:                               POST,
@@ -90,6 +92,21 @@ var strToCode = map[string]Code{
 	`"Abort"`:                              Abort,
 }
 
+func getMaxCodeLen() int {
+	// max uint32 as string binary representation: "0b" + 32 digits
+	max := 34
+	for k := range strToCode {
+		if len(k) > max {
+			max = len(k)
+		}
+	}
+	return max
+}
+
+func init() {
+	_maxCodeLen = getMaxCodeLen()
+}
+
 // UnmarshalJSON unmarshals b into the Code.
 func (c *Code) UnmarshalJSON(b []byte) error {
 	// From json.Unmarshaler: By convention, to approximate the behavior of
@@ -102,6 +119,10 @@ func (c *Code) UnmarshalJSON(b []byte) error {
 		return fmt.Errorf("nil receiver passed to UnmarshalJSON")
 	}
 
+	if len(b) > _maxCodeLen {
+		return fmt.Errorf("invalid code: input too large(length=%d)", len(b))
+	}
+
 	if ci, err := strconv.ParseUint(string(b), 10, 32); err == nil {
 		if ci >= _maxCode {
 			return fmt.Errorf("invalid code: %q", ci)
@@ -115,5 +136,5 @@ func (c *Code) UnmarshalJSON(b []byte) error {
 		*c = jc
 		return nil
 	}
-	return fmt.Errorf("invalid code: %q", string(b))
+	return fmt.Errorf("invalid code: %v", b)
 }
diff --git a/message/codes/codes_test.go b/message/codes/codes_test.go
index 4b62fcb0..8cd2e2c0 100644
--- a/message/codes/codes_test.go
+++ b/message/codes/codes_test.go
@@ -2,6 +2,7 @@ package codes
 
 import (
 	"encoding/json"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -14,6 +15,24 @@ func TestJSONUnmarshal(t *testing.T) {
 	err := json.Unmarshal([]byte(in), &got)
 	require.NoError(t, err)
 	require.Equal(t, want, got)
+
+	inNumeric := "["
+	for i, c := range want {
+		if i > 0 {
+			inNumeric += ","
+		}
+		inNumeric += strconv.FormatUint(uint64(c), 10)
+	}
+	inNumeric += "]"
+	err = json.Unmarshal([]byte(inNumeric), &got)
+	require.NoError(t, err)
+	require.Equal(t, want, got)
+}
+
+func TestUnmarshalJSONNoop(t *testing.T) {
+	var got Code
+	err := got.UnmarshalJSON([]byte("null"))
+	require.NoError(t, err)
 }
 
 func TestUnmarshalJSONNilReceiver(t *testing.T) {
@@ -24,11 +43,18 @@ func TestUnmarshalJSONNilReceiver(t *testing.T) {
 }
 
 func TestUnmarshalJSONUnknownInput(t *testing.T) {
-	var got Code
-	for _, in := range [][]byte{[]byte(""), []byte("xxx"), []byte("Code(17)"), nil} {
+	inputs := [][]byte{nil, []byte(""), []byte("xxx"), []byte("Code(17)"), []byte("255")}
+	for _, in := range inputs {
+		var got Code
 		err := got.UnmarshalJSON(in)
 		require.Error(t, err)
 	}
+
+	var got Code
+	longStr := "This is a very long string that is longer than the max code length"
+	require.True(t, len(longStr) > getMaxCodeLen())
+	err := got.UnmarshalJSON([]byte(longStr))
+	require.Error(t, err)
 }
 
 func TestUnmarshalJSONMarshalUnmarshal(t *testing.T) {
@@ -59,11 +85,16 @@ func TestCodeToString(t *testing.T) {
 }
 
 func FuzzUnmarshalJSON(f *testing.F) {
+	f.Add([]byte("null"))
 	f.Add([]byte("xxx"))
 	f.Add([]byte("Code(17)"))
+	f.Add([]byte("0b101010"))
+	f.Add([]byte("0o52"))
+	f.Add([]byte("0x2a"))
+	f.Add([]byte("42"))
 
 	f.Fuzz(func(t *testing.T, input_data []byte) {
-		var got *Code
+		var got Code
 		_ = got.UnmarshalJSON(input_data)
 	})
 }