Merge pull request #45 from TRON-US/BTFS-1137

refactor ledger grpc client

ledger/ledger.go

@@ -2,65 +2,30 @@ package ledger
 
 import (
 	"context"
-	"crypto/tls"
-	"crypto/x509"
-	"fmt"
 	"time"
 
-	ledgerPb "github.com/tron-us/go-btfs-common/protos/ledger"
-	"github.com/tron-us/go-common/v2/log"
-
 	ic "github.com/libp2p/go-libp2p-core/crypto"
+	ledgerpb "github.com/tron-us/go-btfs-common/protos/ledger"
 	"github.com/tron-us/protobuf/proto"
-	"go.uber.org/zap"
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials"
 )
 
-func LedgerConnection(ledgerAddr, certFile string) (*grpc.ClientConn, error) {
-	var grpcOption grpc.DialOption
-	if certFile == "" {
-		grpcOption = grpc.WithInsecure()
-	} else {
-		b := []byte(certFile)
-		cp := x509.NewCertPool()
-		if !cp.AppendCertsFromPEM(b) {
-			return nil, fmt.Errorf("credentials: failed to append certificates")
-		}
-		credential := credentials.NewTLS(&tls.Config{RootCAs: cp})
-		grpcOption = grpc.WithTransportCredentials(credential)
-	}
-	conn, err := grpc.Dial(ledgerAddr, grpcOption)
-	if err != nil {
-		return nil, err
-	}
-	return conn, nil
-}
-
-func CloseConnection(conn *grpc.ClientConn) {
-	if conn != nil {
-		if err := conn.Close(); err != nil {
-			log.Error("Failed to close connection: ", zap.Error(err))
-		}
-	}
-}
-
-func NewClient(conn *grpc.ClientConn) ledgerPb.ChannelsClient {
-	return ledgerPb.NewChannelsClient(conn)
+func NewClient(conn *grpc.ClientConn) ledgerpb.ChannelsClient {
+	return ledgerpb.NewChannelsClient(conn)
 }
 
-func NewAccount(pubKey ic.PubKey, amount int64) (*ledgerPb.Account, error) {
+func NewAccount(pubKey ic.PubKey, amount int64) (*ledgerpb.Account, error) {
 	addr, err := pubKey.Raw()
 	if err != nil {
 		return nil, err
 	}
-	return &ledgerPb.Account{
-		Address: &ledgerPb.PublicKey{Key: addr},
+	return &ledgerpb.Account{
+		Address: &ledgerpb.PublicKey{Key: addr},
 		Balance: amount,
 	}, nil
 }
 
-func NewChannelCommit(fromKey ic.PubKey, toKey ic.PubKey, amount int64) (*ledgerPb.ChannelCommit, error) {
+func NewChannelCommit(fromKey ic.PubKey, toKey ic.PubKey, amount int64) (*ledgerpb.ChannelCommit, error) {
 	fromAddr, err := fromKey.Raw()
 	if err != nil {
 		return nil, err
@@ -69,66 +34,66 @@ func NewChannelCommit(fromKey ic.PubKey, toKey ic.PubKey, amount int64) (*ledger
 	if err != nil {
 		return nil, err
 	}
-	return &ledgerPb.ChannelCommit{
-		Payer:     &ledgerPb.PublicKey{Key: fromAddr},
-		Recipient: &ledgerPb.PublicKey{Key: toAddr},
+	return &ledgerpb.ChannelCommit{
+		Payer:     &ledgerpb.PublicKey{Key: fromAddr},
+		Recipient: &ledgerpb.PublicKey{Key: toAddr},
 		Amount:    amount,
 		PayerId:   time.Now().UnixNano(),
 	}, err
 }
 
-func NewChannelState(id *ledgerPb.ChannelID, sequence int64, fromAccount *ledgerPb.Account, toAccount *ledgerPb.Account) *ledgerPb.ChannelState {
-	return &ledgerPb.ChannelState{
+func NewChannelState(id *ledgerpb.ChannelID, sequence int64, fromAccount *ledgerpb.Account, toAccount *ledgerpb.Account) *ledgerpb.ChannelState {
+	return &ledgerpb.ChannelState{
 		Id:       id,
 		Sequence: sequence,
 		From:     fromAccount,
 		To:       toAccount,
 	}
 }
 
-func NewSignedChannelState(channelState *ledgerPb.ChannelState, fromSig []byte, toSig []byte) *ledgerPb.SignedChannelState {
-	return &ledgerPb.SignedChannelState{
+func NewSignedChannelState(channelState *ledgerpb.ChannelState, fromSig []byte, toSig []byte) *ledgerpb.SignedChannelState {
+	return &ledgerpb.SignedChannelState{
 		Channel:       channelState,
 		FromSignature: fromSig,
 		ToSignature:   toSig,
 	}
 }
 
-func ImportAccount(ctx context.Context, pubKey ic.PubKey, ledgerClient ledgerPb.ChannelsClient) (*ledgerPb.Account, error) {
+func ImportAccount(ctx context.Context, pubKey ic.PubKey, ledgerClient ledgerpb.ChannelsClient) (*ledgerpb.Account, error) {
 	keyBytes, err := pubKey.Raw()
 	if err != nil {
 		return nil, err
 	}
-	res, err := ledgerClient.CreateAccount(ctx, &ledgerPb.PublicKey{Key: keyBytes})
+	res, err := ledgerClient.CreateAccount(ctx, &ledgerpb.PublicKey{Key: keyBytes})
 	if err != nil {
 		return nil, err
 	}
 	return res.GetAccount(), nil
 }
 
-func ImportSignedAccount(ctx context.Context, privKey ic.PrivKey, pubKey ic.PubKey, ledgerClient ledgerPb.ChannelsClient) (*ledgerPb.SignedCreateAccountResult, error) {
+func ImportSignedAccount(ctx context.Context, privKey ic.PrivKey, pubKey ic.PubKey, ledgerClient ledgerpb.ChannelsClient) (*ledgerpb.SignedCreateAccountResult, error) {
 	pubKeyBytes, err := pubKey.Raw()
 	if err != nil {
 		return nil, err
 	}
-	singedPubKey := &ledgerPb.PublicKey{Key: pubKeyBytes}
+	singedPubKey := &ledgerpb.PublicKey{Key: pubKeyBytes}
 	sigBytes, err := proto.Marshal(singedPubKey)
 	signature, err := privKey.Sign(sigBytes)
 	if err != nil {
 		return nil, err
 	}
-	signedPubkey := &ledgerPb.SignedPublicKey{Key: singedPubKey, Signature: signature}
+	signedPubkey := &ledgerpb.SignedPublicKey{Key: singedPubKey, Signature: signature}
 	return ledgerClient.SignedCreateAccount(ctx, signedPubkey)
 }
 
-func CreateChannel(ctx context.Context, ledgerClient ledgerPb.ChannelsClient, channelCommit *ledgerPb.ChannelCommit, sig []byte) (*ledgerPb.ChannelID, error) {
-	return ledgerClient.CreateChannel(ctx, &ledgerPb.SignedChannelCommit{
+func CreateChannel(ctx context.Context, ledgerClient ledgerpb.ChannelsClient, channelCommit *ledgerpb.ChannelCommit, sig []byte) (*ledgerpb.ChannelID, error) {
+	return ledgerClient.CreateChannel(ctx, &ledgerpb.SignedChannelCommit{
 		Channel:   channelCommit,
 		Signature: sig,
 	})
 }
 
-func CloseChannel(ctx context.Context, ledgerClient ledgerPb.ChannelsClient, signedChannelState *ledgerPb.SignedChannelState) error {
+func CloseChannel(ctx context.Context, ledgerClient ledgerpb.ChannelsClient, signedChannelState *ledgerpb.SignedChannelState) error {
 	_, err := ledgerClient.CloseChannel(ctx, signedChannelState)
 	if err != nil {
 		return err

ledger/sample/demo.go

@@ -5,6 +5,8 @@ import (
 
 	"github.com/tron-us/go-btfs-common/crypto"
 	"github.com/tron-us/go-btfs-common/ledger"
+	ledgerpb "github.com/tron-us/go-btfs-common/protos/ledger"
+	"github.com/tron-us/go-btfs-common/utils/grpc"
 	"github.com/tron-us/go-common/v2/log"
 
 	"go.uber.org/zap"
@@ -16,74 +18,72 @@ const (
 )
 
 func main() {
-	ctx := context.Background()
-	// build connection with ledger
-	clientConn, err := ledger.LedgerConnection("ledger-dev.bt.co:443", "")
+	err := grpc.LedgerClient("https://ledger-dev.bt.co:443").WithContext(context.Background(),
+		func(ctx context.Context, ledgerClient ledgerpb.ChannelsClient) error {
+			// create payer Account
+			payerPrivKey, err := crypto.ToPrivKey(PayerPrivKeyString)
+			if err != nil {
+				log.Panic("can not convert to private key", zap.Error(err))
+			}
+			payerPubKey := payerPrivKey.GetPublic()
+			_, err = ledger.ImportSignedAccount(ctx, payerPrivKey, payerPubKey, ledgerClient)
+			if err != nil {
+				log.Panic("can not create account on ledger", zap.Error(err))
+			}
+			// create receiver account
+			recvPrivKey, err := crypto.ToPrivKey(ReceiverPrivKeyString)
+			if err != nil {
+				log.Panic("can not convert to private key", zap.Error(err))
+			}
+			recvPubKey := recvPrivKey.GetPublic()
+			_, err = ledger.ImportSignedAccount(ctx, recvPrivKey, recvPubKey, ledgerClient)
+			if err != nil {
+				log.Panic("can not create account on ledger", zap.Error(err))
+			}
+			// prepare channel commit
+			amount := int64(1)
+			channelCommit, err := ledger.NewChannelCommit(payerPubKey, recvPubKey, amount)
+			if err != nil {
+				log.Panic("can not create channel commit", zap.Error(err))
+			}
+			// sign for the channel commit
+			fromSig, err := crypto.Sign(payerPrivKey, channelCommit)
+			if err != nil {
+				log.Panic("fail to sign channel commit", zap.Error(err))
+			}
+			// create channel: payer start the channel
+			channelID, err := ledger.CreateChannel(ctx, ledgerClient, channelCommit, fromSig)
+			if err != nil {
+				log.Panic("fail to create channel", zap.Error(err))
+			}
+			// channel state: transfer money from -> to
+			fromAcc, err := ledger.NewAccount(payerPubKey, 0)
+			if err != nil {
+				log.Panic("wrong account on channel", zap.Error(err))
+			}
+			toAcc, err := ledger.NewAccount(recvPubKey, amount)
+			if err != nil {
+				log.Panic("wrong account on channel", zap.Error(err))
+			}
+			channelState := ledger.NewChannelState(channelID, 1, fromAcc, toAcc)
+			// need permission from both account, get signature from both
+			fromSigState, err := crypto.Sign(payerPrivKey, channelState)
+			if err != nil {
+				log.Panic("error when signing the channel state", zap.Error(err))
+			}
+			toSigState, err := crypto.Sign(recvPrivKey, channelState)
+			if err != nil {
+				log.Panic("error when signing the channel state", zap.Error(err))
+			}
+			signedChannelState := ledger.NewSignedChannelState(channelState, fromSigState, toSigState)
+			// close channel
+			err = ledger.CloseChannel(ctx, ledgerClient, signedChannelState)
+			if err != nil {
+				log.Panic("fail to close channel", zap.Error(err))
+			}
+			return nil
+		})
 	if err != nil {
-		log.Panic("fail to connect", zap.Error(err))
-	}
-	defer ledger.CloseConnection(clientConn)
-	// new ledger client
-	ledgerClient := ledger.NewClient(clientConn)
-	// create payer Account
-	payerPrivKey, err := crypto.ToPrivKey(PayerPrivKeyString)
-	if err != nil {
-		log.Panic("can not convert to private key", zap.Error(err))
-	}
-	payerPubKey := payerPrivKey.GetPublic()
-	_, err = ledger.ImportSignedAccount(ctx, payerPrivKey, payerPubKey, ledgerClient)
-	if err != nil {
-		log.Panic("can not create account on ledger", zap.Error(err))
-	}
-	// create receiver account
-	recvPrivKey, err := crypto.ToPrivKey(ReceiverPrivKeyString)
-	if err != nil {
-		log.Panic("can not convert to private key", zap.Error(err))
-	}
-	recvPubKey := recvPrivKey.GetPublic()
-	_, err = ledger.ImportSignedAccount(ctx, recvPrivKey, recvPubKey, ledgerClient)
-	if err != nil {
-		log.Panic("can not create account on ledger", zap.Error(err))
-	}
-	// prepare channel commit
-	amount := int64(1)
-	channelCommit, err := ledger.NewChannelCommit(payerPubKey, recvPubKey, amount)
-	if err != nil {
-		log.Panic("can not create channel commit", zap.Error(err))
-	}
-	// sign for the channel commit
-	fromSig, err := crypto.Sign(payerPrivKey, channelCommit)
-	if err != nil {
-		log.Panic("fail to sign channel commit", zap.Error(err))
-	}
-	// create channel: payer start the channel
-	channelID, err := ledger.CreateChannel(ctx, ledgerClient, channelCommit, fromSig)
-	if err != nil {
-		log.Panic("fail to create channel", zap.Error(err))
-	}
-	// channel state: transfer money from -> to
-	fromAcc, err := ledger.NewAccount(payerPubKey, 0)
-	if err != nil {
-		log.Panic("wrong account on channel", zap.Error(err))
-	}
-	toAcc, err := ledger.NewAccount(recvPubKey, amount)
-	if err != nil {
-		log.Panic("wrong account on channel", zap.Error(err))
-	}
-	channelState := ledger.NewChannelState(channelID, 1, fromAcc, toAcc)
-	// need permission from both account, get signature from both
-	fromSigState, err := crypto.Sign(payerPrivKey, channelState)
-	if err != nil {
-		log.Panic("error when signing the channel state", zap.Error(err))
-	}
-	toSigState, err := crypto.Sign(recvPrivKey, channelState)
-	if err != nil {
-		log.Panic("error when signing the channel state", zap.Error(err))
-	}
-	signedChannelState := ledger.NewSignedChannelState(channelState, fromSigState, toSigState)
-	// close channel
-	err = ledger.CloseChannel(ctx, ledgerClient, signedChannelState)
-	if err != nil {
-		log.Panic("fail to close channel", zap.Error(err))
+		log.Panic(err.Error())
 	}
 }

utils/grpc/client.go

@@ -14,6 +14,7 @@ import (
 	escrowpb "github.com/tron-us/go-btfs-common/protos/escrow"
 	guardpb "github.com/tron-us/go-btfs-common/protos/guard"
 	hubpb "github.com/tron-us/go-btfs-common/protos/hub"
+	ledgerpb "github.com/tron-us/go-btfs-common/protos/ledger"
 	sharedpb "github.com/tron-us/go-btfs-common/protos/shared"
 	statuspb "github.com/tron-us/go-btfs-common/protos/status"
 
@@ -61,6 +62,8 @@ func (g *ClientBuilder) doWithContext(ctx context.Context, f interface{}) error
 		return v(ctx, sharedpb.NewRuntimeServiceClient(conn))
 	case func(ctx context.Context, client grpc_health_v1.HealthClient) error:
 		return v(ctx, grpc_health_v1.NewHealthClient(conn))
+	case func(ctx context.Context, client ledgerpb.ChannelsClient) error:
+		return v(ctx, ledgerpb.NewChannelsClient(conn))
 	default:
 		return fmt.Errorf("illegal function: %T", f)
 	}

utils/grpc/ledger.go

@@ -0,0 +1,19 @@
+package grpc
+
+import (
+	"context"
+	ledgerpb "github.com/tron-us/go-btfs-common/protos/ledger"
+)
+
+func LedgerClient(addr string) *LedgerClientBuilder {
+	return &LedgerClientBuilder{builder(addr)}
+}
+
+type LedgerClientBuilder struct {
+	ClientBuilder
+}
+
+func (g *LedgerClientBuilder) WithContext(ctx context.Context, f func(ctx context.Context,
+	client ledgerpb.ChannelsClient) error) error {
+	return g.doWithContext(ctx, f)
+}