From 054d8a5a3e9395066feb91eaa547f9b09a0a98f1 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Fri, 29 Mar 2019 21:45:43 +0100
Subject: [PATCH] misc: keeping frames in memory as compressed

---
 modules/api_rest/record.go | 74 ++++++++++++++++++++++++++------------
 1 file changed, 52 insertions(+), 22 deletions(-)

diff --git a/modules/api_rest/record.go b/modules/api_rest/record.go
index 4693ce2a..6e061cd6 100644
--- a/modules/api_rest/record.go
+++ b/modules/api_rest/record.go
@@ -15,6 +15,36 @@ import (
 	"github.com/kr/binarydist"
 )
 
+func compress(data []byte) (error, []byte) {
+	compressed := new(bytes.Buffer)
+	compress := gzip.NewWriter(compressed)
+
+	if _, err := compress.Write(data); err != nil {
+		return err, nil
+	} else if err = compress.Flush(); err != nil {
+		return err, nil
+	} else if err = compress.Close(); err != nil {
+		return err, nil
+	}
+
+	return nil, compressed.Bytes()
+}
+
+func decompress(data []byte) (error, []byte) {
+	decompress, err := gzip.NewReader(bytes.NewReader(data))
+	if err != nil {
+		return err, nil
+	}
+	defer decompress.Close()
+
+	raw, err := ioutil.ReadAll(decompress)
+	if err != nil {
+		return err, nil
+	}
+
+	return nil, raw
+}
+
 type patch []byte
 type frame []byte
 
@@ -79,7 +109,7 @@ func (e *RecordEntry) Reset() {
 	e.CurState = 0
 }
 
-func (e *RecordEntry) Compile() error {
+func (e *RecordEntry) Compile() (err error) {
 	e.Lock()
 	defer e.Unlock()
 
@@ -90,7 +120,10 @@ func (e *RecordEntry) Compile() error {
 	e.frames = make([]frame, e.NumStates+1)
 
 	// first is the master frame
-	e.frames[0] = frame(e.Data)
+	if err, e.frames[0] = compress(e.Data); err != nil {
+		return
+	}
+
 	// precompute frames so they can be accessed by index
 	for i := 0; i < e.NumStates; i++ {
 		patch := e.States[i]
@@ -103,7 +136,10 @@ func (e *RecordEntry) Compile() error {
 		}
 
 		e.Cur = newWriter.Bytes()
-		e.frames[i+1] = e.Cur
+
+		if err, e.frames[i+1] = compress(e.Cur); err != nil {
+			return err
+		}
 
 		e.progress(1)
 	}
@@ -143,7 +179,15 @@ func (e *RecordEntry) Next() []byte {
 	defer e.Unlock()
 	cur := e.CurState
 	e.CurState++
-	return e.frames[cur]
+
+	zFrame := e.frames[cur]
+
+	err, data := decompress(zFrame)
+	if err != nil {
+		panic(err)
+	}
+
+	return data
 }
 
 // the Record object represents a recorded session
@@ -187,16 +231,10 @@ func LoadRecord(fileName string, mod *session.SessionModule) (*Record, error) {
 		return nil, fmt.Errorf("error while reading %s: %s", fileName, err)
 	}
 
-	decompress, err := gzip.NewReader(bytes.NewReader(compressed))
+	err, raw := decompress(compressed)
 	if err != nil {
 		return nil, fmt.Errorf("error while reading gzip file %s: %s", fileName, err)
 	}
-	defer decompress.Close()
-
-	raw, err := ioutil.ReadAll(decompress)
-	if err != nil {
-		return nil, fmt.Errorf("error while decompressing %s: %s", fileName, err)
-	}
 
 	rec := &Record{}
 
@@ -244,20 +282,12 @@ func (r *Record) save() error {
 		return err
 	}
 
-	data := buf.Bytes()
-
-	compressed := new(bytes.Buffer)
-	compress := gzip.NewWriter(compressed)
-
-	if _, err := compress.Write(data); err != nil {
-		return err
-	} else if err = compress.Flush(); err != nil {
-		return err
-	} else if err = compress.Close(); err != nil {
+	err, data := compress(buf.Bytes())
+	if err != nil {
 		return err
 	}
 
-	return ioutil.WriteFile(r.fileName, compressed.Bytes(), os.ModePerm)
+	return ioutil.WriteFile(r.fileName, data, os.ModePerm)
 }
 
 func (r *Record) Flush() error {