From 96628c64225303ad5b442a0fe7feda66a346b389 Mon Sep 17 00:00:00 2001
From: n33pm <12273891+n33pm@users.noreply.github.com>
Date: Mon, 8 Apr 2024 16:55:22 +0200
Subject: [PATCH] feat(): add connection options - port/key/network

---
 pkg/peerprotocol/connection.go        | 53 +++++++++++++++++++--------
 pkg/peerprotocol/connectionoptions.go | 25 +++++++++++++
 pkg/protocols/fullnode.go             |  1 -
 3 files changed, 63 insertions(+), 16 deletions(-)

diff --git a/pkg/peerprotocol/connection.go b/pkg/peerprotocol/connection.go
index 1eea00c..32a351a 100644
--- a/pkg/peerprotocol/connection.go
+++ b/pkg/peerprotocol/connection.go
@@ -19,6 +19,7 @@ import (
 type Connection struct {
 	chiaConfig *config.ChiaConfig
 
+	networkID   string
 	peerIP      *net.IP
 	peerPort    uint16
 	peerKeyPair *tls.Certificate
@@ -33,15 +34,8 @@ type PeerResponseHandlerFunc func(*protocols.Message, error)
 
 // NewConnection creates a new connection object with the specified peer
 func NewConnection(ip *net.IP, options ...ConnectionOptionFunc) (*Connection, error) {
-	cfg, err := config.GetChiaConfig()
-	if err != nil {
-		return nil, err
-	}
-
 	c := &Connection{
-		chiaConfig: cfg,
-		peerIP:     ip,
-		peerPort:   cfg.FullNode.Port,
+		peerIP: ip,
 	}
 
 	for _, fn := range options {
@@ -53,21 +47,50 @@ func NewConnection(ip *net.IP, options ...ConnectionOptionFunc) (*Connection, er
 		}
 	}
 
-	err = c.loadKeyPair()
-	if err != nil {
-		return nil, err
+	if c.peerPort == 0 {
+		if err := c.loadChiaConfig(); err != nil {
+			return nil, err
+		}
+		c.peerPort = c.chiaConfig.FullNode.Port
+	}
+
+	if c.peerKeyPair == nil {
+		if err := c.loadChiaConfig(); err != nil {
+			return nil, err
+		}
+		if err := c.loadConfigKeyPair(); err != nil {
+			return nil, err
+		}
+	}
+
+	if len(c.networkID) == 0 {
+		if err := c.loadChiaConfig(); err != nil {
+			return nil, err
+		}
+		c.networkID = c.chiaConfig.SelectedNetwork
 	}
 
 	// Generate the websocket dialer
-	err = c.generateDialer()
-	if err != nil {
+	if err := c.generateDialer(); err != nil {
 		return nil, err
 	}
 
 	return c, nil
 }
 
-func (c *Connection) loadKeyPair() error {
+func (c *Connection) loadChiaConfig() error {
+	if c.chiaConfig != nil {
+		return nil
+	}
+	cfg, err := config.GetChiaConfig()
+	if err != nil {
+		return err
+	}
+	c.chiaConfig = cfg
+	return nil
+}
+
+func (c *Connection) loadConfigKeyPair() error {
 	var err error
 
 	c.peerKeyPair, err = c.chiaConfig.FullNode.SSL.LoadPublicKeyPair(c.chiaConfig.ChiaRoot)
@@ -122,7 +145,7 @@ func (c *Connection) Close() {
 func (c *Connection) Handshake() error {
 	// Handshake
 	handshake := &protocols.Handshake{
-		NetworkID:       c.chiaConfig.SelectedNetwork,
+		NetworkID:       c.networkID,
 		ProtocolVersion: protocols.ProtocolVersion,
 		SoftwareVersion: "2.0.0",
 		ServerPort:      c.peerPort,
diff --git a/pkg/peerprotocol/connectionoptions.go b/pkg/peerprotocol/connectionoptions.go
index 29cd4c8..c6a19c6 100644
--- a/pkg/peerprotocol/connectionoptions.go
+++ b/pkg/peerprotocol/connectionoptions.go
@@ -1,6 +1,7 @@
 package peerprotocol
 
 import (
+	"crypto/tls"
 	"time"
 )
 
@@ -14,3 +15,27 @@ func WithHandshakeTimeout(timeout time.Duration) ConnectionOptionFunc {
 		return nil
 	}
 }
+
+// WithPeerPort sets the port for the peer
+func WithPeerPort(port uint16) ConnectionOptionFunc {
+	return func(c *Connection) error {
+		c.peerPort = port
+		return nil
+	}
+}
+
+// WithPeerKeyPair sets the keypair for the peer
+func WithPeerKeyPair(keypair tls.Certificate) ConnectionOptionFunc {
+	return func(c *Connection) error {
+		c.peerKeyPair = &keypair
+		return nil
+	}
+}
+
+// WithNetworkID sets the network ID for the peer
+func WithNetworkID(networkID string) ConnectionOptionFunc {
+	return func(c *Connection) error {
+		c.networkID = networkID
+		return nil
+	}
+}
diff --git a/pkg/protocols/fullnode.go b/pkg/protocols/fullnode.go
index ad14b76..cedde49 100644
--- a/pkg/protocols/fullnode.go
+++ b/pkg/protocols/fullnode.go
@@ -31,4 +31,3 @@ type RequestBlock struct {
 type RespondBlock struct {
 	Block types.FullBlock `streamable:""`
 }
-