diff options
author | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-16 19:23:18 +0000 |
---|---|---|
committer | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-16 19:23:18 +0000 |
commit | 43a123c1ae6613b3efeed291fa552ecd909d3acf (patch) | |
tree | fd92518b7024bc74031f78a1cf9e454b65e73665 /src/internal/saferio | |
parent | Initial commit. (diff) | |
download | golang-1.20-43a123c1ae6613b3efeed291fa552ecd909d3acf.tar.xz golang-1.20-43a123c1ae6613b3efeed291fa552ecd909d3acf.zip |
Adding upstream version 1.20.14.upstream/1.20.14upstream
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'src/internal/saferio')
-rw-r--r-- | src/internal/saferio/io.go | 135 | ||||
-rw-r--r-- | src/internal/saferio/io_test.go | 136 |
2 files changed, 271 insertions, 0 deletions
diff --git a/src/internal/saferio/io.go b/src/internal/saferio/io.go new file mode 100644 index 0000000..66cc044 --- /dev/null +++ b/src/internal/saferio/io.go @@ -0,0 +1,135 @@ +// Copyright 2022 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package saferio provides I/O functions that avoid allocating large +// amounts of memory unnecessarily. This is intended for packages that +// read data from an [io.Reader] where the size is part of the input +// data but the input may be corrupt, or may be provided by an +// untrustworthy attacker. +package saferio + +import ( + "io" + "reflect" +) + +// chunk is an arbitrary limit on how much memory we are willing +// to allocate without concern. +const chunk = 10 << 20 // 10M + +// ReadData reads n bytes from the input stream, but avoids allocating +// all n bytes if n is large. This avoids crashing the program by +// allocating all n bytes in cases where n is incorrect. +// +// The error is io.EOF only if no bytes were read. +// If an io.EOF happens after reading some but not all the bytes, +// ReadData returns io.ErrUnexpectedEOF. +func ReadData(r io.Reader, n uint64) ([]byte, error) { + if int64(n) < 0 || n != uint64(int(n)) { + // n is too large to fit in int, so we can't allocate + // a buffer large enough. Treat this as a read failure. + return nil, io.ErrUnexpectedEOF + } + + if n < chunk { + buf := make([]byte, n) + _, err := io.ReadFull(r, buf) + if err != nil { + return nil, err + } + return buf, nil + } + + var buf []byte + buf1 := make([]byte, chunk) + for n > 0 { + next := n + if next > chunk { + next = chunk + } + _, err := io.ReadFull(r, buf1[:next]) + if err != nil { + if len(buf) > 0 && err == io.EOF { + err = io.ErrUnexpectedEOF + } + return nil, err + } + buf = append(buf, buf1[:next]...) + n -= next + } + return buf, nil +} + +// ReadDataAt reads n bytes from the input stream at off, but avoids +// allocating all n bytes if n is large. This avoids crashing the program +// by allocating all n bytes in cases where n is incorrect. +func ReadDataAt(r io.ReaderAt, n uint64, off int64) ([]byte, error) { + if int64(n) < 0 || n != uint64(int(n)) { + // n is too large to fit in int, so we can't allocate + // a buffer large enough. Treat this as a read failure. + return nil, io.ErrUnexpectedEOF + } + + if n < chunk { + buf := make([]byte, n) + _, err := r.ReadAt(buf, off) + if err != nil { + // io.SectionReader can return EOF for n == 0, + // but for our purposes that is a success. + if err != io.EOF || n > 0 { + return nil, err + } + } + return buf, nil + } + + var buf []byte + buf1 := make([]byte, chunk) + for n > 0 { + next := n + if next > chunk { + next = chunk + } + _, err := r.ReadAt(buf1[:next], off) + if err != nil { + return nil, err + } + buf = append(buf, buf1[:next]...) + n -= next + off += int64(next) + } + return buf, nil +} + +// SliceCap returns the capacity to use when allocating a slice. +// After the slice is allocated with the capacity, it should be +// built using append. This will avoid allocating too much memory +// if the capacity is large and incorrect. +// +// A negative result means that the value is always too big. +// +// The element type is described by passing a pointer to a value of that type. +// This would ideally use generics, but this code is built with +// the bootstrap compiler which need not support generics. +// We use a pointer so that we can handle slices of interface type. +func SliceCap(v any, c uint64) int { + if int64(c) < 0 || c != uint64(int(c)) { + return -1 + } + typ := reflect.TypeOf(v) + if typ.Kind() != reflect.Ptr { + panic("SliceCap called with non-pointer type") + } + size := uint64(typ.Elem().Size()) + if size > 0 && c > (1<<64-1)/size { + return -1 + } + if c*size > chunk { + c = uint64(chunk / size) + if c == 0 { + c = 1 + } + } + return int(c) +} diff --git a/src/internal/saferio/io_test.go b/src/internal/saferio/io_test.go new file mode 100644 index 0000000..356c9eb --- /dev/null +++ b/src/internal/saferio/io_test.go @@ -0,0 +1,136 @@ +// Copyright 2022 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package saferio + +import ( + "bytes" + "io" + "testing" +) + +func TestReadData(t *testing.T) { + const count = 100 + input := bytes.Repeat([]byte{'a'}, count) + + t.Run("small", func(t *testing.T) { + got, err := ReadData(bytes.NewReader(input), count) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, input) { + t.Errorf("got %v, want %v", got, input) + } + }) + + t.Run("large", func(t *testing.T) { + _, err := ReadData(bytes.NewReader(input), 10<<30) + if err == nil { + t.Error("large read succeeded unexpectedly") + } + }) + + t.Run("maxint", func(t *testing.T) { + _, err := ReadData(bytes.NewReader(input), 1<<62) + if err == nil { + t.Error("large read succeeded unexpectedly") + } + }) + + t.Run("small-EOF", func(t *testing.T) { + _, err := ReadData(bytes.NewReader(nil), chunk-1) + if err != io.EOF { + t.Errorf("ReadData = %v, want io.EOF", err) + } + }) + + t.Run("large-EOF", func(t *testing.T) { + _, err := ReadData(bytes.NewReader(nil), chunk+1) + if err != io.EOF { + t.Errorf("ReadData = %v, want io.EOF", err) + } + }) + + t.Run("large-UnexpectedEOF", func(t *testing.T) { + _, err := ReadData(bytes.NewReader(make([]byte, chunk)), chunk+1) + if err != io.ErrUnexpectedEOF { + t.Errorf("ReadData = %v, want io.ErrUnexpectedEOF", err) + } + }) +} + +func TestReadDataAt(t *testing.T) { + const count = 100 + input := bytes.Repeat([]byte{'a'}, count) + + t.Run("small", func(t *testing.T) { + got, err := ReadDataAt(bytes.NewReader(input), count, 0) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, input) { + t.Errorf("got %v, want %v", got, input) + } + }) + + t.Run("large", func(t *testing.T) { + _, err := ReadDataAt(bytes.NewReader(input), 10<<30, 0) + if err == nil { + t.Error("large read succeeded unexpectedly") + } + }) + + t.Run("maxint", func(t *testing.T) { + _, err := ReadDataAt(bytes.NewReader(input), 1<<62, 0) + if err == nil { + t.Error("large read succeeded unexpectedly") + } + }) + + t.Run("SectionReader", func(t *testing.T) { + // Reading 0 bytes from an io.SectionReader at the end + // of the section will return EOF, but ReadDataAt + // should succeed and return 0 bytes. + sr := io.NewSectionReader(bytes.NewReader(input), 0, 0) + got, err := ReadDataAt(sr, 0, 0) + if err != nil { + t.Fatal(err) + } + if len(got) > 0 { + t.Errorf("got %d bytes, expected 0", len(got)) + } + }) +} + +func TestSliceCap(t *testing.T) { + t.Run("small", func(t *testing.T) { + c := SliceCap((*int)(nil), 10) + if c != 10 { + t.Errorf("got capacity %d, want %d", c, 10) + } + }) + + t.Run("large", func(t *testing.T) { + c := SliceCap((*byte)(nil), 1<<30) + if c < 0 { + t.Error("SliceCap failed unexpectedly") + } else if c == 1<<30 { + t.Errorf("got capacity %d which is too high", c) + } + }) + + t.Run("maxint", func(t *testing.T) { + c := SliceCap((*byte)(nil), 1<<63) + if c >= 0 { + t.Errorf("SliceCap returned %d, expected failure", c) + } + }) + + t.Run("overflow", func(t *testing.T) { + c := SliceCap((*int64)(nil), 1<<62) + if c >= 0 { + t.Errorf("SliceCap returned %d, expected failure", c) + } + }) +} |