From 7be2b0e09de435a7df486e3b96074df9ab90cd0b Mon Sep 17 00:00:00 2001
From: 0xdcarns <dillon.carns@gmail.com>
Date: Fri, 28 Jan 2022 17:49:31 -0500
Subject: [PATCH] adjusted to byte buffer

---
 controllers/node_grpc.go |  2 +-
 database/database.go     | 12 +++++++-----
 logic/traffic.go         | 16 +++++++++++-----
 models/structs.go        |  7 ++++---
 4 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/controllers/node_grpc.go b/controllers/node_grpc.go
index 9221041f..d9dddc99 100644
--- a/controllers/node_grpc.go
+++ b/controllers/node_grpc.go
@@ -87,7 +87,7 @@ func (s *NodeServiceServer) CreateNode(ctx context.Context, req *nodepb.Object)
 
 	node.TrafficKeys = models.TrafficKeys{
 		Mine:   node.TrafficKeys.Mine,
-		Server: key.PublicKey,
+		Server: key,
 	}
 
 	fmt.Printf("finished created node: %v \n", node)
diff --git a/database/database.go b/database/database.go
index 8d054172..cec1a9cd 100644
--- a/database/database.go
+++ b/database/database.go
@@ -1,11 +1,12 @@
 package database
 
 import (
+	"bytes"
 	"crypto/rand"
 	"crypto/rsa"
+	"encoding/gob"
 	"encoding/json"
 	"errors"
-	"fmt"
 	"strings"
 	"time"
 
@@ -216,11 +217,12 @@ func initializeUUID() error {
 	if keyErr != nil {
 		return keyErr
 	}
+	var rsaKey bytes.Buffer
+	if err = gob.NewEncoder(&rsaKey).Encode(rsaPrivKey); err != nil {
+		return err
+	}
 
-	fmt.Printf("key generated: %v \n", rsaPrivKey)
-	fmt.Printf("pub key generate: %v \n", rsaPrivKey.PublicKey)
-
-	telemetry := models.Telemetry{UUID: uuid.NewString(), TrafficKey: *rsaPrivKey}
+	telemetry := models.Telemetry{UUID: uuid.NewString(), TrafficKey: rsaKey}
 	telJSON, err := json.Marshal(&telemetry)
 	if err != nil {
 		return err
diff --git a/logic/traffic.go b/logic/traffic.go
index 088a1761..ac187bf3 100644
--- a/logic/traffic.go
+++ b/logic/traffic.go
@@ -2,15 +2,21 @@ package logic
 
 import (
 	"crypto/rsa"
+	"encoding/gob"
 	"fmt"
 )
 
-// RetrieveTrafficKey - retrieves key based on node
-func RetrieveTrafficKey() (rsa.PrivateKey, error) {
+// RetrieveTrafficKey - retrieves public key based on node
+func RetrieveTrafficKey() (rsa.PublicKey, error) {
 	var telRecord, err = fetchTelemetryRecord()
 	if err != nil {
-		return rsa.PrivateKey{}, err
+		return rsa.PublicKey{}, err
 	}
-	fmt.Printf("retrieved key: %v \n", telRecord.TrafficKey)
-	return telRecord.TrafficKey, nil
+	var key = rsa.PrivateKey{}
+	if err = gob.NewDecoder(&telRecord.TrafficKey).Decode(&key); err != nil {
+		return rsa.PublicKey{}, err
+	}
+	fmt.Printf("retrieved key: %v \n", key.PublicKey)
+
+	return key.PublicKey, nil
 }
diff --git a/models/structs.go b/models/structs.go
index b8a18d15..0ee55a7b 100644
--- a/models/structs.go
+++ b/models/structs.go
@@ -1,6 +1,7 @@
 package models
 
 import (
+	"bytes"
 	"crypto/rsa"
 
 	jwt "github.com/golang-jwt/jwt/v4"
@@ -170,9 +171,9 @@ type ServerUpdateData struct {
 
 // Telemetry - contains UUID of the server and timestamp of last send to posthog
 type Telemetry struct {
-	UUID       string         `json:"uuid" bson:"uuid"`
-	LastSend   int64          `json:"lastsend" bson:"lastsend"`
-	TrafficKey rsa.PrivateKey `json:"traffickey" bson:"traffickey"`
+	UUID       string       `json:"uuid" bson:"uuid"`
+	LastSend   int64        `json:"lastsend" bson:"lastsend"`
+	TrafficKey bytes.Buffer `json:"traffickey" bson:"traffickey"`
 }
 
 // ServerAddr - to pass to clients to tell server addresses and if it's the leader or not