zstd: option to decompress a single frame/ get total compressed frame length

[katexochen]

I'm trying to unpack a bunch of initrds, of which some are compressed zstd, others not (other compression or raw cpio), all concatenated together with some padding in between (format as expected by the kernel).

To unpack this, I either need to

• Get the size of the first zstd archive/frame, so I can limit the reader used by the decompressor to this range or
• Have the decompressor fail with an identifiable error type and with file seek at the right position (currently points to the byte after the expected magic bytes, so one have to manually seek back 4 bytes)

Is this something you'd be interested to support?

(My current implementation iterates over the block headers to calculate the total size of the frame)

[klauspost]

There is a function that allows you to decode a frame header -and first block header: https://pkg.go.dev/github.com/klauspost/compress/zstd#Header

You will need to read the first bytes yourself, but that should be pretty trivial, and you can provide both (maybe using io.MultiReader) if you decide to move on.

Note that frame size is optional, so check HasFCS.

I can't really tell what seeking you are trying to do, but you would have to do that yourself, since it sounds like a very case-specific thing.

[klauspost]

(compressed size is not recorded in the header - it is assumed you manage that yourself)

[katexochen]

compressed size is not recorded in the header

Yes, that's exactly my problem. 😅 I don't know the compressed size. So what I'm currently doing is:

func main() {
       f, err := os.OpenFile(path, os.O_RDONLY, 0)
	... // checked the magic byte, it's zstd
	decompressZSTD(f, someNewFile)
        ... // remove padding, next magic byte is gzip
        decompressGZIP(f, anotherNewFile)
}
	

func decompressZSTD(in io.ReadSeeker, out io.Writer) error {
	frameSize, err := zstdFrameSize(in)
	if err != nil {
		return fmt.Errorf("getting zstd frame size: %w", err)
	}
	limitIn := io.LimitReader(in, int64(frameSize))
	d, err := zstd.NewReader(limitIn)
	if err != nil {
		return err
	}
	defer d.Close()
	if _, err := io.Copy(out, d); err != nil {
		return fmt.Errorf("decompressing zstd: %w", err)
	}
	fmt.Println("successfully decompressed zstd")
	return nil
}

func zstdFrameSize(in io.ReadSeeker) (int, error) {
	const zsdtBlockHeaderSize = 3

	pos, err := in.Seek(0, io.SeekCurrent)
	if err != nil {
		return 0, fmt.Errorf("getting current position: %w", err)
	}
	defer in.Seek(pos, io.SeekStart)

	var header zstd.Header
	buf, err := peek(in, 18)
	if err != nil {
		return 0, fmt.Errorf("reading zstd header: %w", err)
	}
	if err := header.Decode(buf); err != nil {
		return 0, fmt.Errorf("decoding zstd header: %w", err)
	}
	size := header.HeaderSize
	in.Seek(int64(header.HeaderSize), io.SeekCurrent)
	if header.HasCheckSum {
		size += 4
	}

	for {
		bhBytes, err := peek(in, zsdtBlockHeaderSize)
		if err != nil {
			return 0, fmt.Errorf("reading zstd block header: %w", err)
		}
		// https://github.com/klauspost/compress/blob/2a46d6bf5d0fb5d9f44b815438ce43470706f73f/zstd/blockdec.go#L139
		blockHeader := uint32(bhBytes[0]) | (uint32(bhBytes[1]) << 8) | (uint32(bhBytes[2]) << 16)
		blockType := zstdBlockType((blockHeader >> 1) & 3)
		blockIsLast := blockHeader&1 != 0
		blockSize := zsdtBlockHeaderSize
		switch blockType {
		case zsdtBlockTypeRLE:
			blockSize += 1
		case zsdtBlockTypeCompressed:
			fallthrough
		case zsdtBlockTypeRaw:
			blockSize += int(blockHeader >> 3)
		default:
			panic("Invalid block type")
		}
		size += blockSize
		in.Seek(int64(blockSize), io.SeekCurrent)
		if blockIsLast {
			break
		}
	}

	return size, nil
}

It would be nice if I could do

func decompressZSTDfunc(in io.ReadSeeker, out io.Writer) error {
	d, err := zstd.NewReader(in)
	if err != nil {
		return err
	}
	defer d.Close()
	if _, err := io.Copy(out, d);  errors.Is(zstd.ErrNonZstdTrashAfterFrame) {
		return nil 
        else err != nil {
		return fmt.Errorf("decompressing zstd: %w", err)
	}
	fmt.Println("successfully decompressed zstd")
	return nil
}

Which would require

• An error type to check for (zstd.ErrNonZstdTrashAfterFrame)
• zstd lib to only peek at the magic bytes of the next frame. The way it is currently implemented, compressGZIP will not find the magic bytes for gzip as they were already read by the zstd.Compressor.

I understand if this is a specific use case and might not be of interest for the upstream library.

[katexochen]

This could also be implemented as decompressor option: func OnlyFirstFrame(b bool) DOption