From c29a8a797455a2698dbace65bcf578ae72bafac8 Mon Sep 17 00:00:00 2001
From: Aurora Gaffney <aurora@blinklabs.io>
Date: Fri, 22 Mar 2024 17:02:01 -0500
Subject: [PATCH] feat: expose handshake protocol version

Fixes #558
---
 connection.go | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/connection.go b/connection.go
index 2fa92357..5e391a3f 100644
--- a/connection.go
+++ b/connection.go
@@ -59,6 +59,8 @@ type Connection struct {
 	errorChan             chan error
 	protoErrorChan        chan error
 	handshakeFinishedChan chan interface{}
+	handshakeVersion      uint16
+	handshakeVersionData  protocol.VersionData
 	doneChan              chan interface{}
 	waitGroup             sync.WaitGroup
 	onceClose             sync.Once
@@ -205,6 +207,11 @@ func (c *Connection) TxSubmission() *txsubmission.TxSubmission {
 	return c.txSubmission
 }
 
+// ProtocolVersion returns the negotiated protocol version and the version data from the remote peer
+func (c *Connection) ProtocolVersion() (uint16, protocol.VersionData) {
+	return c.handshakeVersion, c.handshakeVersionData
+}
+
 // shutdown performs cleanup operations when the connection is shutdown, either due to explicit Close() or an error
 func (c *Connection) shutdown() {
 	// Gracefully stop the muxer
@@ -285,12 +292,12 @@ func (c *Connection) setupConnection() error {
 		protocol.QueryModeDisabled,
 	)
 	// Perform handshake
-	var handshakeVersion uint16
 	var handshakeFullDuplex bool
 	handshakeConfig := handshake.NewConfig(
 		handshake.WithProtocolVersionMap(protoVersions),
 		handshake.WithFinishedFunc(func(version uint16, versionData protocol.VersionData) error {
-			handshakeVersion = version
+			c.handshakeVersion = version
+			c.handshakeVersionData = versionData
 			if c.useNodeToNodeProto {
 				if versionData.DiffusionMode() == protocol.DiffusionModeInitiatorAndResponder {
 					handshakeFullDuplex = true
@@ -319,7 +326,7 @@ func (c *Connection) setupConnection() error {
 		// This is purposely empty, but we need this case to break out when this channel is closed
 	}
 	// Provide the negotiated protocol version to the various mini-protocols
-	protoOptions.Version = handshakeVersion
+	protoOptions.Version = c.handshakeVersion
 	// Start Goroutine to pass along errors from the mini-protocols
 	c.waitGroup.Add(1)
 	go func() {
@@ -340,7 +347,7 @@ func (c *Connection) setupConnection() error {
 	}()
 	// Configure the relevant mini-protocols
 	if c.useNodeToNodeProto {
-		versionNtN := protocol.GetProtocolVersion(handshakeVersion)
+		versionNtN := protocol.GetProtocolVersion(c.handshakeVersion)
 		protoOptions.Mode = protocol.ProtocolModeNodeToNode
 		c.chainSync = chainsync.New(protoOptions, c.chainSyncConfig)
 		c.blockFetch = blockfetch.New(protoOptions, c.blockFetchConfig)
@@ -377,7 +384,7 @@ func (c *Connection) setupConnection() error {
 			}
 		}
 	} else {
-		versionNtC := protocol.GetProtocolVersion(handshakeVersion)
+		versionNtC := protocol.GetProtocolVersion(c.handshakeVersion)
 		protoOptions.Mode = protocol.ProtocolModeNodeToClient
 		c.chainSync = chainsync.New(protoOptions, c.chainSyncConfig)
 		c.localTxSubmission = localtxsubmission.New(protoOptions, c.localTxSubmissionConfig)