Add comments and clean up encryption

4 files changed, 52 insertions, 29 deletions

diff --git a/backend/api/routes.go b/backend/api/routes.go
index ff43714..474fdda 100644
--- a/backend/api/routes.go
+++ b/backend/api/routes.go
@@ -3,6 +3,7 @@ package api
 import (
 	"encoding/json"
 	"fmt"
+	"github.com/jackyzha0/ctrl-v/security"
 	"net/http"
 	"time"
 
@@ -87,6 +88,13 @@ func handleGetPaste(w http.ResponseWriter, r *http.Request, parsedPassword strin
 		return
 	}
 
+	// if internal error with encryption
+	if err == security.EncryptionError {
+		w.WriteHeader(http.StatusInternalServerError)
+		fmt.Fprintf(w, "%s", err)
+		return
+	}
+
 	// otherwise, return paste content, title, and current time
 	w.Header().Set("Content-Type", "application/json")
 	pasteMap := map[string]interface{}{
diff --git a/backend/cache/cache.go b/backend/cache/cache.go
index 6d5eb42..43e615a 100644
--- a/backend/cache/cache.go
+++ b/backend/cache/cache.go
@@ -17,7 +17,6 @@ var C *Cache
 
 var PasteNotFound = errors.New("could not find a paste with that hash")
 var UserUnauthorized = errors.New("paste is password protected")
-var EncryptionError = errors.New("could not encrypt the given content")
 
 func init() {
 	C = &Cache{
@@ -52,17 +51,17 @@ func (c *Cache) Get(hash, userPassword string) (db.Paste, error) {
 		}
 
 		// if password matches, decrypt content
-		key, _, err := security.DeriveKey([]byte(userPassword), p.Salt)
+		key, _, err := security.DeriveKey(userPassword, p.Salt)
 		if err != nil {
-			return db.Paste{}, EncryptionError
+			return db.Paste{}, security.EncryptionError
 		}
 
-		decryptedBytes, err := security.Decrypt(key, []byte(p.Content))
+		decryptedContent, err := security.Decrypt(key, p.Content)
 		if err != nil {
-			return db.Paste{}, EncryptionError
+			return db.Paste{}, security.EncryptionError
 		}
 
-		p.Content = string(decryptedBytes)
+		p.Content = decryptedContent
 	}
 
 	return p, nil
diff --git a/backend/db/db.go b/backend/db/db.go
index b18eddf..df112d0 100644
--- a/backend/db/db.go
+++ b/backend/db/db.go
@@ -48,18 +48,18 @@ func New(ip, content, expiry, title, password string) (string, error) {
 	// if there is a password, encrypt content and hash the password
 	if password != "" {
 		// use pass to encrypt content
-		key, salt, err := security.DeriveKey([]byte(password), nil)
+		key, salt, err := security.DeriveKey(password, nil)
 		if err != nil {
 			return "", fmt.Errorf("could not generate key: %s", err.Error())
 		}
 		new.Salt = salt
 
-		encryptedBytes, err := security.Encrypt(key, []byte(new.Content))
+		encryptedContent, err := security.Encrypt(key, new.Content)
 		if err != nil {
 			return "", fmt.Errorf("could not encrypt content: %s", err.Error())
 		}
 
-		new.Content = string(encryptedBytes)
+		new.Content = encryptedContent
 
 		// hash given password
 		hashedPass, err := security.HashPassword(password)
diff --git a/backend/security/encrypt.go b/backend/security/encrypt.go
index fff027c..4af5e3c 100644
--- a/backend/security/encrypt.go
+++ b/backend/security/encrypt.go
@@ -4,62 +4,78 @@ import (
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/rand"
+	"errors"
 	"golang.org/x/crypto/scrypt"
 )
 
-func Encrypt(key, data []byte) ([]byte, error) {
-	blockCipher, err := aes.NewCipher(key)
+var EncryptionError = errors.New("could not encrypt the given content")
+
+func Encrypt(key, data string) (string, error) {
+	// initialize aes block cipher with given key
+	blockCipher, err := aes.NewCipher([]byte(key))
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
+	// wrap block cipher with Galois Counter Mode and standard nonce length
 	gcm, err := cipher.NewGCM(blockCipher)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
+	// generate nonce (number once used) unique to the given key
 	nonce := make([]byte, gcm.NonceSize())
 	if _, err = rand.Read(nonce); err != nil {
-		return nil, err
+		return "", err
 	}
 
-	cipherText := gcm.Seal(nonce, nonce, data, nil)
+	// seal nonce with data to use during decryption
+	cipherText := gcm.Seal(nonce, nonce, []byte(data), nil)
 
-	return cipherText, nil
+	return string(cipherText), nil
 }
 
-func Decrypt(key, data []byte) ([]byte, error) {
-	blockCipher, err := aes.NewCipher(key)
+func Decrypt(key, data string) (string, error) {
+	// similar to encrypt, create cipher and wrap with GCM
+	blockCipher, err := aes.NewCipher([]byte(key))
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
 	gcm, err := cipher.NewGCM(blockCipher)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
+	// extract the nonce from the data
 	nonce, cipherText := data[:gcm.NonceSize()], data[gcm.NonceSize():]
-	plaintext, err := gcm.Open(nil, nonce, cipherText, nil)
+
+	// use nonce to decrypt the data
+	plaintext, err := gcm.Open(nil, []byte(nonce), []byte(cipherText), nil)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 
-	return plaintext, nil
+	return string(plaintext), nil
 }
 
-func DeriveKey(password, salt []byte) ([]byte, []byte, error) {
+const keyBytes = 16
+const iterations = 16384
+const relativeMemoryCost = 8
+const relativeCPUCost = 1
+
+func DeriveKey(password string, salt []byte) (string, []byte, error) {
 	if salt == nil {
-		salt = make([]byte, 16)
+		salt = make([]byte, keyBytes)
 		if _, err := rand.Read(salt); err != nil {
-			return nil, nil, err
+			return "", nil, err
 		}
 	}
 
-	key, err := scrypt.Key(password, salt, 16384, 8, 1, 16)
+	key, err := scrypt.Key([]byte(password), salt, iterations, relativeMemoryCost, relativeCPUCost, keyBytes)
 	if err != nil {
-		return nil, nil, err
+		return "", nil, err
 	}
 
-	return key, salt, nil
+	return string(key), salt, nil
 }