github.com/grafana/pyroscope@v1.18.0/pkg/symbolizer/reader.go (about)

     1  package symbolizer
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  
     8  	"github.com/klauspost/compress/gzip"
     9  	"github.com/klauspost/compress/zstd"
    10  )
    11  
    12  func NewReaderAtCloser(data []byte) interface {
    13  	io.ReadCloser
    14  	io.ReaderAt
    15  } {
    16  	bytesReader := bytes.NewReader(data)
    17  	return struct {
    18  		io.ReadCloser
    19  		io.ReaderAt
    20  	}{
    21  		ReadCloser: io.NopCloser(bytesReader),
    22  		ReaderAt:   bytesReader,
    23  	}
    24  }
    25  
    26  // memoryBuffer implements io.WriteSeeker for writing to an in-memory buffer with seeking capabilities
    27  type memoryBuffer struct {
    28  	data []byte
    29  	pos  int64
    30  }
    31  
    32  func newMemoryBuffer(initialCapacity int) *memoryBuffer {
    33  	capacity := initialCapacity
    34  	if capacity < 0 {
    35  		capacity = 0
    36  	}
    37  
    38  	return &memoryBuffer{
    39  		data: make([]byte, 0, capacity),
    40  		pos:  0,
    41  	}
    42  }
    43  
    44  func (m *memoryBuffer) Write(p []byte) (n int, err error) {
    45  	if m.pos > int64(len(m.data)) {
    46  		m.data = append(m.data, make([]byte, m.pos-int64(len(m.data)))...)
    47  	}
    48  
    49  	if m.pos+int64(len(p)) > int64(len(m.data)) {
    50  		m.data = append(m.data, make([]byte, m.pos+int64(len(p))-int64(len(m.data)))...)
    51  	}
    52  
    53  	n = copy(m.data[m.pos:], p)
    54  	m.pos += int64(n)
    55  	return n, nil
    56  }
    57  
    58  func (m *memoryBuffer) Seek(offset int64, whence int) (int64, error) {
    59  	var newPos int64
    60  	switch whence {
    61  	case io.SeekStart:
    62  		newPos = offset
    63  	case io.SeekCurrent:
    64  		newPos = m.pos + offset
    65  	case io.SeekEnd:
    66  		newPos = int64(len(m.data)) + offset
    67  	default:
    68  		return 0, fmt.Errorf("invalid whence: %d", whence)
    69  	}
    70  
    71  	if newPos < 0 {
    72  		return 0, fmt.Errorf("negative position: %d", newPos)
    73  	}
    74  
    75  	m.pos = newPos
    76  	return m.pos, nil
    77  }
    78  
    79  func (m *memoryBuffer) Bytes() []byte {
    80  	return m.data
    81  }
    82  
    83  func readAllWithLimit(r io.Reader, typ string, maxBytes int64) ([]byte, error) {
    84  	// maxBytes == 0 means unlimited body size
    85  	if maxBytes > 0 {
    86  		r = io.LimitReader(r, maxBytes+1) // +1 to detect if limit is exceeded
    87  	}
    88  
    89  	decompressed, err := io.ReadAll(r)
    90  	if err != nil {
    91  		return nil, fmt.Errorf("decompress %s data: %w", typ, err)
    92  	}
    93  
    94  	// Check if we hit the size limit
    95  	if maxBytes > 0 && int64(len(decompressed)) > maxBytes {
    96  		return nil, &ErrSymbolSizeBytesExceedsLimit{Limit: maxBytes}
    97  	}
    98  
    99  	return decompressed, nil
   100  }
   101  
   102  // detectCompression checks if data is compressed and decompresses it if needed
   103  func detectCompression(data []byte, maxBytes int64) ([]byte, error) {
   104  	if len(data) >= 2 && data[0] == 0x1f && data[1] == 0x8b {
   105  		r, err := gzip.NewReader(bytes.NewReader(data))
   106  		if err != nil {
   107  			return nil, fmt.Errorf("create gzip reader: %w", err)
   108  		}
   109  		defer r.Close()
   110  		return readAllWithLimit(r, "gzip", maxBytes)
   111  	}
   112  
   113  	// Check for zstd (magic bytes: 0x28, 0xb5, 0x2f, 0xfd)
   114  	if len(data) >= 4 && data[0] == 0x28 && data[1] == 0xb5 && data[2] == 0x2f && data[3] == 0xfd {
   115  		r, err := zstd.NewReader(bytes.NewReader(data))
   116  		if err != nil {
   117  			return nil, fmt.Errorf("create zstd reader: %w", err)
   118  		}
   119  		defer r.Close()
   120  		return readAllWithLimit(r, "zstd", maxBytes)
   121  	}
   122  
   123  	return data, nil
   124  }