From dde07638408eace87387b5ad237c6b12d6917045 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Thu, 18 Jun 2020 11:32:38 +0300
Subject: [PATCH] network: fix compression handling functions

Turns out, Neo uses block compression and not a streaming lz4 format. And it
doesn't contain uncompressed payload size which makes guessing it somewhat
suboptimal.
---
 pkg/network/compress.go | 24 ++++++++++--------------
 1 file changed, 10 insertions(+), 14 deletions(-)

diff --git a/pkg/network/compress.go b/pkg/network/compress.go
index 7b6d6c3d0..ddde5e9d7 100644
--- a/pkg/network/compress.go
+++ b/pkg/network/compress.go
@@ -1,33 +1,29 @@
 package network
 
 import (
-	"bytes"
-	"io"
-
 	"github.com/pierrec/lz4"
 )
 
 // compress compresses bytes using lz4.
 func compress(source []byte) ([]byte, error) {
-	dest := new(bytes.Buffer)
-	w := lz4.NewWriter(dest)
-	_, err := io.Copy(w, bytes.NewReader(source))
+	dest := make([]byte, lz4.CompressBlockBound(len(source)))
+	size, err := lz4.CompressBlock(source, dest, nil)
 	if err != nil {
 		return nil, err
 	}
-	if w.Close() != nil {
-		return nil, err
-	}
-	return dest.Bytes(), nil
+	return dest[:size], nil
 }
 
 // decompress decompresses bytes using lz4.
 func decompress(source []byte) ([]byte, error) {
-	dest := new(bytes.Buffer)
-	r := lz4.NewReader(bytes.NewReader(source))
-	_, err := io.Copy(dest, r)
+	maxSize := len(source) * 255
+	if maxSize > PayloadMaxSize {
+		maxSize = PayloadMaxSize
+	}
+	dest := make([]byte, maxSize)
+	size, err := lz4.UncompressBlock(source, dest)
 	if err != nil {
 		return nil, err
 	}
-	return dest.Bytes(), nil
+	return dest[:size], nil
 }