feat: Implement BTCHeaderBytes and BTCHeaderHashBytes types (#29)

proto/babylon/btclightclient/btclightclient.proto

@@ -1,16 +1,16 @@
 syntax = "proto3";
 package babylon.btclightclient.v1;
 
-option go_package = "github.com/babylonchain/babylon/x/btclightclient/types";
+import "gogoproto/gogo.proto";
 
-// BTCHeader defines the structure of the incoming message bytes
-message BTCHeaderBytes {
-  bytes header_bytes = 1;
-}
+option go_package = "github.com/babylonchain/babylon/x/btclightclient/types";
 
 // BaseBTCHeader corresponds to the oldest BTC header maintained in storage
 // It is denoted by the header bytes and the height
 message BaseBTCHeader {
-  BTCHeaderBytes header = 1;
+  bytes header = 1 [
+    (gogoproto.customtype) = "github.com/babylonchain/babylon/types.BTCHeaderBytes"
+  ];
   uint64 height = 2;
 }
+

proto/babylon/btclightclient/genesis.proto

@@ -1,7 +1,6 @@
 syntax = "proto3";
 package babylon.btclightclient.v1;
 
-
 import "gogoproto/gogo.proto";
 import "babylon/btclightclient/params.proto";
 import "babylon/btclightclient/btclightclient.proto";

proto/babylon/btclightclient/query.proto

@@ -5,6 +5,7 @@ import "gogoproto/gogo.proto";
 import "google/api/annotations.proto";
 import "cosmos/base/query/v1beta1/pagination.proto";
 import "babylon/btclightclient/params.proto";
+import "babylon/btclightclient/btclightclient.proto";
 
 option go_package = "github.com/babylonchain/babylon/x/btclightclient/types";
 
@@ -42,14 +43,18 @@ message QueryHashesRequest {
 
 // QueryHashesResponse is response type for the Query/Hashes RPC method.
 message QueryHashesResponse {
-  repeated bytes hashes = 1;
+  repeated bytes hashes = 1 [
+    (gogoproto.customtype) = "github.com/babylonchain/babylon/types.BTCHeaderHashBytes"
+  ];
 
   cosmos.base.query.v1beta1.PageResponse pagination = 2;
 }
 // QueryContainsRequest is request type for the Query/Contains RPC method.
 // It involves checking whether a hash is maintained by the module.
 message QueryContainsRequest {
-  bytes hash = 1;
+  bytes hash = 1 [
+    (gogoproto.customtype) = "github.com/babylonchain/babylon/types.BTCHeaderHashBytes"
+  ];
 }
 
 // QueryContainsResponse is response type for the Query/Contains RPC method.

proto/babylon/btclightclient/tx.proto

@@ -1,6 +1,7 @@
 syntax = "proto3";
 package babylon.btclightclient.v1;
 
+import "gogoproto/gogo.proto";
 import "babylon/btclightclient/btclightclient.proto";
 
 
@@ -14,6 +15,8 @@ service Msg {
 // MsgInsertHeader defines the message for incoming header bytes
 message MsgInsertHeader {
   string signer = 1;
-  BTCHeaderBytes header = 2;
+  bytes header = 2 [
+    (gogoproto.customtype) = "github.com/babylonchain/babylon/types.BTCHeaderBytes"
+  ];
 }
 message MsgInsertHeaderResponse {}

types/btc_header_bytes.go

@@ -0,0 +1,100 @@
+package types
+
+import (
+	"bytes"
+	"encoding/hex"
+	"encoding/json"
+	"github.com/btcsuite/btcd/wire"
+)
+
+type BTCHeaderBytes []byte
+
+func NewBTCHeaderBytesFromHex(hex string) (BTCHeaderBytes, error) {
+	var headerBytes BTCHeaderBytes
+	err := headerBytes.UnmarshalHex(hex)
+	if err != nil {
+		return nil, err
+	}
+	return headerBytes, nil
+}
+
+func (m BTCHeaderBytes) MarshalJSON() ([]byte, error) {
+	hex, err := m.MarshalHex()
+	if err != nil {
+		return nil, err
+	}
+	return json.Marshal(hex)
+}
+
+func (m *BTCHeaderBytes) UnmarshalJSON(bz []byte) error {
+	var headerHexStr string
+	err := json.Unmarshal(bz, &headerHexStr)
+
+	if err != nil {
+		return err
+	}
+
+	return m.UnmarshalHex(headerHexStr)
+}
+
+func (m BTCHeaderBytes) Marshal() ([]byte, error) {
+	return m, nil
+}
+
+func (m *BTCHeaderBytes) Unmarshal(data []byte) error {
+	*m = data
+	return nil
+}
+
+func (m BTCHeaderBytes) MarshalHex() (string, error) {
+	btcdHeader, err := m.ToBlockHeader()
+	if err != nil {
+		return "", err
+	}
+
+	var buf bytes.Buffer
+	btcdHeader.Serialize(&buf)
+	return hex.EncodeToString(buf.Bytes()), nil
+}
+
+func (m *BTCHeaderBytes) UnmarshalHex(header string) error {
+	// Decode the hash string from hex
+	decoded, err := hex.DecodeString(header)
+	if err != nil {
+		return err
+	}
+
+	return m.Unmarshal(decoded)
+}
+
+func (m BTCHeaderBytes) MarshalTo(data []byte) (int, error) {
+	copy(data, m)
+	return len(data), nil
+}
+
+func (m *BTCHeaderBytes) Size() int {
+	bz, _ := m.Marshal()
+	return len(bz)
+}
+
+func (m BTCHeaderBytes) ToBlockHeader() (*wire.BlockHeader, error) {
+	// Create an empty header
+	header := &wire.BlockHeader{}
+
+	// The Deserialize method expects an io.Reader instance
+	reader := bytes.NewReader(m)
+	// Decode the header bytes
+	err := header.Deserialize(reader)
+	// There was a parsing error
+	if err != nil {
+		return nil, err
+	}
+	return header, nil
+}
+
+func (m *BTCHeaderBytes) FromBlockHeader(header *wire.BlockHeader) {
+	var buf bytes.Buffer
+	header.Serialize(&buf)
+
+	*m = buf.Bytes()
+}

types/btc_header_hash_bytes.go

@@ -0,0 +1,91 @@
+package types
+
+import (
+	"encoding/json"
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
+)
+
+type BTCHeaderHashBytes []byte
+
+func NewBTCHeaderHashBytesFromHex(hex string) (BTCHeaderHashBytes, error) {
+	var hashBytes BTCHeaderHashBytes
+	err := hashBytes.UnmarshalHex(hex)
+	if err != nil {
+		return nil, err
+	}
+	return hashBytes, nil
+}
+
+func (m BTCHeaderHashBytes) MarshalJSON() ([]byte, error) {
+	hex, err := m.MarshalHex()
+	if err != nil {
+		return nil, err
+	}
+	// Marshal the JSON from hex format
+	return json.Marshal(hex)
+}
+
+func (m *BTCHeaderHashBytes) UnmarshalJSON(bz []byte) error {
+	var headerHashStr string
+	err := json.Unmarshal(bz, &headerHashStr)
+	if err != nil {
+		return err
+	}
+
+	return m.UnmarshalHex(headerHashStr)
+}
+
+func (m BTCHeaderHashBytes) Marshal() ([]byte, error) {
+	// Just return the bytes
+	return m, nil
+}
+
+func (m *BTCHeaderHashBytes) Unmarshal(bz []byte) error {
+	*m = bz
+	return nil
+}
+
+func (m *BTCHeaderHashBytes) MarshalHex() (string, error) {
+	chHash, err := m.ToChainhash()
+	if err != nil {
+		return "", err
+	}
+
+	return chHash.String(), nil
+}
+
+func (m *BTCHeaderHashBytes) UnmarshalHex(hash string) error {
+	decoded, err := chainhash.NewHashFromStr(hash)
+	if err != nil {
+		return err
+	}
+
+	// Copy the bytes into the instance
+	return m.Unmarshal(decoded[:])
+}
+
+func (m BTCHeaderHashBytes) MarshalTo(data []byte) (int, error) {
+	copy(data, m)
+	return len(data), nil
+}
+
+func (m *BTCHeaderHashBytes) Size() int {
+	bz, _ := m.Marshal()
+	return len(bz)
+}
+
+func (m BTCHeaderHashBytes) ToChainhash() (*chainhash.Hash, error) {
+	return chainhash.NewHash(m)
+}
+
+func (m *BTCHeaderHashBytes) FromChainhash(hash *chainhash.Hash) {
+	var headerHashBytes BTCHeaderHashBytes
+	headerHashBytes.Unmarshal(hash[:])
+	*m = headerHashBytes
+}
+
+func (m BTCHeaderHashBytes) reverse() {
+	for i := 0; i < chainhash.HashSize/2; i++ {
+		m[i], m[chainhash.HashSize-1-i] = m[chainhash.HashSize-1-i], m[i]
+	}
+}

x/btclightclient/client/cli/tx.go

@@ -35,8 +35,10 @@ func CmdTxInsertHeader() *cobra.Command {
 				return err
 			}
 
-			headerBytes := []byte(args[0])
-			msg := types.NewMsgInsertHeader(clientCtx.GetFromAddress(), headerBytes)
+			msg, err := types.NewMsgInsertHeader(clientCtx.GetFromAddress(), args[0])
+			if err != nil {
+				return err
+			}
 
 			return tx.GenerateOrBroadcastTxCLI(clientCtx, cmd.Flags(), msg)
 		},

x/btclightclient/genesis_test.go

@@ -1,7 +1,7 @@
 package btclightclient_test
 
 import (
-	"encoding/hex"
+	bbl "github.com/babylonchain/babylon/types"
 	"testing"
 
 	keepertest "github.com/babylonchain/babylon/testutil/keeper"
@@ -12,7 +12,8 @@ import (
 )
 
 func TestGenesis(t *testing.T) {
-	headerBytes, _ := hex.DecodeString(types.DefaultBaseHeaderHex)
+	headerBytes, _ := bbl.NewBTCHeaderBytesFromHex(types.DefaultBaseHeaderHex)
+
 	genesisState := types.GenesisState{
 		Params:        types.DefaultParams(),
 		BaseBtcHeader: types.DefaultBaseBTCHeader(headerBytes, types.DefaultBaseHeaderHeight),

x/btclightclient/keeper/base_btc_header.go

@@ -1,6 +1,7 @@
 package keeper
 
 import (
+	bbl "github.com/babylonchain/babylon/types"
 	"github.com/babylonchain/babylon/x/btclightclient/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
@@ -23,8 +24,9 @@ func (k Keeper) GetBaseBTCHeader(ctx sdk.Context) types.BaseBTCHeader {
 		return types.BaseBTCHeader{}
 	}
 
-	headerBytes := types.BtcdHeaderToBytes(baseBtcdHeader)
-	return types.BaseBTCHeader{Header: headerBytes, Height: height}
+	var headerBytes bbl.BTCHeaderBytes
+	headerBytes.FromBlockHeader(baseBtcdHeader)
+	return types.BaseBTCHeader{Header: &headerBytes, Height: height}
 }
 
 // SetBaseBTCHeader checks whether a base BTC header exists and
@@ -35,7 +37,7 @@ func (k Keeper) SetBaseBTCHeader(ctx sdk.Context, baseBTCHeader types.BaseBTCHea
 		panic("A base BTC Header has already been set")
 	}
 
-	btcdHeader, err := types.BytesToBtcdHeader(baseBTCHeader.Header)
+	btcdHeader, err := baseBTCHeader.Header.ToBlockHeader()
 	if err != nil {
 		panic("Base BTC Header bytes do not correspond to btcd header")
 	}

x/btclightclient/keeper/grpc_query.go

@@ -2,6 +2,7 @@ package keeper
 
 import (
 	"context"
+	bbl "github.com/babylonchain/babylon/types"
 	"github.com/babylonchain/babylon/x/btclightclient/types"
 	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -25,7 +26,7 @@ func (k Keeper) Hashes(ctx context.Context, req *types.QueryHashesRequest) (*typ
 	if req == nil {
 		return nil, status.Error(codes.InvalidArgument, "invalid request")
 	}
-	var hashes [][]byte
+	var hashes []bbl.BTCHeaderHashBytes
 
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 
@@ -49,7 +50,7 @@ func (k Keeper) Contains(ctx context.Context, req *types.QueryContainsRequest) (
 		return nil, status.Error(codes.InvalidArgument, "invalid request")
 	}
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	chHash, err := types.BytesToChainhash(req.Hash)
+	chHash, err := req.Hash.ToChainhash()
 	if err != nil {
 		return nil, err
 	}

x/btclightclient/keeper/msg_server.go

@@ -22,7 +22,7 @@ func (m msgServer) InsertHeader(ctx context.Context, msg *types.MsgInsertHeader)
 	// 		 that will get rejected.
 
 	// Get Btcd header from bytes
-	btcdHeader, err := types.BytesToBtcdHeader(msg.Header)
+	btcdHeader, err := msg.Header.ToBlockHeader()
 	if err != nil {
 		return nil, err
 	}