Session mod files encryption (#766)

* feat(backend): added session mod files encryption

backend/cmd/ender/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"log"
+	"openreplay/backend/internal/storage"
 	"os"
 	"os/signal"
 	"syscall"
@@ -78,6 +79,15 @@ func main() {
 						currDuration, newDuration)
 					return true
 				}
+				if cfg.UseEncryption {
+					if key := storage.GenerateEncryptionKey(); key != nil {
+						if err := pg.InsertSessionEncryptionKey(sessionID, key); err != nil {
+							log.Printf("can't save session encryption key: %s, session will not be encrypted", err)
+						} else {
+							msg.EncryptionKey = string(key)
+						}
+					}
+				}
 				if err := producer.Produce(cfg.TopicRawWeb, sessionID, msg.Encode()); err != nil {
 					log.Printf("can't send sessionEnd to topic: %s; sessID: %d", err, sessionID)
 					return false

backend/cmd/storage/main.go

@@ -44,7 +44,7 @@ func main() {
 		messages.NewMessageIterator(
 			func(msg messages.Message) {
 				sesEnd := msg.(*messages.SessionEnd)
-				if err := srv.UploadSessionFiles(msg.SessionID()); err != nil {
+				if err := srv.UploadSessionFiles(sesEnd); err != nil {
 					log.Printf("can't find session: %d", msg.SessionID())
 					sessionFinder.Find(msg.SessionID(), sesEnd.Timestamp)
 				}

backend/internal/config/ender/config.go

@@ -14,6 +14,7 @@ type Config struct {
 	TopicRawWeb                string `env:"TOPIC_RAW_WEB,required"`
 	ProducerTimeout            int    `env:"PRODUCER_TIMEOUT,default=2000"`
 	PartitionsNumber           int    `env:"PARTITIONS_NUMBER,required"`
+	UseEncryption              bool   `env:"USE_ENCRYPTION,default=false"`
 }
 
 func New() *Config {

backend/internal/storage/encryptor.go

@@ -0,0 +1,17 @@
+package storage
+
+import (
+	"errors"
+)
+
+func GenerateEncryptionKey() []byte {
+	return nil
+}
+
+func EncryptData(data, fullKey []byte) ([]byte, error) {
+	return nil, errors.New("not supported")
+}
+
+func DecryptData(data, fullKey []byte) ([]byte, error) {
+	return nil, errors.New("not supported")
+}

backend/internal/storage/storage.go

@@ -8,6 +8,7 @@ import (
 	"log"
 	config "openreplay/backend/internal/config/storage"
 	"openreplay/backend/pkg/flakeid"
+	"openreplay/backend/pkg/messages"
 	"openreplay/backend/pkg/monitoring"
 	"openreplay/backend/pkg/storage"
 	"os"
@@ -68,19 +69,19 @@ func New(cfg *config.Config, s3 *storage.S3, metrics *monitoring.Metrics) (*Stor
 	}, nil
 }
 
-func (s *Storage) UploadSessionFiles(sessID uint64) error {
+func (s *Storage) UploadSessionFiles(msg *messages.SessionEnd) error {
-	sessionDir := strconv.FormatUint(sessID, 10)
+	sessionDir := strconv.FormatUint(msg.SessionID(), 10)
-	if err := s.uploadKey(sessID, sessionDir+"/dom.mob", true, 5); err != nil {
+	if err := s.uploadKey(msg.SessionID(), sessionDir+"/dom.mob", true, 5, msg.EncryptionKey); err != nil {
 		return err
 	}
-	if err := s.uploadKey(sessID, sessionDir+"/devtools.mob", false, 4); err != nil {
+	if err := s.uploadKey(msg.SessionID(), sessionDir+"/devtools.mob", false, 4, msg.EncryptionKey); err != nil {
 		return err
 	}
 	return nil
 }
 
 // TODO: make a bit cleaner
-func (s *Storage) uploadKey(sessID uint64, key string, shouldSplit bool, retryCount int) error {
+func (s *Storage) uploadKey(sessID uint64, key string, shouldSplit bool, retryCount int, encryptionKey string) error {
 	if retryCount <= 0 {
 		return nil
 	}
@@ -95,6 +96,14 @@ func (s *Storage) uploadKey(sessID uint64, key string, shouldSplit bool, retryCo
 	}
 	defer file.Close()
 
+	var fileSize int64 = 0
+	fileInfo, err := file.Stat()
+	if err != nil {
+		log.Printf("can't get file info: %s", err)
+	} else {
+		fileSize = fileInfo.Size()
+	}
+	var encryptedData []byte
 	if shouldSplit {
 		nRead, err := file.Read(s.startBytes)
 		if err != nil {
@@ -105,45 +114,104 @@ func (s *Storage) uploadKey(sessID uint64, key string, shouldSplit bool, retryCo
 				time.UnixMilli(int64(flakeid.ExtractTimestamp(sessID))),
 			)
 			time.AfterFunc(s.cfg.RetryTimeout, func() {
-				s.uploadKey(sessID, key, shouldSplit, retryCount-1)
+				s.uploadKey(sessID, key, shouldSplit, retryCount-1, encryptionKey)
 			})
 			return nil
 		}
 		s.readingTime.Record(context.Background(), float64(time.Now().Sub(start).Milliseconds()))
 
 		start = time.Now()
-		startReader := bytes.NewBuffer(s.startBytes[:nRead])
+		// Encrypt session file if we have encryption key
+		if encryptionKey != "" {
+			encryptedData, err = EncryptData(s.startBytes[:nRead], []byte(encryptionKey))
+			if err != nil {
+				log.Printf("can't encrypt data: %s", err)
+				encryptedData = s.startBytes[:nRead]
+			}
+		} else {
+			encryptedData = s.startBytes[:nRead]
+		}
+		// Compress and save to s3
+		startReader := bytes.NewBuffer(encryptedData)
 		if err := s.s3.Upload(s.gzipFile(startReader), key+"s", "application/octet-stream", true); err != nil {
 			log.Fatalf("Storage: start upload failed.  %v\n", err)
 		}
+		// TODO: fix possible error (if we read less then FileSplitSize)
 		if nRead == s.cfg.FileSplitSize {
-			if err := s.s3.Upload(s.gzipFile(file), key+"e", "application/octet-stream", true); err != nil {
+			restPartSize := fileSize - int64(nRead)
+			fileData := make([]byte, restPartSize)
+			nRead, err = file.Read(fileData)
+			if err != nil {
+				log.Printf("File read error: %s; sessID: %s, part: %d, sessStart: %s",
+					err,
+					key,
+					sessID%16,
+					time.UnixMilli(int64(flakeid.ExtractTimestamp(sessID))),
+				)
+				return nil
+			}
+			if int64(nRead) != restPartSize {
+				log.Printf("can't read the rest part of file")
+			}
+
+			// Encrypt session file if we have encryption key
+			if encryptionKey != "" {
+				encryptedData, err = EncryptData(fileData, []byte(encryptionKey))
+				if err != nil {
+					log.Printf("can't encrypt data: %s", err)
+					encryptedData = fileData
+				}
+			} else {
+				encryptedData = fileData
+			}
+			// Compress and save to s3
+			endReader := bytes.NewBuffer(encryptedData)
+			if err := s.s3.Upload(s.gzipFile(endReader), key+"e", "application/octet-stream", true); err != nil {
 				log.Fatalf("Storage: end upload failed. %v\n", err)
 			}
 		}
 		s.archivingTime.Record(context.Background(), float64(time.Now().Sub(start).Milliseconds()))
 	} else {
 		start = time.Now()
-		if err := s.s3.Upload(s.gzipFile(file), key+"s", "application/octet-stream", true); err != nil {
+		fileData := make([]byte, fileSize)
+		nRead, err := file.Read(fileData)
+		if err != nil {
+			log.Printf("File read error: %s; sessID: %s, part: %d, sessStart: %s",
+				err,
+				key,
+				sessID%16,
+				time.UnixMilli(int64(flakeid.ExtractTimestamp(sessID))),
+			)
+			return nil
+		}
+		if int64(nRead) != fileSize {
+			log.Printf("can't read the rest part of file")
+		}
+
+		// Encrypt session file if we have encryption key
+		if encryptionKey != "" {
+			encryptedData, err = EncryptData(fileData, []byte(encryptionKey))
+			if err != nil {
+				log.Printf("can't encrypt data: %s", err)
+				encryptedData = fileData
+			}
+		} else {
+			encryptedData = fileData
+		}
+		endReader := bytes.NewBuffer(encryptedData)
+		if err := s.s3.Upload(s.gzipFile(endReader), key+"s", "application/octet-stream", true); err != nil {
 			log.Fatalf("Storage: end upload failed. %v\n", err)
 		}
 		s.archivingTime.Record(context.Background(), float64(time.Now().Sub(start).Milliseconds()))
 	}
 
 	// Save metrics
-	var fileSize float64 = 0
-	fileInfo, err := file.Stat()
-	if err != nil {
-		log.Printf("can't get file info: %s", err)
-	} else {
-		fileSize = float64(fileInfo.Size())
-	}
 	ctx, _ := context.WithTimeout(context.Background(), time.Millisecond*200)
 	if shouldSplit {
 		s.totalSessions.Add(ctx, 1)
-		s.sessionDOMSize.Record(ctx, fileSize)
+		s.sessionDOMSize.Record(ctx, float64(fileSize))
 	} else {
-		s.sessionDevtoolsSize.Record(ctx, fileSize)
+		s.sessionDevtoolsSize.Record(ctx, float64(fileSize))
 	}
 
 	return nil

backend/pkg/db/cache/messages-common.go

@@ -11,6 +11,10 @@ func (c *PGCache) InsertSessionEnd(sessionID uint64, timestamp uint64) (uint64,
 	return c.Conn.InsertSessionEnd(sessionID, timestamp)
 }
 
+func (c *PGCache) InsertSessionEncryptionKey(sessionID uint64, key []byte) error {
+	return c.Conn.InsertSessionEncryptionKey(sessionID, key)
+}
+
 func (c *PGCache) HandleSessionEnd(sessionID uint64) error {
 	if err := c.Conn.HandleSessionEnd(sessionID); err != nil {
 		log.Printf("can't handle session end: %s", err)

backend/pkg/db/postgres/messages-common.go

@@ -82,6 +82,10 @@ func (conn *Conn) InsertSessionEnd(sessionID uint64, timestamp uint64) (uint64,
 	return dur, nil
 }
 
+func (conn *Conn) InsertSessionEncryptionKey(sessionID uint64, key []byte) error {
+	return conn.c.Exec(`UPDATE sessions SET file_key = $2 WHERE session_id = $1`, sessionID, string(key))
+}
+
 func (conn *Conn) HandleSessionEnd(sessionID uint64) error {
 	sqlRequest := `
 	UPDATE sessions

backend/pkg/messages/filters.go

@@ -11,4 +11,4 @@ func IsIOSType(id int) bool {
 
 func IsDOMType(id int) bool {
 	return 0 == id || 4 == id || 5 == id || 6 == id || 7 == id || 8 == id || 9 == id || 10 == id || 11 == id || 12 == id || 13 == id || 14 == id || 15 == id || 16 == id || 18 == id || 19 == id || 20 == id || 37 == id || 38 == id || 49 == id || 54 == id || 55 == id || 59 == id || 60 == id || 61 == id || 67 == id || 69 == id || 70 == id || 71 == id || 72 == id || 73 == id || 74 == id || 75 == id || 76 == id || 77 == id || 90 == id || 93 == id || 96 == id || 100 == id || 102 == id || 103 == id || 105 == id
-}
+}

ee/backend/internal/storage/encryptor.go

@@ -0,0 +1,65 @@
+package storage
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"errors"
+	"fmt"
+	"math/rand"
+)
+
+const letterSet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+
+func GenerateEncryptionKey() []byte {
+	return append(generateRandomBytes(16), generateRandomBytes(16)...)
+}
+
+func generateRandomBytes(size int) []byte {
+	b := make([]byte, size)
+	for i := range b {
+		b[i] = letterSet[rand.Int63()%int64(len(letterSet))]
+	}
+	return b
+}
+
+func fillLastBlock(rawText []byte, blockSize int) []byte {
+	padding := blockSize - len(rawText)%blockSize
+	padText := bytes.Repeat([]byte{byte(padding)}, padding)
+	return append(rawText, padText...)
+}
+
+func EncryptData(data, fullKey []byte) ([]byte, error) {
+	if len(fullKey) != 32 {
+		return nil, errors.New("wrong format of encryption key")
+	}
+	key, iv := fullKey[:16], fullKey[16:]
+	// Fill the last block of data by zeros
+	paddedData := fillLastBlock(data, aes.BlockSize)
+	// Create new AES cipher with CBC encryptor
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, fmt.Errorf("cbc encryptor failed: %s", err)
+	}
+	mode := cipher.NewCBCEncrypter(block, iv)
+	// Encrypting data
+	ciphertext := make([]byte, len(paddedData))
+	mode.CryptBlocks(ciphertext, paddedData)
+	// Return encrypted data
+	return ciphertext, nil
+}
+
+func DecryptData(data, fullKey []byte) ([]byte, error) {
+	if len(fullKey) != 32 {
+		return nil, errors.New("wrong format of encryption key")
+	}
+	key, iv := fullKey[:16], fullKey[16:]
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, fmt.Errorf("cbc encryptor failed: %s", err)
+	}
+	cbc := cipher.NewCBCDecrypter(block, iv)
+	res := make([]byte, len(data))
+	cbc.CryptBlocks(res, data)
+	return res, nil
+}

ee/backend/pkg/failover/failover.go

@@ -89,7 +89,9 @@ func (s *sessionFinderImpl) worker() {
 }
 
 func (s *sessionFinderImpl) findSession(sessionID, timestamp, partition uint64) {
-	err := s.storage.UploadSessionFiles(sessionID)
+	sessEnd := &messages.SessionEnd{Timestamp: timestamp}
+	sessEnd.SetSessionID(sessionID)
+	err := s.storage.UploadSessionFiles(sessEnd)
 	if err == nil {
 		log.Printf("found session: %d in partition: %d, original: %d",
 			sessionID, partition, sessionID%numberOfPartitions)

ee/connectors/msgcodec/messages.py

@@ -66,8 +66,9 @@ class SessionStart(Message):
 class SessionEnd(Message):
     __id__ = 3
 
-    def __init__(self, timestamp):
+    def __init__(self, timestamp, encryption_key):
         self.timestamp = timestamp
+        self.encryption_key = encryption_key
 
 
 class SetPageLocation(Message):

ee/connectors/msgcodec/msgcodec.py

@@ -125,7 +125,8 @@ class MessageCodec(Codec):
 
         if message_id == 3:
             return SessionEnd(
-                timestamp=self.read_uint(reader)
+                timestamp=self.read_uint(reader),
+                encryption_key=self.read_string(reader)
             )
 
         if message_id == 4:

mobs/messages.rb

@@ -48,6 +48,7 @@ end
 # end
 message 3, 'SessionEnd', :tracker => false, :replayer => false do
   uint 'Timestamp'
+  string 'EncryptionKey'
 end
 message 4, 'SetPageLocation' do
   string 'URL'