futriix/internal/compression/compression.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))
}