Updates to the onion API and node handling, as well as a new `pi_t_fu…

…nctions` package with symmetric AES key generation and encryption.

Description:

This commit updates the onion API and node handling to fix several issues. In the `internal/api/onionApi.go` file, the `QueueOnion` and `Receive` functions have been modified to correctly handle onions and messages. In `node.go`, the `startRun` function has been fixed to correctly handle active nodes and message queuing.

Additionally, the `pi_t` package has been broken out into its own module, `pi_t`, and relocated to the `pkg` directory. This change allows the `pi_t` package to be easily reused in other projects. The package includes a new function, `GenerateSymmetricKey`, which generates a random 256-bit AES key for encryption. The `EncryptWithAES` and `DecryptWithAES` functions have been added to the `pi_t` package to encrypt and decrypt data using the generated AES key.

The commit also updates the `internal/node/node_handler.go` file to correctly handle onions and messages using the new `pi_t` package. The `HandleReceive` and `HandleClientRequest` functions have been modified to use the new `pi_t` functions to encrypt and decrypt onions.

undefined

internal/api/onionApi.go

@@ -0,0 +1,5 @@
+package api
+
+type OnionApi struct {
+	Onion string
+}

internal/node/node.go

@@ -48,7 +48,7 @@ func qoLess(a, b QueuedOnion) bool {
 
 // NewNode creates a new node
 func NewNode(id int, host string, port int, bulletinBoardUrl string) (*Node, error) {
-	if publicKey, privateKey, err := pi_t.KeyGen(); err != nil {
+	if privateKey, publicKey, err := pi_t.KeyGen(); err != nil {
 		return nil, fmt.Errorf("node.NewNode(): failed to generate key pair: %w", err)
 	} else {
 		n := &Node{
@@ -128,12 +128,14 @@ func (n *Node) getRandomNode() *api.PublicNodeApi {
 
 func (n *Node) QueueOnion(msg api.Message, pathLength int) error {
 	timeReceived := time.Now()
+	n.mu.Lock()
+	defer n.mu.Unlock()
 	if msgString, err := json.Marshal(msg); err != nil {
-		return fmt.Errorf("NewOnion(): failed to marshal message: %w", err)
+		return fmt.Errorf("QueueOnion(): failed to marshal message: %w", err)
 	} else if to := n.getNode(msg.To); to == nil {
-		return fmt.Errorf("NewOnion(): failed to get node with id %d", msg.To)
+		return fmt.Errorf("QueueOnion(): failed to get node with id %d", msg.To)
 	} else if routingPath, err2 := n.DetermineRoutingPath(pathLength); err2 != nil {
-		return fmt.Errorf("NewOnion(): failed to determine routing path: %w", err2)
+		return fmt.Errorf("QueueOnion(): failed to determine routing path: %w", err2)
 	} else {
 		publicKeys := utils.Map(routingPath, func(node api.PublicNodeApi) string {
 			return node.PublicKey
@@ -177,17 +179,18 @@ func (n *Node) IDsMatch(nodeApi api.PublicNodeApi) bool {
 }
 
 func (n *Node) startRun(activeNodes []api.PublicNodeApi) (didParticipate bool, e error) {
-	n.wg.Wait()
-	n.wg.Add(1)
-	defer n.wg.Done()
+	//n.wg.Wait()
+	//n.wg.Add(1)
+	//defer n.wg.Done()
 
 	n.mu.Lock()
 	if len(activeNodes) == 0 {
 		n.mu.Unlock()
 		return false, fmt.Errorf("startRun(): no active nodes")
 	}
 	n.ActiveNodes = utils.Copy(activeNodes)
-	onionsToSend := n.OnionQueue.Drain()
+	onionsToSend := n.OnionQueue.Values()
+	n.OnionQueue.Clear()
 	n.mu.Unlock()
 
 	slog.Info("Starting run with", "num_onions", len(onionsToSend))
@@ -209,23 +212,38 @@ func (n *Node) Receive(o string) error {
 	if destination, payload, err := pi_t.PeelOnion(o, n.PrivateKey); err != nil {
 		return fmt.Errorf("node.Receive(): failed to remove layer: %w", err)
 	} else {
-		bruised, err2 := pi_t.BruiseOnion(payload)
-		if err2 != nil {
-			return fmt.Errorf("node.Receive(): failed to bruise onion: %w", err2)
-		}
-		if err3 := sendToNode(QueuedOnion{
-			ConstructedOnion:   bruised,
-			DestinationAddress: destination,
-		}); err != nil {
-			return fmt.Errorf("node.Receive(): failed to send to next node: %w", err3)
+		if destination == "" {
+			var msg api.Message
+			if err2 := json.Unmarshal([]byte(payload), &msg); err2 != nil {
+				return fmt.Errorf("node.Receive(): failed to unmarshal message: %w", err2)
+			}
+			slog.Info("Received message", "from", msg.From, "to", msg.To, "msg", msg.Msg)
+
+		} else {
+			slog.Info("Received onion", "destination", destination)
+			//bruised, err2 := pi_t.BruiseOnion(payload)
+			//if err2 != nil {
+			//	return fmt.Errorf("node.Receive(): failed to bruise onion: %w", err2)
+			//}
+			if err3 := sendToNode(QueuedOnion{
+				ConstructedOnion:   payload,
+				DestinationAddress: destination,
+			}); err != nil {
+				return fmt.Errorf("node.Receive(): failed to send to next node: %w", err3)
+			}
 		}
 	}
 	return nil
 }
 
 func sendToNode(onion QueuedOnion) error {
-	url := fmt.Sprintf("http://%s/receive", onion.DestinationAddress)
-	if resp, err2 := http.Post(url, "application/json", bytes.NewBuffer([]byte(onion.ConstructedOnion))); err2 != nil {
+	url := fmt.Sprintf("%s/receive", onion.DestinationAddress)
+	o := api.OnionApi{
+		Onion: onion.ConstructedOnion,
+	}
+	if data, err := json.Marshal(o); err != nil {
+		slog.Error("failed to marshal msgs", err)
+	} else if resp, err2 := http.Post(url, "application/json", bytes.NewBuffer(data)); err2 != nil {
 		return fmt.Errorf("sendToNode(): failed to send POST request with onion to next node: %w", err2)
 	} else {
 		defer func(Body io.ReadCloser) {
@@ -236,6 +254,6 @@ func sendToNode(onion QueuedOnion) error {
 		if resp.StatusCode != http.StatusOK {
 			return fmt.Errorf("sendToNode(): Failed to send to next node, status code: %d, status: %s", resp.StatusCode, resp.Status)
 		}
-		return nil
 	}
+	return nil
 }

internal/node/node_handler.go

@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/HannahMarsh/pi_t-experiment/internal/api"
-	"github.com/HannahMarsh/pi_t-experiment/pkg/utils"
 	"golang.org/x/exp/slog"
 	"io"
 	"net/http"
@@ -14,13 +13,13 @@ import (
 
 func (n *Node) HandleReceive(w http.ResponseWriter, r *http.Request) {
 	slog.Info("Received onion")
-	var o string
+	var o api.OnionApi
 	if err := json.NewDecoder(r.Body).Decode(&o); err != nil {
 		slog.Error("Error decoding onion", err)
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
-	if err := n.Receive(o); err != nil {
+	if err := n.Receive(o.Onion); err != nil {
 		slog.Error("Error receiving onion", err)
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
@@ -36,7 +35,7 @@ func (n *Node) HandleStartRun(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
-	slog.Info("Active nodes", "activeNodes", activeNodes)
+	//slog.Info("Active nodes", "activeNodes", activeNodes)
 	go func() {
 		if didParticipate, err := n.startRun(activeNodes); err != nil {
 			slog.Error("Error starting run", err)
@@ -54,11 +53,11 @@ func (n *Node) HandleClientRequest(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
-	slog.Info("Received client request", "num_messages", len(msgs), "destinations", utils.Map(msgs, func(m api.Message) int { return m.To }))
+	//slog.Info("Received client request", "num_messages", len(msgs), "destinations", utils.Map(msgs, func(m api.Message) int { return m.To }))
 	//slog.Info("Enqueuing messages", "num_messages", len(msgs))
 	for _, msg := range msgs {
 		if err := n.QueueOnion(msg, 2); err != nil {
-			slog.Error("Error queuing message", err)
+			slog.Error("Error queueing message", err)
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
 		}

internal/pi_t/pi_t_functions.go

@@ -1,6 +1,8 @@
 package pi_t
 
 import (
+	"crypto/aes"
+	"crypto/cipher"
 	"crypto/rand"
 	"crypto/rsa"
 	"crypto/x509"
@@ -9,6 +11,8 @@ import (
 	"encoding/pem"
 	"errors"
 	"fmt"
+	"io"
+	"strings"
 )
 
 // KeyGen generates an RSA key pair and returns the public and private keys in PEM format
@@ -43,21 +47,84 @@ func KeyGen() (privateKeyPEM, publicKeyPEM string, err error) {
 	return privateKeyPEM, publicKeyPEM, nil
 }
 
+// GenerateSymmetricKey generates a random AES key for encryption
+func GenerateSymmetricKey() ([]byte, error) {
+	key := make([]byte, 32) // AES-256
+	if _, err := rand.Read(key); err != nil {
+		return nil, err
+	}
+	return key, nil
+}
+
+// EncryptWithAES encrypts plaintext using AES encryption
+func EncryptWithAES(key, plaintext []byte) (string, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return "", err
+	}
+
+	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
+	iv := ciphertext[:aes.BlockSize]
+	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
+		return "", err
+	}
+
+	stream := cipher.NewCFBEncrypter(block, iv)
+	stream.XORKeyStream(ciphertext[aes.BlockSize:], plaintext)
+
+	return base64.StdEncoding.EncodeToString(ciphertext), nil
+}
+
+// DecryptWithAES decrypts ciphertext using AES encryption
+func DecryptWithAES(key []byte, ct string) ([]byte, error) {
+	ciphertext, _ := base64.StdEncoding.DecodeString(ct)
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(ciphertext) < aes.BlockSize {
+		return nil, errors.New("ciphertext too short")
+	}
+	iv := ciphertext[:aes.BlockSize]
+	ciphertext = ciphertext[aes.BlockSize:]
+
+	stream := cipher.NewCFBDecrypter(block, iv)
+	stream.XORKeyStream(ciphertext, ciphertext)
+
+	return ciphertext, nil
+}
+
 type OnionLayer struct {
 	NextHop string
 	Payload string
 }
 
 // FormOnion creates an onion by encapsulating a message in multiple encryption layers
 func FormOnion(payload []byte, publicKeys []string, routingPath []string) (string, string, error) {
-
 	for i := len(publicKeys) - 1; i >= 0; i-- {
-		layer := OnionLayer{
-			NextHop: routingPath[i],
-			Payload: base64.StdEncoding.EncodeToString(payload),
+		var layerBytes []byte
+		var err error
+		if i == len(publicKeys)-1 {
+			layerBytes = payload
+		} else {
+			layer := OnionLayer{
+				NextHop: routingPath[i+1],
+				Payload: base64.StdEncoding.EncodeToString(payload),
+			}
+
+			layerBytes, err = json.Marshal(layer)
+			if err != nil {
+				return "", "", err
+			}
 		}
 
-		layerBytes, err := json.Marshal(layer)
+		symmetricKey, err := GenerateSymmetricKey()
+		if err != nil {
+			return "", "", err
+		}
+
+		encryptedPayload, err := EncryptWithAES(symmetricKey, layerBytes)
 		if err != nil {
 			return "", "", err
 		}
@@ -72,7 +139,20 @@ func FormOnion(payload []byte, publicKeys []string, routingPath []string) (strin
 			return "", "", err
 		}
 
-		payload, err = rsa.EncryptPKCS1v15(rand.Reader, pubKey.(*rsa.PublicKey), layerBytes)
+		encryptedKey, err := rsa.EncryptPKCS1v15(rand.Reader, pubKey.(*rsa.PublicKey), symmetricKey)
+		if err != nil {
+			return "", "", err
+		}
+
+		combinedPayload := struct {
+			Key     string
+			Payload string
+		}{
+			Key:     base64.StdEncoding.EncodeToString(encryptedKey),
+			Payload: encryptedPayload,
+		}
+
+		payload, err = json.Marshal(combinedPayload)
 		if err != nil {
 			return "", "", err
 		}
@@ -98,15 +178,37 @@ func PeelOnion(onion string, privateKeyPEM string) (string, string, error) {
 		return "", "", err
 	}
 
-	decryptedBytes, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, onionBytes)
+	var combinedPayload struct {
+		Key     string
+		Payload string
+	}
+	if err = json.Unmarshal(onionBytes, &combinedPayload); err != nil {
+		return "", "", err
+	}
+
+	encryptedKey, err := base64.StdEncoding.DecodeString(combinedPayload.Key)
 	if err != nil {
 		return "", "", err
 	}
 
+	symmetricKey, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, encryptedKey)
+	if err != nil {
+		return "", "", err
+	}
+
+	decryptedBytes, err := DecryptWithAES(symmetricKey, combinedPayload.Payload)
+	if err != nil {
+		return "", "", err
+	}
+
+	if !strings.HasPrefix(string(decryptedBytes), "{\"NextHop\":") {
+		return "", string(decryptedBytes), nil
+	}
+
 	var layer OnionLayer
 	err = json.Unmarshal(decryptedBytes, &layer)
 	if err != nil {
-		return "", "", err
+		return "", string(decryptedBytes), nil
 	}
 
 	return layer.NextHop, layer.Payload, nil
@@ -122,6 +224,8 @@ func BruiseOnion(onion string) (string, error) {
 	// Introduce bruising by modifying a small portion of the payload
 	if len(onionBytes) > 0 {
 		onionBytes[0] ^= 0xFF
+	} else {
+		return "", errors.New("empty onion")
 	}
 
 	return base64.StdEncoding.EncodeToString(onionBytes), nil