From 96f636c2daaf9c7822748d492c6fc508dff2912b Mon Sep 17 00:00:00 2001 From: Evan Simkowitz Date: Tue, 12 Dec 2023 16:51:19 -0800 Subject: [PATCH 1/4] Fix potential overflow in promptenc arithmetic --- wavesrv/pkg/promptenc/promptenc.go | 9 +++-- wavesrv/pkg/utilfn/utilfn.go | 31 +++++++++++++++++ wavesrv/pkg/utilfn/utilfn_test.go | 54 ++++++++++++++++++++++++++++++ 3 files changed, 92 insertions(+), 2 deletions(-) diff --git a/wavesrv/pkg/promptenc/promptenc.go b/wavesrv/pkg/promptenc/promptenc.go index 23bfe1bd..64ff2e40 100644 --- a/wavesrv/pkg/promptenc/promptenc.go +++ b/wavesrv/pkg/promptenc/promptenc.go @@ -12,6 +12,7 @@ import ( "io" "reflect" + "github.com/wavetermdev/waveterm/wavesrv/pkg/utilfn" ccp "golang.org/x/crypto/chacha20poly1305" ) @@ -65,9 +66,13 @@ func MakeEncryptorB64(key64 string) (*Encryptor, error) { } func (enc *Encryptor) EncryptData(plainText []byte, odata string) ([]byte, error) { - outputBuf := make([]byte, enc.AEAD.NonceSize()+enc.AEAD.Overhead()+len(plainText)) + bufSize, err := utilfn.AddIntSlice(enc.AEAD.NonceSize(), enc.AEAD.Overhead(), len(plainText)) + if err != nil { + return nil, err + } + outputBuf := make([]byte, bufSize) nonce := outputBuf[0:enc.AEAD.NonceSize()] - _, err := io.ReadFull(rand.Reader, nonce) + _, err = io.ReadFull(rand.Reader, nonce) if err != nil { return nil, err } diff --git a/wavesrv/pkg/utilfn/utilfn.go b/wavesrv/pkg/utilfn/utilfn.go index dbab5c29..5c870ecf 100644 --- a/wavesrv/pkg/utilfn/utilfn.go +++ b/wavesrv/pkg/utilfn/utilfn.go @@ -6,6 +6,8 @@ package utilfn import ( "crypto/sha1" "encoding/base64" + "errors" + "math" "regexp" "strings" "unicode/utf8" @@ -209,3 +211,32 @@ func ChunkSlice[T any](s []T, chunkSize int) [][]T { } return rtn } + +var ErrOverflow = errors.New("integer overflow") + +// Add two int values, returning an error if the result overflows. +func AddInt(left, right int) (int, error) { + if right > 0 { + if left > math.MaxInt-right { + return 0, ErrOverflow + } + } else { + if left < math.MaxInt-right { + return 0, ErrOverflow + } + } + return left + right, nil +} + +// Add a slice of ints, returning an error if the result overflows. +func AddIntSlice(vals ...int) (int, error) { + var rtn int + for _, v := range vals { + var err error + rtn, err = AddInt(rtn, v) + if err != nil { + return 0, err + } + } + return rtn, nil +} diff --git a/wavesrv/pkg/utilfn/utilfn_test.go b/wavesrv/pkg/utilfn/utilfn_test.go index c6956ab9..39f87150 100644 --- a/wavesrv/pkg/utilfn/utilfn_test.go +++ b/wavesrv/pkg/utilfn/utilfn_test.go @@ -5,6 +5,7 @@ package utilfn import ( "fmt" + "math" "testing" ) @@ -49,3 +50,56 @@ func TestDiff(t *testing.T) { testDiff(t, Str1, Str3) testDiff(t, Str3, Str1) } + +const unexpectedError = "unexpected error" +const expectedError = "expected error" +const wrongRetVal = "wrong return value" + +func testAddInt(t *testing.T, a int, b int, shouldError bool, expected int) { + retVal, err := AddInt(a, b) + if err != nil { + if !shouldError { + t.Errorf(unexpectedError) + } + return + } + if shouldError { + t.Errorf(expectedError) + return + } + if retVal != expected { + t.Errorf(wrongRetVal) + } +} + +func TestAddInt(t *testing.T) { + testAddInt(t, 1, 2, false, 3) + testAddInt(t, 1, math.MaxInt, true, 0) +} + +func testAddIntSlice(t *testing.T, shouldError bool, expected int, vals ...int) { + retVal, err := AddIntSlice(vals...) + if err != nil { + if !shouldError { + t.Errorf(unexpectedError) + } + return + } + if shouldError { + t.Errorf(expectedError) + return + } + if retVal != expected { + t.Errorf(wrongRetVal) + } +} + +func TestAddIntSlice(t *testing.T) { + testAddIntSlice(t, false, 0) + testAddIntSlice(t, false, 1, 1) + testAddIntSlice(t, false, 3, 1, 2) + testAddIntSlice(t, false, 6, 1, 2, 3) + testAddIntSlice(t, true, 0, 1, math.MaxInt) + testAddIntSlice(t, true, 0, 1, 2, math.MaxInt) + testAddIntSlice(t, true, 0, math.MaxInt, 2, 1) +} From 5d39815dea54b6c27831a15957a57726c1340a40 Mon Sep 17 00:00:00 2001 From: Evan Simkowitz Date: Tue, 12 Dec 2023 17:39:19 -0800 Subject: [PATCH 2/4] refactor a bit --- wavesrv/pkg/utilfn/utilfn_test.go | 33 ++++++++++--------------------- 1 file changed, 10 insertions(+), 23 deletions(-) diff --git a/wavesrv/pkg/utilfn/utilfn_test.go b/wavesrv/pkg/utilfn/utilfn_test.go index 39f87150..6523b7db 100644 --- a/wavesrv/pkg/utilfn/utilfn_test.go +++ b/wavesrv/pkg/utilfn/utilfn_test.go @@ -51,47 +51,34 @@ func TestDiff(t *testing.T) { testDiff(t, Str3, Str1) } -const unexpectedError = "unexpected error" -const expectedError = "expected error" -const wrongRetVal = "wrong return value" - -func testAddInt(t *testing.T, a int, b int, shouldError bool, expected int) { - retVal, err := AddInt(a, b) +func testArithmetic(t *testing.T, fn func() (int, error), shouldError bool, expected int) { // nolint: unparam + retVal, err := fn() if err != nil { if !shouldError { - t.Errorf(unexpectedError) + t.Errorf("unexpected error") } return } if shouldError { - t.Errorf(expectedError) + t.Errorf("expected error") return } if retVal != expected { - t.Errorf(wrongRetVal) + t.Errorf("wrong return value") } } +func testAddInt(t *testing.T, a int, b int, shouldError bool, expected int) { + testArithmetic(t, func() (int, error) { return AddInt(a, b) }, shouldError, expected) +} + func TestAddInt(t *testing.T) { testAddInt(t, 1, 2, false, 3) testAddInt(t, 1, math.MaxInt, true, 0) } func testAddIntSlice(t *testing.T, shouldError bool, expected int, vals ...int) { - retVal, err := AddIntSlice(vals...) - if err != nil { - if !shouldError { - t.Errorf(unexpectedError) - } - return - } - if shouldError { - t.Errorf(expectedError) - return - } - if retVal != expected { - t.Errorf(wrongRetVal) - } + testArithmetic(t, func() (int, error) { return AddIntSlice(vals...) }, shouldError, expected) } func TestAddIntSlice(t *testing.T) { From 414937911e42b9ffb260ccb5658d6b646996e925 Mon Sep 17 00:00:00 2001 From: Evan Simkowitz Date: Tue, 12 Dec 2023 17:40:05 -0800 Subject: [PATCH 3/4] remove nolint --- wavesrv/pkg/utilfn/utilfn_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wavesrv/pkg/utilfn/utilfn_test.go b/wavesrv/pkg/utilfn/utilfn_test.go index 6523b7db..1faf40bf 100644 --- a/wavesrv/pkg/utilfn/utilfn_test.go +++ b/wavesrv/pkg/utilfn/utilfn_test.go @@ -51,7 +51,7 @@ func TestDiff(t *testing.T) { testDiff(t, Str3, Str1) } -func testArithmetic(t *testing.T, fn func() (int, error), shouldError bool, expected int) { // nolint: unparam +func testArithmetic(t *testing.T, fn func() (int, error), shouldError bool, expected int) { retVal, err := fn() if err != nil { if !shouldError { From 1feffc8a114ea7355825d2d8c2560830639e7483 Mon Sep 17 00:00:00 2001 From: Evan Simkowitz Date: Tue, 12 Dec 2023 20:36:16 -0800 Subject: [PATCH 4/4] address comment, clean up tests --- wavesrv/pkg/utilfn/utilfn.go | 2 +- wavesrv/pkg/utilfn/utilfn_test.go | 17 ++++++++++++++--- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/wavesrv/pkg/utilfn/utilfn.go b/wavesrv/pkg/utilfn/utilfn.go index 5c870ecf..8df22fe7 100644 --- a/wavesrv/pkg/utilfn/utilfn.go +++ b/wavesrv/pkg/utilfn/utilfn.go @@ -221,7 +221,7 @@ func AddInt(left, right int) (int, error) { return 0, ErrOverflow } } else { - if left < math.MaxInt-right { + if left < math.MinInt-right { return 0, ErrOverflow } } diff --git a/wavesrv/pkg/utilfn/utilfn_test.go b/wavesrv/pkg/utilfn/utilfn_test.go index 1faf40bf..a65b4606 100644 --- a/wavesrv/pkg/utilfn/utilfn_test.go +++ b/wavesrv/pkg/utilfn/utilfn_test.go @@ -68,13 +68,18 @@ func testArithmetic(t *testing.T, fn func() (int, error), shouldError bool, expe } } -func testAddInt(t *testing.T, a int, b int, shouldError bool, expected int) { +func testAddInt(t *testing.T, shouldError bool, expected int, a int, b int) { testArithmetic(t, func() (int, error) { return AddInt(a, b) }, shouldError, expected) } func TestAddInt(t *testing.T) { - testAddInt(t, 1, 2, false, 3) - testAddInt(t, 1, math.MaxInt, true, 0) + testAddInt(t, false, 3, 1, 2) + testAddInt(t, true, 0, 1, math.MaxInt) + testAddInt(t, true, 0, math.MinInt, -1) + testAddInt(t, false, math.MaxInt-1, math.MaxInt, -1) + testAddInt(t, false, math.MinInt+1, math.MinInt, 1) + testAddInt(t, false, math.MaxInt, math.MaxInt, 0) + testAddInt(t, true, 0, math.MinInt, -1) } func testAddIntSlice(t *testing.T, shouldError bool, expected int, vals ...int) { @@ -89,4 +94,10 @@ func TestAddIntSlice(t *testing.T) { testAddIntSlice(t, true, 0, 1, math.MaxInt) testAddIntSlice(t, true, 0, 1, 2, math.MaxInt) testAddIntSlice(t, true, 0, math.MaxInt, 2, 1) + testAddIntSlice(t, false, math.MaxInt, 0, 0, math.MaxInt) + testAddIntSlice(t, true, 0, math.MinInt, -1) + testAddIntSlice(t, false, math.MaxInt, math.MaxInt-3, 1, 2) + testAddIntSlice(t, true, 0, math.MaxInt-2, 1, 2) + testAddIntSlice(t, false, math.MinInt, math.MinInt+3, -1, -2) + testAddIntSlice(t, true, 0, math.MinInt+2, -1, -2) }