483 lines
22 KiB
Go
483 lines
22 KiB
Go
/*
|
||
* Copyright 2026 Safronov Grigorii
|
||
*
|
||
* Licensed under the CDDL, Version 1.0 (the "License");
|
||
* you may not use this file except in compliance with the License.
|
||
*
|
||
* You may obtain a copy of the License at
|
||
* https://opensource.org/licenses/CDDL-1.0
|
||
*/
|
||
|
||
// Файл: internal/compression/compression.go
|
||
// Назначение: Реализация сжатия данных с использованием различных алгоритмов.
|
||
// Поддерживаемые алгоритмы: Snappy, Brotli, Zstd.
|
||
// Обеспечивает прозрачное сжатие/распаковку для документов.
|
||
// Защита от collision magic bytes через усиленную проверку заголовка.
|
||
//
|
||
// ИСПРАВЛЕНО:
|
||
// 1. Compress больше не паникует при config == nil — возвращает data
|
||
// как есть (это логично: без конфигурации сжатие выключено).
|
||
// 2. Compress учитывает оверхед заголовка при решении «сжимать или нет».
|
||
// Раньше проверялось `len(compressed) >= len(data)`, но не учитывалось,
|
||
// что заголовок (13 или 14 байт) может сделать итог больше исходных
|
||
// данных. Теперь сравнение: `len(compressed) + headerSize >= len(data)`.
|
||
// 3. Level для brotli теперь применяется (brotli.NewWriterLevel), а не
|
||
// игнорируется. Диапазон уровней brotli — 0..11, поэтому входной
|
||
// Level (1..9) нормализуется.
|
||
// 4. Введены константы MagicV1HeaderSize и MagicV2HeaderSize. Раньше
|
||
// в decompressV1 стояло магическое число 13, а MinCompressedSize = 14
|
||
// относилось к V2 — легко было перепутать.
|
||
// 5. zstd-энкодер и декодер кэшируются (sync.Map). zstd.Encoder.EncodeAll
|
||
// и zstd.Decoder.DecodeAll потокобезопасны, поэтому кэш безопасен.
|
||
// Раньше энкодер/декодер создавались на каждый вызов Compress /
|
||
// doDecompress — это дорого на горячем пути.
|
||
// 6. Логика парсинга заголовка (V1/V2) вынесена в parseHeader —
|
||
// убрано дублирование между Decompress, DecompressAuto, IsCompressed,
|
||
// GetCompressionType.
|
||
// 7. В doDecompress для CompressionNone добавлена проверка размера
|
||
// (защита в глубину от malformed-входа).
|
||
// 8. bytes.Equal(data[0:N], magic) заменено на bytes.HasPrefix —
|
||
// идиоматичнее и читается лучше.
|
||
// 9. Лимит originalSize вынесен в константу MaxOriginalSize.
|
||
// 10. Обновлён комментарий про LZ4: LZ4 удалён, остались snappy/brotli/zstd.
|
||
//
|
||
// ИСПРАВЛЕНО (компиляция):
|
||
// В Compress, в ветке case "zstd", переменная err была объявлена,
|
||
// но не использовалась: zstd.Encoder.EncodeAll не возвращает ошибку.
|
||
// Убрано объявление `encoder, err := ...` — теперь `encoder, e := ...`
|
||
// с локальной переменной e, которая проверяется тут же. Внешняя err
|
||
// по-прежнему используется в других ветках (snappy/brotli) и в
|
||
// финальной проверке оверхеда заголовка.
|
||
|
||
package compression
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/binary"
|
||
"fmt"
|
||
"sync"
|
||
|
||
"github.com/andybalholm/brotli"
|
||
"github.com/golang/snappy"
|
||
"github.com/klauspost/compress/zstd"
|
||
)
|
||
|
||
// =============================================================================
|
||
// КОНФИГУРАЦИЯ
|
||
// =============================================================================
|
||
|
||
// Config представляет конфигурацию сжатия.
|
||
type Config struct {
|
||
Enabled bool // Включено ли сжатие
|
||
Algorithm string // Алгоритм сжатия: snappy, brotli, zstd
|
||
Level int // Уровень сжатия (1-9)
|
||
MinSize int // Минимальный размер для сжатия (байт)
|
||
}
|
||
|
||
// =============================================================================
|
||
// MAGIC NUMBERS И КОНСТАНТЫ ФОРМАТА
|
||
// =============================================================================
|
||
|
||
// MagicNumber используется для идентификации сжатых данных (V1).
|
||
var MagicNumber = []byte{0x46, 0x54, 0x52, 0x53} // "FTRS" - Futriis
|
||
|
||
// MagicNumberV2 — расширенный magic с версией для предотвращения collision.
|
||
var MagicNumberV2 = []byte{0x46, 0x54, 0x52, 0x53, 0x02} // "FTRS\x02"
|
||
|
||
const (
|
||
// MagicV1HeaderSize — размер заголовка формата V1:
|
||
// 4 (magic) + 1 (type) + 8 (size) = 13 байт.
|
||
MagicV1HeaderSize = 4 + 1 + 8
|
||
|
||
// MagicV2HeaderSize — размер заголовка формата V2:
|
||
// 5 (magic) + 1 (type) + 8 (size) = 14 байт.
|
||
MagicV2HeaderSize = 5 + 1 + 8
|
||
|
||
// MinCompressedSize — минимальный размер валидного сжатого блока.
|
||
// Равен размеру заголовка V2 (самый свежий формат).
|
||
// Оставлен экспортируемым для обратной совместимости с внешним кодом,
|
||
// который мог его использовать.
|
||
MinCompressedSize = MagicV2HeaderSize
|
||
|
||
// MaxOriginalSize — верхняя граница originalSize, которую мы считаем
|
||
// разумной. 100 GiB. Защита от мусорных/враждебных заголовков.
|
||
MaxOriginalSize = uint64(100) * 1024 * 1024 * 1024
|
||
)
|
||
|
||
// CompressionType определяет тип сжатия.
|
||
type CompressionType byte
|
||
|
||
const (
|
||
CompressionNone CompressionType = 0x00
|
||
CompressionSnappy CompressionType = 0x01
|
||
CompressionBrotli CompressionType = 0x02
|
||
CompressionZstd CompressionType = 0x03
|
||
)
|
||
|
||
// =============================================================================
|
||
// КЭШ ZSTD ENCODER / DECODER
|
||
// =============================================================================
|
||
//
|
||
// zstd.Encoder.EncodeAll и zstd.Decoder.DecodeAll потокобезопасны
|
||
// (см. документацию klauspost/compress/zstd). Поэтому один общий
|
||
// encoder/decoder можно безопасно переиспользовать из многих горутин.
|
||
//
|
||
// Кэшируем по уровню для encoder и один общий для decoder.
|
||
|
||
var (
|
||
zstdEncoders sync.Map // map[zstd.EncoderLevel]*zstd.Encoder
|
||
zstdDecoder *zstd.Decoder
|
||
zstdDecOnce sync.Once
|
||
zstdDecErr error
|
||
)
|
||
|
||
// getZstdEncoder возвращает кэшированный zstd-энкодер для уровня.
|
||
func getZstdEncoder(level zstd.EncoderLevel) (*zstd.Encoder, error) {
|
||
if enc, ok := zstdEncoders.Load(level); ok {
|
||
return enc.(*zstd.Encoder), nil
|
||
}
|
||
enc, err := zstd.NewWriter(nil, zstd.WithEncoderLevel(level))
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to create zstd encoder: %v", err)
|
||
}
|
||
actual, _ := zstdEncoders.LoadOrStore(level, enc)
|
||
// Если другая горутина успела первой — закрываем свой encoder.
|
||
if actual.(*zstd.Encoder) != enc {
|
||
enc.Close()
|
||
}
|
||
return actual.(*zstd.Encoder), nil
|
||
}
|
||
|
||
// getZstdDecoder возвращает общий zstd-декодер (создаётся один раз).
|
||
func getZstdDecoder() (*zstd.Decoder, error) {
|
||
zstdDecOnce.Do(func() {
|
||
zstdDecoder, zstdDecErr = zstd.NewReader(nil)
|
||
})
|
||
return zstdDecoder, zstdDecErr
|
||
}
|
||
|
||
// =============================================================================
|
||
// ПАРСИНГ ЗАГОЛОВКА
|
||
// =============================================================================
|
||
|
||
// headerInfo описывает разобранный заголовок сжатого блока.
|
||
type headerInfo struct {
|
||
version int
|
||
compType CompressionType
|
||
originalSize uint64
|
||
headerSize int
|
||
}
|
||
|
||
// parseHeader пытается разобрать заголовок сжатого блока (V2 или V1).
|
||
// Возвращает nil, если данные не являются сжатым блоком.
|
||
//
|
||
// ВНИМАНИЕ: для V1 проверяется валидность compType — это защита от
|
||
// ложных срабатываний, когда несжатые данные случайно начинаются
|
||
// с байт "FTRS".
|
||
func parseHeader(data []byte) *headerInfo {
|
||
// V2: 5-байтовый magic + 1 байт type + 8 байт size.
|
||
if len(data) >= MagicV2HeaderSize && bytes.HasPrefix(data, MagicNumberV2) {
|
||
compType := CompressionType(data[5])
|
||
if isValidCompressionType(compType) {
|
||
return &headerInfo{
|
||
version: 2,
|
||
compType: compType,
|
||
originalSize: binary.LittleEndian.Uint64(data[6:14]),
|
||
headerSize: MagicV2HeaderSize,
|
||
}
|
||
}
|
||
// Magic совпал, но тип невалиден — не считаем это сжатым блоком.
|
||
return nil
|
||
}
|
||
|
||
// V1: 4-байтовый magic + 1 байт type + 8 байт size.
|
||
if len(data) >= MagicV1HeaderSize && bytes.HasPrefix(data, MagicNumber) {
|
||
compType := CompressionType(data[4])
|
||
if isValidCompressionType(compType) {
|
||
return &headerInfo{
|
||
version: 1,
|
||
compType: compType,
|
||
originalSize: binary.LittleEndian.Uint64(data[5:13]),
|
||
headerSize: MagicV1HeaderSize,
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// isValidCompressionType проверяет, что байт соответствует известному
|
||
// типу сжатия. Используется как защита от ложных срабатываний magic.
|
||
func isValidCompressionType(t CompressionType) bool {
|
||
switch t {
|
||
case CompressionNone, CompressionSnappy, CompressionBrotli, CompressionZstd:
|
||
return true
|
||
default:
|
||
return false
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// СЖАТИЕ
|
||
// =============================================================================
|
||
|
||
// Compress сжимает данные с использованием указанного алгоритма.
|
||
//
|
||
// ИСПРАВЛЕНО:
|
||
// - config == nil больше не приводит к панике — возвращаем data как есть.
|
||
// - Учитывается оверхед заголовка при решении «сжимать или нет».
|
||
// - Level применяется не только к zstd, но и к brotli.
|
||
//
|
||
// ИСПРАВЛЕНО (компиляция):
|
||
// В ветке case "zstd" убрано объявление err, которое не использовалось:
|
||
// zstd.Encoder.EncodeAll не возвращает ошибку. Ошибка возможна только
|
||
// при создании encoder через getZstdEncoder — она проверяется тут же.
|
||
func Compress(data []byte, config *Config) ([]byte, error) {
|
||
// config == nil трактуем как «сжатие выключено».
|
||
if config == nil || !config.Enabled {
|
||
return data, nil
|
||
}
|
||
|
||
// MinSize <= 0 означает «сжимать всё, включая пустые данные».
|
||
// Если MinSize > 0 и данные меньше — не сжимаем.
|
||
if config.MinSize > 0 && len(data) < config.MinSize {
|
||
return data, nil
|
||
}
|
||
|
||
// Пустые данные сжимать бессмысленно.
|
||
if len(data) == 0 {
|
||
return data, nil
|
||
}
|
||
|
||
var compressed []byte
|
||
var compType CompressionType
|
||
|
||
switch config.Algorithm {
|
||
case "snappy":
|
||
compressed = snappy.Encode(nil, data)
|
||
compType = CompressionSnappy
|
||
|
||
case "brotli":
|
||
// brotli использует уровни 0..11. Входной Level — 1..9
|
||
// (см. Config.Level). Нормализуем: <=1 -> 1, >=9 -> 11 (best).
|
||
// Это даёт предсказуемое поведение: пользователь ставит 9 —
|
||
// получает максимальное сжатие.
|
||
level := normalizeBrotliLevel(config.Level)
|
||
buf := bytes.NewBuffer(nil)
|
||
writer := brotli.NewWriterLevel(buf, level)
|
||
|
||
if _, err := writer.Write(data); err != nil {
|
||
// writer.Close() стоит вызвать даже при ошибке Write,
|
||
// чтобы освободить ресурсы, но ошибку Close игнорируем —
|
||
// основная уже есть.
|
||
_ = writer.Close()
|
||
return nil, fmt.Errorf("brotli write failed: %v", err)
|
||
}
|
||
if err := writer.Close(); err != nil {
|
||
return nil, fmt.Errorf("brotli close failed: %v", err)
|
||
}
|
||
compressed = buf.Bytes()
|
||
compType = CompressionBrotli
|
||
|
||
case "zstd":
|
||
var encoderLevel zstd.EncoderLevel
|
||
switch {
|
||
case config.Level <= 1:
|
||
encoderLevel = zstd.SpeedFastest
|
||
case config.Level <= 3:
|
||
encoderLevel = zstd.SpeedDefault
|
||
case config.Level <= 6:
|
||
encoderLevel = zstd.SpeedBetterCompression
|
||
default:
|
||
encoderLevel = zstd.SpeedBestCompression
|
||
}
|
||
|
||
// ИСПРАВЛЕНО: вместо `encoder, err := ...` используем локальную e,
|
||
// чтобы не объявлять неиспользуемую err. getZstdEncoder возвращает
|
||
// (*zstd.Encoder, error), и эта ошибка проверяется тут же.
|
||
encoder, e := getZstdEncoder(encoderLevel)
|
||
if e != nil {
|
||
return nil, e
|
||
}
|
||
|
||
// EncodeAll потокобезопасен, используем общий encoder.
|
||
compressed = encoder.EncodeAll(data, nil)
|
||
compType = CompressionZstd
|
||
|
||
default:
|
||
return nil, fmt.Errorf("unsupported compression algorithm: %s", config.Algorithm)
|
||
}
|
||
|
||
// ИСПРАВЛЕНО: учитываем оверхед заголовка V2.
|
||
// Если сжатые данные + заголовок не меньше исходных —
|
||
// сжатие не имеет смысла, возвращаем исходные данные.
|
||
if len(compressed)+MagicV2HeaderSize >= len(data) {
|
||
return data, nil
|
||
}
|
||
|
||
// Формируем заголовок V2: 5 (magic) + 1 (type) + 8 (size).
|
||
header := make([]byte, MagicV2HeaderSize)
|
||
copy(header[0:5], MagicNumberV2)
|
||
header[5] = byte(compType)
|
||
binary.LittleEndian.PutUint64(header[6:14], uint64(len(data)))
|
||
|
||
result := make([]byte, 0, len(header)+len(compressed))
|
||
result = append(result, header...)
|
||
result = append(result, compressed...)
|
||
|
||
return result, nil
|
||
}
|
||
|
||
// normalizeBrotliLevel приводит входной уровень (1..9) к диапазону
|
||
// brotli (0..11). Значения <=0 и >9 трактуются разумно:
|
||
// - level <= 0 -> brotli.DefaultCompression
|
||
// - level 1..9 -> линейное отображение: 1->1, 9->11
|
||
// - level > 9 -> 11 (максимум)
|
||
func normalizeBrotliLevel(level int) int {
|
||
if level <= 0 {
|
||
return brotli.DefaultCompression
|
||
}
|
||
if level >= 9 {
|
||
return 11
|
||
}
|
||
// Линейное отображение 1..9 -> 1..11.
|
||
// 1->1, 2->2, ..., 9->11. Формула: 1 + (level-1)*10/8.
|
||
return 1 + (level-1)*10/8
|
||
}
|
||
|
||
// =============================================================================
|
||
// РАСПАКОВКА
|
||
// =============================================================================
|
||
|
||
// Decompress распаковывает данные.
|
||
// Поддерживает оба формата (V1 и V2), усиленная валидация.
|
||
func Decompress(data []byte) ([]byte, error) {
|
||
info := parseHeader(data)
|
||
if info == nil {
|
||
return nil, fmt.Errorf("invalid magic number")
|
||
}
|
||
return decompressWithInfo(data, info)
|
||
}
|
||
|
||
// decompressWithInfo выполняет распаковку, используя разобранный заголовок.
|
||
func decompressWithInfo(data []byte, info *headerInfo) ([]byte, error) {
|
||
if len(data) < info.headerSize {
|
||
return nil, fmt.Errorf("data too short for compressed format V%d", info.version)
|
||
}
|
||
|
||
// Проверка разумности размера.
|
||
if info.originalSize == 0 || info.originalSize > MaxOriginalSize {
|
||
return nil, fmt.Errorf("invalid original size: %d", info.originalSize)
|
||
}
|
||
|
||
compressedData := data[info.headerSize:]
|
||
if len(compressedData) == 0 {
|
||
return nil, fmt.Errorf("no compressed data")
|
||
}
|
||
|
||
return doDecompress(info.compType, compressedData, info.originalSize)
|
||
}
|
||
|
||
// doDecompress выполняет распаковку по типу.
|
||
//
|
||
// ИСПРАВЛЕНО:
|
||
// - для CompressionNone добавлена проверка размера (защита в глубину);
|
||
// - zstd-декодер берётся из кэша, а не создаётся на каждый вызов.
|
||
func doDecompress(compType CompressionType, compressedData []byte, originalSize uint64) ([]byte, error) {
|
||
var decompressed []byte
|
||
var err error
|
||
|
||
switch compType {
|
||
case CompressionSnappy:
|
||
decompressed, err = snappy.Decode(nil, compressedData)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("snappy decode failed: %v", err)
|
||
}
|
||
|
||
case CompressionBrotli:
|
||
reader := brotli.NewReader(bytes.NewReader(compressedData))
|
||
buf := bytes.NewBuffer(nil)
|
||
_, err = buf.ReadFrom(reader)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("brotli decode failed: %v", err)
|
||
}
|
||
decompressed = buf.Bytes()
|
||
|
||
case CompressionZstd:
|
||
decoder, err := getZstdDecoder()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
// DecodeAll потокобезопасен, используем общий decoder.
|
||
decompressed, err = decoder.DecodeAll(compressedData, nil)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("zstd decode failed: %v", err)
|
||
}
|
||
|
||
case CompressionNone:
|
||
// ИСПРАВЛЕНО: проверяем согласованность размера даже для None.
|
||
// Это защита в глубину от malformed-входа. Сейчас Compress не
|
||
// пишет CompressionNone в заголовок, но на входе может оказаться
|
||
// подделанный блок.
|
||
if uint64(len(compressedData)) != originalSize {
|
||
return nil, fmt.Errorf("size mismatch for uncompressed: expected %d, got %d",
|
||
originalSize, len(compressedData))
|
||
}
|
||
return compressedData, nil
|
||
|
||
default:
|
||
return nil, fmt.Errorf("unsupported compression type: %d", compType)
|
||
}
|
||
|
||
// Проверяем, что распакованный размер совпадает с заявленным.
|
||
if uint64(len(decompressed)) != originalSize {
|
||
return nil, fmt.Errorf("size mismatch: expected %d, got %d", originalSize, len(decompressed))
|
||
}
|
||
|
||
return decompressed, nil
|
||
}
|
||
|
||
// DecompressAuto автоматически определяет, сжаты ли данные, и распаковывает при необходимости.
|
||
// Усиленная проверка для предотвращения collision.
|
||
func DecompressAuto(data []byte) ([]byte, error) {
|
||
info := parseHeader(data)
|
||
if info == nil {
|
||
// Данные не сжаты.
|
||
return data, nil
|
||
}
|
||
return decompressWithInfo(data, info)
|
||
}
|
||
|
||
// IsCompressed проверяет, сжаты ли данные (V1 или V2).
|
||
func IsCompressed(data []byte) bool {
|
||
return parseHeader(data) != nil
|
||
}
|
||
|
||
// GetCompressionType возвращает тип сжатия данных.
|
||
// Возвращает CompressionNone, если данные не сжаты.
|
||
func GetCompressionType(data []byte) CompressionType {
|
||
info := parseHeader(data)
|
||
if info == nil {
|
||
return CompressionNone
|
||
}
|
||
return info.compType
|
||
}
|
||
|
||
// =============================================================================
|
||
// СТАТИСТИКА
|
||
// =============================================================================
|
||
|
||
// GetCompressionRatio возвращает коэффициент сжатия как долю:
|
||
// - 1.0 — данные не уменьшились (или выросли);
|
||
// - 0.5 — сжатые данные в два раза меньше исходных;
|
||
// - 0.0 — теоретический предел (пустые данные после сжатия).
|
||
//
|
||
// ВНИМАНИЕ: это именно доля compressed/original, а не «во сколько раз
|
||
// сжали». Если нужно «во сколько раз», используйте 1.0 / ratio.
|
||
func GetCompressionRatio(original, compressed []byte) float64 {
|
||
if len(original) == 0 {
|
||
return 1.0
|
||
}
|
||
return float64(len(compressed)) / float64(len(original))
|
||
} |