Skip to content

Commit

Permalink
Add NewHeader constructor function. Remove duplicat functions MakeHea…
Browse files Browse the repository at this point in the history 
…der and buildHeader.

Use NewHeader everywhere is possible.
  • Loading branch information
@alexeykiselev
alexeykiselev committed Jan 9, 2025
1 parent ac521c6 commit fd6db0a
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 132 deletions.
21 changes: 4 additions & 17 deletions pkg/proto/microblock.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -252,7 +252,7 @@ func (a *MicroBlockMessage) WriteTo(w io.Writer) (int64, error) {
return 0, err
}

h, err := MakeHeader(ContentIDMicroblock, buf.Bytes())
h, err := NewHeader(ContentIDMicroblock, buf.Bytes())
if err != nil {
return 0, err
}
Expand Down Expand Up @@ -308,16 +308,10 @@ func (a *MicroBlockInvMessage) ReadFrom(_ io.Reader) (n int64, err error) {
}

func (a *MicroBlockInvMessage) WriteTo(w io.Writer) (n int64, err error) {
var h Header
h.Length = maxHeaderLength + uint32(len(a.Body)) - 4
h.Magic = headerMagic
h.ContentID = ContentIDInvMicroblock
h.payloadLength = common.SafeIntToUint32(len(a.Body))
dig, err := crypto.FastHash(a.Body)
h, err := NewHeader(ContentIDInvMicroblock, a.Body)
if err != nil {
return 0, err
}
copy(h.PayloadChecksum[:], dig[:4])
n1, err := h.WriteTo(w)
if err != nil {
return 0, err
Expand Down Expand Up @@ -351,21 +345,14 @@ func (a *MicroBlockRequestMessage) ReadFrom(_ io.Reader) (n int64, err error) {
}

func (a *MicroBlockRequestMessage) WriteTo(w io.Writer) (int64, error) {
var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(a.TotalBlockSig)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = ContentIDMicroblockRequest
h.payloadLength = common.SafeIntToUint32(len(a.TotalBlockSig))
dig, err := crypto.FastHash(a.TotalBlockSig)
h, err := NewHeader(ContentIDMicroblockRequest, a.TotalBlockSig)
if err != nil {
return 0, err
}
copy(h.PayloadChecksum[:], dig[:4])
n2, err := h.WriteTo(w)
if err != nil {
return 0, err
}

n3, err := w.Write(a.TotalBlockSig)
if err != nil {
return 0, err
Expand Down Expand Up @@ -542,7 +529,7 @@ func (a *PBMicroBlockMessage) WriteTo(w io.Writer) (int64, error) {
return 0, err
}

h, err := MakeHeader(ContentIDPBMicroBlock, buf.Bytes())
h, err := NewHeader(ContentIDPBMicroBlock, buf.Bytes())
if err != nil {
return 0, err
}
Expand Down
153 changes: 38 additions & 115 deletions pkg/proto/proto.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ import (
"strconv"
"strings"

"github.com/ccoveille/go-safecast"

"github.com/pkg/errors"
"github.com/valyala/bytebufferpool"

Expand Down Expand Up @@ -79,6 +81,26 @@ type Header struct {
PayloadChecksum [headerChecksumLen]byte
}

func NewHeader(contentID PeerMessageID, body []byte) (Header, error) {
bl, err := safecast.ToUint32(len(body))
if err != nil {
return Header{}, fmt.Errorf("failed to create header: %w", err)
}
h := Header{
Length: maxHeaderLength + bl - headerChecksumLen,
Magic: headerMagic,
ContentID: contentID,
payloadLength: bl,
PayloadChecksum: [4]byte{},
}
dig, err := crypto.FastHash(body)
if err != nil {
return Header{}, fmt.Errorf("failed to create header: %w", err)
}
copy(h.PayloadChecksum[:], dig[:headerChecksumLen])
return h, nil
}

func (h *Header) MarshalBinary() ([]byte, error) {
data := make([]byte, h.HeaderLength())
if _, err := h.Copy(data); err != nil {
Expand Down Expand Up @@ -1107,16 +1129,10 @@ func (m *GetSignaturesMessage) MarshalBinary() ([]byte, error) {
body = append(body, b[:]...)
}

var h Header
h.Length = maxHeaderLength + uint32(len(body)) - 4
h.Magic = headerMagic
h.ContentID = ContentIDGetSignatures
h.payloadLength = common.SafeIntToUint32(len(body))
dig, err := crypto.FastHash(body)
h, err := NewHeader(ContentIDGetSignatures, body)
if err != nil {
return nil, err
}
copy(h.PayloadChecksum[:], dig[:4])

hdr, err := h.MarshalBinary()
if err != nil {
Expand Down Expand Up @@ -1198,16 +1214,10 @@ func (m *SignaturesMessage) MarshalBinary() ([]byte, error) {
body = append(body, b[:]...)
}

var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(body)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = ContentIDSignatures
h.payloadLength = common.SafeIntToUint32(len(body))
dig, err := crypto.FastHash(body)
h, err := NewHeader(ContentIDSignatures, body)
if err != nil {
return nil, err
}
copy(h.PayloadChecksum[:], dig[:4])

hdr, err := h.MarshalBinary()
if err != nil {
Expand Down Expand Up @@ -1285,17 +1295,10 @@ type GetBlockMessage struct {
func (m *GetBlockMessage) MarshalBinary() ([]byte, error) {
body := m.BlockID.Bytes()

var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(body)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = ContentIDGetBlock
h.payloadLength = common.SafeIntToUint32(len(body))
dig, err := crypto.FastHash(body)
h, err := NewHeader(ContentIDGetBlock, body)
if err != nil {
return nil, err
}
copy(h.PayloadChecksum[:], dig[:4])

hdr, err := h.MarshalBinary()
if err != nil {
return nil, err
Expand Down Expand Up @@ -1381,17 +1384,10 @@ type BlockMessage struct {

// MarshalBinary encodes BlockMessage to binary form
func (m *BlockMessage) MarshalBinary() ([]byte, error) {
var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(m.BlockBytes)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = ContentIDBlock
h.payloadLength = common.SafeIntToUint32(len(m.BlockBytes))
dig, err := crypto.FastHash(m.BlockBytes)
h, err := NewHeader(ContentIDBlock, m.BlockBytes)
if err != nil {
return nil, err
}
copy(h.PayloadChecksum[:], dig[:4])

hdr, err := h.MarshalBinary()
if err != nil {
return nil, err
Expand All @@ -1400,20 +1396,6 @@ func (m *BlockMessage) MarshalBinary() ([]byte, error) {
return hdr, nil
}

func MakeHeader(contentID PeerMessageID, payload []byte) (Header, error) {
var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(payload)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = contentID
h.payloadLength = common.SafeIntToUint32(len(payload))
dig, err := crypto.FastHash(payload)
if err != nil {
return Header{}, err
}
copy(h.PayloadChecksum[:], dig[:4])
return h, nil
}

// UnmarshalBinary decodes BlockMessage from binary from
func (m *BlockMessage) UnmarshalBinary(data []byte) error {
var h Header
Expand Down Expand Up @@ -1464,17 +1446,10 @@ type ScoreMessage struct {

// MarshalBinary encodes ScoreMessage to binary form
func (m *ScoreMessage) MarshalBinary() ([]byte, error) {
var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(m.Score)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = ContentIDScore
h.payloadLength = common.SafeIntToUint32(len(m.Score))
dig, err := crypto.FastHash(m.Score)
h, err := NewHeader(ContentIDScore, m.Score)
if err != nil {
return nil, err
}
copy(h.PayloadChecksum[:], dig[:4])

hdr, err := h.MarshalBinary()
if err != nil {
return nil, err
Expand Down Expand Up @@ -1531,17 +1506,10 @@ type TransactionMessage struct {

// MarshalBinary encodes TransactionMessage to binary form
func (m *TransactionMessage) MarshalBinary() ([]byte, error) {
var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(m.Transaction)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = ContentIDTransaction
h.payloadLength = common.SafeIntToUint32(len(m.Transaction))
dig, err := crypto.FastHash(m.Transaction)
h, err := NewHeader(ContentIDTransaction, m.Transaction)
if err != nil {
return nil, err
}
copy(h.PayloadChecksum[:], dig[:4])

hdr, err := h.MarshalBinary()
if err != nil {
return nil, err
Expand Down Expand Up @@ -1619,16 +1587,10 @@ func (m *CheckPointMessage) MarshalBinary() ([]byte, error) {
body = append(body, c.Signature[:]...)
}

var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(body)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = ContentIDCheckpoint
h.payloadLength = common.SafeIntToUint32(len(body))
dig, err := crypto.FastHash(body)
h, err := NewHeader(ContentIDCheckpoint, body)
if err != nil {
return nil, err
}
copy(h.PayloadChecksum[:], dig[:4])

hdr, err := h.MarshalBinary()
if err != nil {
Expand Down Expand Up @@ -1703,17 +1665,10 @@ type PBBlockMessage struct {

// MarshalBinary encodes PBBlockMessage to binary form
func (m *PBBlockMessage) MarshalBinary() ([]byte, error) {
var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(m.PBBlockBytes)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = ContentIDPBBlock
h.payloadLength = common.SafeIntToUint32(len(m.PBBlockBytes))
dig, err := crypto.FastHash(m.PBBlockBytes)
h, err := NewHeader(ContentIDPBBlock, m.PBBlockBytes)
if err != nil {
return nil, err
}
copy(h.PayloadChecksum[:], dig[:4])

hdr, err := h.MarshalBinary()
if err != nil {
return nil, err
Expand Down Expand Up @@ -1772,17 +1727,10 @@ type PBTransactionMessage struct {

// MarshalBinary encodes PBTransactionMessage to binary form
func (m *PBTransactionMessage) MarshalBinary() ([]byte, error) {
var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(m.Transaction)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = ContentIDPBTransaction
h.payloadLength = common.SafeIntToUint32(len(m.Transaction))
dig, err := crypto.FastHash(m.Transaction)
h, err := NewHeader(ContentIDPBTransaction, m.Transaction)
if err != nil {
return nil, err
}
copy(h.PayloadChecksum[:], dig[:4])

hdr, err := h.MarshalBinary()
if err != nil {
return nil, err
Expand Down Expand Up @@ -1912,16 +1860,10 @@ func (m *GetBlockIdsMessage) MarshalBinary() ([]byte, error) {
body = append(body, b...)
}

var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(body)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = ContentIDGetBlockIDs
h.payloadLength = common.SafeIntToUint32(len(body))
dig, err := crypto.FastHash(body)
h, err := NewHeader(ContentIDGetBlockIDs, body)
if err != nil {
return nil, err
}
copy(h.PayloadChecksum[:], dig[:4])

hdr, err := h.MarshalBinary()
if err != nil {
Expand Down Expand Up @@ -2008,16 +1950,10 @@ func (m *BlockIdsMessage) MarshalBinary() ([]byte, error) {
body = append(body, b...)
}

var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(body)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = ContentIDBlockIDs
h.payloadLength = common.SafeIntToUint32(len(body))
dig, err := crypto.FastHash(body)
h, err := NewHeader(ContentIDBlockIDs, body)
if err != nil {
return nil, err
}
copy(h.PayloadChecksum[:], dig[:4])

hdr, err := h.MarshalBinary()
if err != nil {
Expand Down Expand Up @@ -2123,20 +2059,6 @@ type MiningLimits struct {
MaxTxsSizeInBytes int
}

func buildHeader(body []byte, messID PeerMessageID) (Header, error) {
var h Header
h.Length = maxHeaderLength + common.SafeIntToUint32(len(body)) - headerChecksumLen
h.Magic = headerMagic
h.ContentID = messID
h.payloadLength = common.SafeIntToUint32(len(body))
dig, err := crypto.FastHash(body)
if err != nil {
return Header{}, err
}
copy(h.PayloadChecksum[:], dig[:headerChecksumLen])
return h, nil
}

type GetBlockSnapshotMessage struct {
BlockID BlockID
}
Expand Down Expand Up @@ -2173,7 +2095,8 @@ func (m *GetBlockSnapshotMessage) UnmarshalBinary(data []byte) error {

func (m *GetBlockSnapshotMessage) MarshalBinary() ([]byte, error) {
body := m.BlockID.Bytes()
h, err := buildHeader(body, ContentIDGetBlockSnapshot)

h, err := NewHeader(ContentIDGetBlockSnapshot, body)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -2231,7 +2154,7 @@ func (m *BlockSnapshotMessage) UnmarshalBinary(data []byte) error {
func (m *BlockSnapshotMessage) MarshalBinary() ([]byte, error) {
body := m.Bytes

h, err := buildHeader(body, ContentIDBlockSnapshot)
h, err := NewHeader(ContentIDBlockSnapshot, body)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -2290,7 +2213,7 @@ func (m *MicroBlockSnapshotMessage) UnmarshalBinary(data []byte) error {
func (m *MicroBlockSnapshotMessage) MarshalBinary() ([]byte, error) {
body := m.Bytes

h, err := buildHeader(body, ContentIDMicroBlockSnapshot)
h, err := NewHeader(ContentIDMicroBlockSnapshot, body)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -2340,7 +2263,7 @@ func (m *MicroBlockSnapshotRequestMessage) UnmarshalBinary(data []byte) error {
func (m *MicroBlockSnapshotRequestMessage) MarshalBinary() ([]byte, error) {
body := m.BlockIDBytes

h, err := buildHeader(body, ContentIDMicroBlockSnapshotRequest)
h, err := NewHeader(ContentIDMicroBlockSnapshotRequest, body)
if err != nil {
return nil, err
}
Expand Down

0 comments on commit fd6db0a

Please sign in to comment.