feat(): add harvester protocol

Chia-Network · Apr 18, 2024 · fd0ad2f · fd0ad2f
1 parent 83fc55f
commit fd0ad2f
Show file tree

Hide file tree

Showing 11 changed files with 238 additions and 59 deletions.
diff --git a/pkg/peerprotocol/connection.go b/pkg/peerprotocol/connection.go
@@ -19,14 +19,16 @@ import (
 type Connection struct {
 	chiaConfig *config.ChiaConfig
 
-	networkID   string
-	peerIP      *net.IP
-	peerPort    uint16
-	peerKeyPair *tls.Certificate
-	peerDialer  *websocket.Dialer
+	networkID    string
+	peerIP       *net.IP
+	peerPort     uint16
+	peerKeyPair  *tls.Certificate
+	peerDialer   *websocket.Dialer
+	PeerCertHash []byte
 
 	handshakeTimeout time.Duration
 	conn             *websocket.Conn
+	serverPort       uint16
 }
 
 // PeerResponseHandlerFunc is a function that will be called when a response is returned from a peer
@@ -47,13 +49,21 @@ func NewConnection(ip *net.IP, options ...ConnectionOptionFunc) (*Connection, er
 		}
 	}
 
+	// peerPort != local Full Node Port
 	if c.peerPort == 0 {
 		if err := c.loadChiaConfig(); err != nil {
 			return nil, err
 		}
 		c.peerPort = c.chiaConfig.FullNode.Port
 	}
 
+	if c.serverPort == 0 {
+		if err := c.loadChiaConfig(); err != nil {
+			return nil, err
+		}
+		c.serverPort = c.chiaConfig.FullNode.Port
+	}
+
 	if c.peerKeyPair == nil {
 		if err := c.loadChiaConfig(); err != nil {
 			return nil, err
@@ -153,6 +163,8 @@ func (c *Connection) handshake(nodeType protocols.NodeType) error {
 		ProtocolVersion: protocols.ProtocolVersion,
 		SoftwareVersion: "2.0.0",
 		ServerPort:      c.peerPort,
+		SoftwareVersion: "2.2.1",
+		ServerPort:      c.serverPort,
 		NodeType:        nodeType,
 		Capabilities: []protocols.Capability{
 			{

diff --git a/pkg/peerprotocol/connectionoptions.go b/pkg/peerprotocol/connectionoptions.go
@@ -24,6 +24,14 @@ func WithPeerPort(port uint16) ConnectionOptionFunc {
 	}
 }
 
+// WithServerPort sets the port for the peer
+func WithServerPort(port uint16) ConnectionOptionFunc {
+	return func(c *Connection) error {
+		c.serverPort = port
+		return nil
+	}
+}
+
 // WithPeerKeyPair sets the keypair for the peer
 func WithPeerKeyPair(keypair tls.Certificate) ConnectionOptionFunc {
 	return func(c *Connection) error {

diff --git a/pkg/peerprotocol/harvester.go b/pkg/peerprotocol/harvester.go
@@ -1,6 +1,11 @@
 package peerprotocol
 
-import "github.com/chia-network/go-chia-libs/pkg/protocols"
+import (
+	"errors"
+	"github.com/chia-network/go-chia-libs/pkg/protocols"
+)
+
+var ErrInvalidNodeType = errors.New("invalid node type")
 
 // HarvesterProtocol is for interfacing with full nodes via the peer protocol
 type HarvesterProtocol struct {
@@ -13,11 +18,23 @@ func NewHarvesterProtocol(connection *Connection) (*HarvesterProtocol, error) {
 }
 
 // Handshake performs the handshake with the peer
-func (c *HarvesterProtocol) Handshake(data *protocols.HarvesterHandshake) error {
+func (c *HarvesterProtocol) Handshake(nodeType protocols.NodeType) error {
+	if nodeType != protocols.NodeTypeHarvester && nodeType != protocols.NodeTypeFarmer {
+		return ErrInvalidNodeType
+	}
+	return c.handshake(nodeType)
+}
+
+// HarvesterHandshake performs the handshake with the peer
+func (c *HarvesterProtocol) HarvesterHandshake(data *protocols.HarvesterHandshake) error {
 	return c.Do(protocols.ProtocolMessageTypeHarvesterHandshake, data)
 }
 
 // NewSignagePointHarvester sends a new signage point to the harvester
 func (c *HarvesterProtocol) NewSignagePointHarvester(data *protocols.NewSignagePointHarvester) error {
 	return c.Do(protocols.ProtocolMessageTypeNewSignagePointHarvester, data)
 }
+
+func (c *HarvesterProtocol) RequestSignatures(data *protocols.RequestSignatures) error {
+	return c.Do(protocols.ProtocolMessageTypeRequestSignatures, data)
+}
diff --git a/pkg/protocols/farmer.go b/pkg/protocols/farmer.go
@@ -38,3 +38,15 @@ type SignagePointSourceData struct {
 	SubSlotData mo.Option[SPSubSlotSourceData] `streamable:""`
 	VDFData     mo.Option[SPVDFSourceData]     `streamable:""`
 }
+
+// FarmingInfo is the format for the farming_info response
+// https://github.com/Chia-Network/chia-blockchain/blob/main/chia/protocols/farmer_protocol.py#L87
+type FarmingInfo struct {
+	ChallengeHash types.Bytes32 `streamable:""`
+	SPHash        types.Bytes32 `streamable:""`
+	Timestamp     uint64        `streamable:""`
+	Passed        uint32        `streamable:""`
+	Proofs        uint32        `streamable:""`
+	TotalPlots    uint32        `streamable:""`
+	LookupTime    uint64        `streamable:""`
+}
diff --git a/pkg/protocols/farmer_test.go b/pkg/protocols/farmer_test.go
@@ -80,26 +80,6 @@ func TestNewSignagePoint(t *testing.T) {
 			}
 		})
 	}
-	//hexStr := ""
-	//
-	//// Hex to bytes
-	//encodedBytes, err := hex.DecodeString(hexStr)
-	//assert.NoError(t, err)
-	//
-	//rp := &protocols.NewSignagePoint{}
-	//
-	//err = streamable.Unmarshal(encodedBytes, rp)
-	//assert.NoError(t, err)
-
-	//assert.Equal(t, "69171fb97a11a983e1c45f01393a0755e3b65016be6e92ea776ab8ee5b24b66a", hex.EncodeToString(rp.ChallengeHash[:]))
-	//assert.Equal(t, "73165326a79bf653220e33573ef2cace709b31e8f9939cd9e7e7df7f009d0881", hex.EncodeToString(rp.ChallengeChainSP[:]))
-	//assert.Equal(t, "5c0f452f044d025bf0bb4d51744737e09723ab323731a9bbeed9c687e369ec17", hex.EncodeToString(rp.RewardChainSP[:]))
-	//assert.Equal(t, uint64(11776), rp.Difficulty)
-	//assert.Equal(t, uint64(578813952), rp.SubSlotIters)
-	//assert.Equal(t, uint8(26), rp.SignagePointIndex)
-	//assert.Equal(t, uint32(5198109), rp.PeakHeight)
-	//
-	//assert.NotNil(t, rp.SPSourceData)
 
 	// todo test SPSourceData
 }
diff --git a/pkg/protocols/harvester.go b/pkg/protocols/harvester.go
@@ -6,6 +6,14 @@ import (
 	"github.com/chia-network/go-chia-libs/pkg/types"
 )
 
+type HarvesterMode uint8
+
+// https://github.com/Chia-Network/chia-blockchain/blob/main/chia/farmer/farmer_api.py#L97
+const (
+	HarvesterModeCPU HarvesterMode = 1
+	HarvesterModeGPU HarvesterMode = 2
+)
+
 // PoolDifficulty
 // https://github.com/Chia-Network/chia-blockchain/blob/main/chia/protocols/harvester_protocol.py#L23
 type PoolDifficulty struct {
@@ -33,6 +41,75 @@ type NewSignagePointHarvester struct {
 	FilterPrefixBits  uint8            `streamable:""`
 }
 
+// ProofOfSpaceFeeInfo
+// https://github.com/Chia-Network/chia-blockchain/blob/main/chia/protocols/harvester_protocol.py#L50
+type ProofOfSpaceFeeInfo struct {
+	AppliedFeeThreshold uint32 `streamable:""`
+}
+
+// NewProofOfSpace
+// https://github.com/Chia-Network/chia-blockchain/blob/main/chia/protocols/harvester_protocol.py#L56
+type NewProofOfSpace struct {
+	ChallengeHash               types.Bytes32                  `streamable:""`
+	SPHash                      types.Bytes32                  `streamable:""`
+	PlotIdentifier              string                         `streamable:""`
+	Proof                       types.ProofOfSpace             `streamable:""`
+	SignagePointIndex           uint8                          `streamable:""`
+	IncludeSourceSignatureData  bool                           `streamable:""`
+	FarmerRewardAddressOverride mo.Option[types.Bytes32]       `streamable:""`
+	FeeInfo                     mo.Option[ProofOfSpaceFeeInfo] `streamable:""`
+}
+
+// SigningDataKind
+// https://github.com/Chia-Network/chia-blockchain/blob/main/chia/protocols/harvester_protocol.py#L68
+type SigningDataKind uint8
+
+const (
+	FOLIAGE_BLOCK_DATA SigningDataKind = iota + 1
+	FOLIAGE_TRANSACTION_BLOCK
+	CHALLENGE_CHAIN_VDF
+	REWARD_CHAIN_VDF
+	CHALLENGE_CHAIN_SUB_SLOT
+	REWARD_CHAIN_SUB_SLOT
+	PARTIAL
+)
+
+// SignatureRequestSourceData
+// https://github.com/Chia-Network/chia-blockchain/blob/main/chia/protocols/harvester_protocol.py#L80
+type SignatureRequestSourceData struct {
+	Kind SigningDataKind `streamable:""`
+	Data []byte          `streamable:""`
+}
+
+// RequestSignatures
+// https://github.com/Chia-Network/chia-blockchain/blob/main/chia/protocols/harvester_protocol.py#L89
+type RequestSignatures struct {
+	PlotIdentifier    string                                             `streamable:""`
+	ChallengeHash     types.Bytes32                                      `streamable:""`
+	SPHash            types.Bytes32                                      `streamable:""`
+	Messages          []types.Bytes32                                    `streamable:""`
+	MessageData       mo.Option[[]mo.Option[SignatureRequestSourceData]] `streamable:""`
+	RCBlockUnfinished mo.Option[types.RewardChainBlockUnfinished]        `streamable:""`
+}
+
+type MessageSignature struct {
+	types.Bytes32   `streamable:""`
+	types.G2Element `streamable:""`
+}
+
+// RespondSignatures
+// https://github.com/Chia-Network/chia-blockchain/blob/main/chia/protocols/harvester_protocol.py#L101
+type RespondSignatures struct {
+	PlotIdentifier              string                   `streamable:""`
+	ChallengeHash               types.Bytes32            `streamable:""`
+	SPHash                      types.Bytes32            `streamable:""`
+	LocalPK                     types.G1Element          `streamable:""`
+	FarmerPK                    types.G1Element          `streamable:""`
+	MessageSignatures           []MessageSignature       `streamable:""`
+	IncludeSourceSignatureData  bool                     `streamable:""`
+	FarmerRewardAddressOverride mo.Option[types.Bytes32] `streamable:""`
+}
+
 // Plot is the plot definition in the harvester protocol
 // https://github.com/Chia-Network/chia-blockchain/blob/main/chia/protocols/harvester_protocol.py#L114
 type Plot struct {
@@ -46,3 +123,21 @@ type Plot struct {
 	TimeModified           types.Timestamp            `json:"time_modified"`
 	CompressionLevel       mo.Option[uint8]           `json:"compression_level"`
 }
+
+// PlotSyncStart is the plot sync start message
+// https://github.com/Chia-Network/chia-blockchain/blob/main/chia/protocols/harvester_protocol.py#L150
+type PlotSyncStart struct {
+	Identifier     PlotSyncIdentifier `streamable:""`
+	Initial        bool               `streamable:""`
+	LastSyncID     uint64             `streamable:""`
+	PlotFileCount  uint32             `streamable:""`
+	HarvestingMode uint8              `streamable:""`
+}
+
+// PlotSyncIdentifier
+// https://github.com/Chia-Network/chia-blockchain/blob/main/chia/protocols/harvester_protocol.py#L142
+type PlotSyncIdentifier struct {
+	Timestamp uint64 `streamable:""`
+	SyncID    uint64 `streamable:""`
+	MessageID uint64 `streamable:""`
+}
diff --git a/pkg/protocols/harvester_test.go b/pkg/protocols/harvester_test.go
@@ -0,0 +1,48 @@
+package protocols_test
+
+import (
+	"encoding/hex"
+	"github.com/chia-network/go-chia-libs/pkg/protocols"
+	"github.com/chia-network/go-chia-libs/pkg/streamable"
+	"github.com/chia-network/go-chia-libs/pkg/types"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestHarvesterHandshakeMarshal(t *testing.T) {
+	// Hex to bytes
+	fpk1, err := hex.DecodeString("c31c4d1b4cc006eb8ea294656aba445f43e4e3d78910392c6880a1e9bd806d0419acd4e60c4abd533f11dafafbf23e7f")
+	assert.NoError(t, err)
+
+	fpk2, err := hex.DecodeString("af397582efbcdc276bbe5ee3b73981c0c9ff5b3fafeae5c34a0f9ea27a67023966a53d6a3f5c37ee9272cdbe3f3cfc9c")
+	assert.NoError(t, err)
+
+	ppk1, err := hex.DecodeString("0da874d9ede1dc629731c05f853c7d20eb987b00465f1acdd84d1256b17762a2b22b3bb2cbd25ac3178515b4d75311fd")
+	assert.NoError(t, err)
+
+	ppk2, err := hex.DecodeString("d74f2a6f8e9960417df6665fefd7598dbdc5a5fb75b9e0b0f2f3ffc9b270de54dc368af2fe63aa122e7b558f792bc891")
+	assert.NoError(t, err)
+
+	rp := protocols.HarvesterHandshake{
+		FarmerPublicKeys: []types.G1Element{types.G1Element(fpk1), types.G1Element(fpk2)},
+		PoolPublicKeys:   []types.G1Element{types.G1Element(ppk1), types.G1Element(ppk2)},
+	}
+
+	b, err := streamable.Marshal(rp)
+	assert.NoError(t, err)
+
+	assert.Len(t, b, 200)
+
+	// Unmarshal
+	rp2 := &protocols.HarvesterHandshake{}
+	err = streamable.Unmarshal(b, rp2)
+	assert.NoError(t, err)
+
+	assert.Len(t, rp2.FarmerPublicKeys, 2)
+	assert.Len(t, rp2.PoolPublicKeys, 2)
+
+	assert.Equal(t, fpk1, rp2.FarmerPublicKeys[0][:])
+	assert.Equal(t, fpk2, rp2.FarmerPublicKeys[1][:])
+	assert.Equal(t, ppk1, rp2.PoolPublicKeys[0][:])
+	assert.Equal(t, ppk2, rp2.PoolPublicKeys[1][:])
+}
diff --git a/pkg/protocols/messagetypes.go b/pkg/protocols/messagetypes.go
@@ -5,17 +5,29 @@ package protocols
 type ProtocolMessageType uint8
 
 const (
+	// there are many more of these in Chia - only listing the ones current is use for now
+
 	// ProtocolMessageTypeHandshake handshake
 	ProtocolMessageTypeHandshake ProtocolMessageType = 1
 
 	// ProtocolMessageTypeHarvesterHandshake harvester_handshake
 	ProtocolMessageTypeHarvesterHandshake ProtocolMessageType = 3
 
-	// there are many more of these in Chia - only listing the ones current is use for now
+	// ProtocolMessageTypeNewProofOfSpace new_proof_of_space
+	ProtocolMessageTypeNewProofOfSpace ProtocolMessageType = 5
+
+	// ProtocolMessageTypeRequestSignatures request_signatures
+	ProtocolMessageTypeRequestSignatures ProtocolMessageType = 6
+
+	// ProtocolMessageTypeRespondSignatures respond_signatures
+	ProtocolMessageTypeRespondSignatures ProtocolMessageType = 7
 
 	// ProtocolMessageTypeNewSignagePoint new_signage_point
 	ProtocolMessageTypeNewSignagePoint ProtocolMessageType = 8
 
+	// ProtocolMessageTypeFarmingInfo farming_info
+	ProtocolMessageTypeFarmingInfo ProtocolMessageType = 12
+
 	// ProtocolMessageTypeNewPeak new_peak
 	ProtocolMessageTypeNewPeak ProtocolMessageType = 20
 
@@ -33,4 +45,7 @@ const (
 
 	// ProtocolMessageTypeNewSignagePointHarvester new_signage_point_harvester
 	ProtocolMessageTypeNewSignagePointHarvester ProtocolMessageType = 66
+
+	// ProtocolMessageTypePlotSyncStart plot_sync_start
+	ProtocolMessageTypePlotSyncStart ProtocolMessageType = 78
 )
diff --git a/pkg/protocols/shared.go b/pkg/protocols/shared.go
@@ -3,7 +3,7 @@ package protocols
 // ProtocolVersion Current supported Protocol Version
 // Not all of this is supported, but this was the current version at the time
 // This library was started
-const ProtocolVersion string = "0.0.33"
+const ProtocolVersion string = "0.0.36"
 
 // NodeType is the type of peer (farmer, full node, etc)
 // Source for node types is chia/server/outbound_messages.py