1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
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)
}
|
