From eb1b5a076c11dc15e773183afe3a9a375fbe9cd5 Mon Sep 17 00:00:00 2001
From: Zack Scholl <Zack Scholl>
Date: Tue, 17 Oct 2017 21:40:20 -0600
Subject: [PATCH] Store file data in memory when transfering

---
 connect.go | 38 +++++++++++++++++++-------------------
 main.go    |  9 ++-------
 2 files changed, 21 insertions(+), 26 deletions(-)

diff --git a/connect.go b/connect.go
index 7af6190..8ad556a 100644
--- a/connect.go
+++ b/connect.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -90,9 +91,6 @@ func runClient(connectionType string, codePhrase string) {
 		} else {
 			fmt.Printf("\nDownloaded %s!", fileName)
 		}
-	} else {
-		log.Info("cleaning up")
-		os.Remove(fileName + ".encrypted")
 	}
 }
 
@@ -189,26 +187,28 @@ func sendFile(id int, connection net.Conn, codePhrase string) {
 	})
 	defer connection.Close()
 
-	// Open the file that needs to be send to the client
-	file, err := os.Open(fileName + ".encrypted")
-	if err != nil {
-		logger.Error(err)
-		return
-	}
-	defer file.Close()
-	// Get the filename and filesize
-	fileInfo, err := file.Stat()
-	if err != nil {
-		logger.Error(err)
-		return
-	}
+	var err error
 
-	numChunks := math.Ceil(float64(fileInfo.Size()) / float64(BUFFERSIZE))
+	// // Open the file that needs to be send to the client
+	// file, err := os.Open(fileName + ".encrypted")
+	// if err != nil {
+	// 	logger.Error(err)
+	// 	return
+	// }
+	// defer file.Close()
+	// // Get the filename and filesize
+	// fileInfo, err := file.Stat()
+	// if err != nil {
+	// 	logger.Error(err)
+	// 	return
+	// }
+
+	numChunks := math.Ceil(float64(len(fileBytes)) / float64(BUFFERSIZE))
 	chunksPerWorker := int(math.Ceil(numChunks / float64(numberConnections)))
 
 	bytesPerConnection := int64(chunksPerWorker * BUFFERSIZE)
 	if id+1 == numberConnections {
-		bytesPerConnection = fileInfo.Size() - (numberConnections-1)*bytesPerConnection
+		bytesPerConnection = int64(len(fileBytes)) - (numberConnections-1)*bytesPerConnection
 	}
 	fileSize := fillString(strconv.FormatInt(int64(bytesPerConnection), 10), 10)
 
@@ -242,7 +242,7 @@ func sendFile(id int, connection net.Conn, codePhrase string) {
 	connection.Write([]byte(fillString(fileHash, BUFFERSIZE)))
 
 	sendBuffer := make([]byte, BUFFERSIZE)
-
+	file := bytes.NewBuffer(fileBytes)
 	chunkI := 0
 	for {
 		_, err = file.Read(sendBuffer)
diff --git a/main.go b/main.go
index d5b24d9..0d552ac 100644
--- a/main.go
+++ b/main.go
@@ -21,6 +21,7 @@ var server, file string
 var serverAddress, fileName, codePhraseFlag, connectionTypeFlag string
 var runAsRelay, debugFlag bool
 var fileSalt, fileIV, fileHash string
+var fileBytes []byte
 
 func main() {
 	flag.BoolVar(&runAsRelay, "relay", false, "run as relay")
@@ -63,13 +64,7 @@ func main() {
 			log.Fatal(err)
 			return
 		}
-		var encrypted []byte
-		encrypted, fileSalt, fileIV = Encrypt(fdata, codePhraseFlag)
-		if err != nil {
-			log.Fatal(err)
-			return
-		}
-		ioutil.WriteFile(fileName+".encrypted", encrypted, 0644)
+		fileBytes, fileSalt, fileIV = Encrypt(fdata, codePhraseFlag)
 		fileHash = HashBytes(fdata)
 	}