62 lines
1.3 KiB
Go
62 lines
1.3 KiB
Go
package protocol
|
|
|
|
import (
|
|
"bytes"
|
|
"compress/gzip"
|
|
"io"
|
|
|
|
"spark-server/internal/pkg/errs"
|
|
)
|
|
|
|
func Compress(body []byte) ([]byte, error) {
|
|
var buf bytes.Buffer
|
|
w := gzip.NewWriter(&buf)
|
|
if _, err := w.Write(body); err != nil {
|
|
_ = w.Close()
|
|
return nil, err
|
|
}
|
|
if err := w.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
return buf.Bytes(), nil
|
|
}
|
|
|
|
func Decompress(body []byte, limit int) ([]byte, error) {
|
|
r, err := gzip.NewReader(bytes.NewReader(body))
|
|
if err != nil {
|
|
return nil, errs.ErrInvalidFrame
|
|
}
|
|
defer r.Close()
|
|
out, err := io.ReadAll(io.LimitReader(r, int64(limit)+1))
|
|
if err != nil {
|
|
return nil, errs.ErrInvalidFrame
|
|
}
|
|
if len(out) > limit {
|
|
return nil, errs.ErrFrameTooLarge
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// PrepareBody gzip-compresses body when it is at least minCompress bytes and smaller after gzip.
|
|
func PrepareBody(body []byte, minCompress int) ([]byte, uint16, error) {
|
|
if minCompress <= 0 || len(body) < minCompress {
|
|
return body, 0, nil
|
|
}
|
|
out, err := Compress(body)
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
if len(out) >= len(body) {
|
|
return body, 0, nil
|
|
}
|
|
return out, FlagCompress, nil
|
|
}
|
|
|
|
// DecodeBody inflates a gzip body when FlagCompress is set.
|
|
func DecodeBody(body []byte, flags uint16, limit int) ([]byte, error) {
|
|
if flags&FlagCompress == 0 {
|
|
return body, nil
|
|
}
|
|
return Decompress(body, limit)
|
|
}
|