github.com/mvdan/u-root-coreutils@v0.0.0-20230122170626-c2eef2898555/pkg/uio/archivereader.go (about)

     1  // Copyright 2021 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package uio
     6  
     7  import (
     8  	"bytes"
     9  	"errors"
    10  	"io"
    11  
    12  	"github.com/pierrec/lz4/v4"
    13  )
    14  
    15  const (
    16  	// preReadSizeBytes is the num of bytes pre-read from a io.Reader that will
    17  	// be used to match against archive header.
    18  	defaultArchivePreReadSizeBytes = 1024
    19  )
    20  
    21  var ErrPreReadError = errors.New("pre-read nothing")
    22  
    23  // ArchiveReader reads from a io.Reader, decompresses source bytes
    24  // when applicable.
    25  //
    26  // It allows probing for multiple archive format, while still able
    27  // to read from beginning, by pre-reading a small number of bytes.
    28  //
    29  // Always use newArchiveReader to initialize.
    30  type ArchiveReader struct {
    31  	// src is where we read source bytes.
    32  	src io.Reader
    33  	// buf stores pre-read bytes from original io.Reader. Archive format
    34  	// detection will be done against it.
    35  	buf []byte
    36  
    37  	// preReadSizeBytes is how many bytes we pre-read for magic number
    38  	// matching for each archive type. This should be greater than or
    39  	// equal to the largest header frame size of each supported archive
    40  	// format.
    41  	preReadSizeBytes int
    42  }
    43  
    44  func NewArchiveReader(r io.Reader) (ArchiveReader, error) {
    45  	ar := ArchiveReader{
    46  		src: r,
    47  		// Randomly chosen, should be enough for most types:
    48  		//
    49  		// e.g. gzip with 10 byte header, lz4 with a header size
    50  		// between 7 and 19 bytes.
    51  		preReadSizeBytes: defaultArchivePreReadSizeBytes,
    52  	}
    53  	pbuf := make([]byte, ar.preReadSizeBytes)
    54  
    55  	nr, err := io.ReadFull(r, pbuf)
    56  	// In case the image is smaller pre-read block size, 1kb for now.
    57  	// Ever possible ? probably not in case a compression is needed!
    58  	ar.buf = pbuf[:nr]
    59  	if err == io.EOF {
    60  		// If we could not pre-read anything, we can't determine if
    61  		// it is a compressed file.
    62  		ar.src = io.MultiReader(bytes.NewReader(pbuf[:nr]), r)
    63  		return ar, ErrPreReadError
    64  	}
    65  
    66  	// Try each supported compression type, return upon first match.
    67  
    68  	// Try lz4.
    69  	// magic number error will be thrown if source is not a lz4 archive.
    70  	// e.g. "lz4: bad magic number".
    71  	if ok, err := lz4.ValidFrameHeader(ar.buf); err == nil && ok {
    72  		ar.src = lz4.NewReader(io.MultiReader(bytes.NewReader(ar.buf), r))
    73  		return ar, nil
    74  	}
    75  
    76  	// Try other archive types here, gzip, xz, etc when needed.
    77  
    78  	// Last resort, read as is.
    79  	ar.src = io.MultiReader(bytes.NewReader(ar.buf), r)
    80  	return ar, nil
    81  }
    82  
    83  func (ar ArchiveReader) Read(p []byte) (n int, err error) {
    84  	return ar.src.Read(p)
    85  }