tinymq/conn/util/compress.go

package util

import (
	"bytes"
	"compress/gzip"
	"io"
	"log"
)

// 压缩数据
// 注意：如果没有启用压缩，和数据为空不要调用该函数
func CompressData(data []byte) (out []byte, ok bool) {
	var buf bytes.Buffer
	zw := gzip.NewWriter(&buf)
	if _, err := zw.Write(data); err != nil {
		log.Println("[CompressData Write ERROR]", err)
		return data, false
	}
	if err := zw.Close(); err != nil {
		log.Println("[CompressData Close ERROR]", err)
		return data, false
	}
	compressedData := buf.Bytes()
	if len(compressedData) >= len(data) {
		// log.Println("user origin data:", string(data))
		return data, false
	} else {
		return compressedData, true
	}
}

// 解压缩
func UncompressData(compressedData []byte) (data []byte, ok bool) {
	// 没有压缩，直接返回
	if compressedData[0]&0x01 != 0 {
		zr, err := gzip.NewReader(bytes.NewReader(compressedData[1:]))
		if err != nil {
			log.Println("[UncompressData NewReader ERROR]", err)
			return compressedData, false
		}
		defer zr.Close()
		uncompressedData, err := io.ReadAll(zr)
		if err != nil {
			log.Println("[UncompressData ReadAll ERROR]", err)
			return compressedData, false
		}
		return uncompressedData, true
	} else {
		return compressedData[1:], true
	}
}