github.com/google/syzkaller@v0.0.0-20251211124644-a066d2bc4b02/pkg/image/compression.go (about)

     1  // Copyright 2022 syzkaller project authors. All rights reserved.
     2  // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
     3  
     4  package image
     5  
     6  import (
     7  	"bytes"
     8  	"compress/zlib"
     9  	"encoding/base64"
    10  	"fmt"
    11  	"io"
    12  	"sync/atomic"
    13  )
    14  
    15  var (
    16  	// Total amount of images in memory and consumed memory (in bytes).
    17  	// Currently maintained only by the optimized implementation.
    18  	// Cannot import stats package due to import cycles.
    19  	StatImages atomic.Int64
    20  	StatMemory atomic.Int64
    21  )
    22  
    23  func Compress(rawData []byte) []byte {
    24  	var buffer bytes.Buffer
    25  	zlibWriter := zlib.NewWriter(&buffer)
    26  
    27  	_, err := zlibWriter.Write(rawData)
    28  	if err != nil {
    29  		panic(fmt.Sprintf("could not compress with zlib: %v", err))
    30  	}
    31  
    32  	err = zlibWriter.Close()
    33  	if err != nil {
    34  		panic(fmt.Sprintf("could not finalize compression with zlib: %v", err))
    35  	}
    36  
    37  	return buffer.Bytes()
    38  }
    39  
    40  func MustDecompress(compressed []byte) (data []byte, dtor func()) {
    41  	if len(compressed) == 0 {
    42  		return nil, func() {}
    43  	}
    44  	return mustDecompress(compressed)
    45  }
    46  
    47  func DecompressCheck(compressed []byte) error {
    48  	return decompressWriter(io.Discard, compressed)
    49  }
    50  
    51  func decompressWriter(w io.Writer, compressed []byte) error {
    52  	if len(compressed) == 0 {
    53  		return nil
    54  	}
    55  	zlibReader, err := zlib.NewReader(bytes.NewReader(compressed))
    56  	if err != nil {
    57  		return fmt.Errorf("could not initialise zlib: %w", err)
    58  	}
    59  
    60  	if _, err := io.Copy(w, zlibReader); err != nil {
    61  		return fmt.Errorf("could not read data with zlib: %w", err)
    62  	}
    63  
    64  	return zlibReader.Close()
    65  }
    66  
    67  func DecodeB64(b64Data []byte) ([]byte, error) {
    68  	decoder := base64.NewDecoder(base64.StdEncoding, bytes.NewReader(b64Data))
    69  	rawData, err := io.ReadAll(decoder)
    70  	if err != nil {
    71  		return nil, fmt.Errorf("could not decode Base64: %w", err)
    72  	}
    73  	return rawData, nil
    74  }
    75  
    76  func EncodeB64(rawData []byte) []byte {
    77  	var buf bytes.Buffer
    78  	encoder := base64.NewEncoder(base64.StdEncoding, &buf)
    79  	_, err := encoder.Write(rawData)
    80  	if err != nil {
    81  		panic(fmt.Sprintf("could not encode Base64: %v", err))
    82  	}
    83  	err = encoder.Close()
    84  	if err != nil {
    85  		panic(fmt.Sprintf("could not finalize encoding to Base64: %v", err))
    86  	}
    87  	return buf.Bytes()
    88  }