From f82c267d8141013f2af6910bc960278540d883bd Mon Sep 17 00:00:00 2001
From: Milos Gajdos <milosgajdos83@gmail.com>
Date: Mon, 25 Nov 2019 15:34:41 +0000
Subject: [PATCH] Encrypt session communication

---
 tunnel/crypto.go  |  6 +++---
 tunnel/session.go | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/tunnel/crypto.go b/tunnel/crypto.go
index 89a0bdbf..f34f9c15 100644
--- a/tunnel/crypto.go
+++ b/tunnel/crypto.go
@@ -8,7 +8,7 @@ import (
 	"io"
 )
 
-// Encrypt encrypts data and returns encrypted payload
+// Encrypt encrypts data and returns the encrypted data
 func Encrypt(data []byte, key string) ([]byte, error) {
 	// generate a new AES cipher using our 32 byte key
 	c, err := aes.NewCipher(hash(key))
@@ -37,7 +37,7 @@ func Encrypt(data []byte, key string) ([]byte, error) {
 	return gcm.Seal(nonce, nonce, data, nil), nil
 }
 
-// Decrypt decrypts the payload and returns decrypted data
+// Decrypt decrypts the payload and returns the decrypted data
 func Decrypt(data []byte, key string) ([]byte, error) {
 	// generate AES cipher for decrypting the message
 	c, err := aes.NewCipher(hash(key))
@@ -64,7 +64,7 @@ func Decrypt(data []byte, key string) ([]byte, error) {
 }
 
 // hash hahes the data into 32 bytes key and returns it
-// hash uses sha256 to hash the passed in string.
+// hash uses sha256 underneath to hash the supplied key
 func hash(key string) []byte {
 	hasher := sha256.New()
 	hasher.Write([]byte(key))
diff --git a/tunnel/session.go b/tunnel/session.go
index 59cbdb5e..09042fd1 100644
--- a/tunnel/session.go
+++ b/tunnel/session.go
@@ -301,13 +301,27 @@ func (s *session) Send(m *transport.Message) error {
 		// no op
 	}
 
+	// get the token
+	token, ok := m.Header["Micro-Tunnel-Token"]
+	if !ok {
+		// TODO: should we continue or return error
+		log.Debugf("no token found, insecure channel")
+	}
+
+	// encrypt the transport message payload
+	body, err := Encrypt(m.Body, token+s.channel+s.session)
+	if err != nil {
+		return err
+	}
+
 	// make copy
 	data := &transport.Message{
 		Header: make(map[string]string),
-		Body:   m.Body,
+		Body:   body,
 	}
 
 	for k, v := range m.Header {
+		// TODO: should we also encrypt headers?
 		data.Header[k] = v
 	}
 
@@ -352,7 +366,22 @@ func (s *session) Recv(m *transport.Message) error {
 	default:
 	}
 
+	// TODO: if we encrypt headers we will have to decrypt them here
+	token, ok := msg.data.Header["Micro-Tunnel-Token"]
+	if !ok {
+		// TODO: should we continue or return error
+		log.Debugf("no token found, insecure channel")
+	}
+
 	log.Tracef("Received %+v from recv backlog", msg)
+
+	// decrypt the received payload using the token
+	body, err := Decrypt(msg.data.Body, token+s.channel+s.session)
+	if err != nil {
+		return err
+	}
+	msg.data.Body = body
+
 	// set message
 	*m = *msg.data
 	// return nil