From c3bec5f80c626a314d6e0773145586a65edd61f3 Mon Sep 17 00:00:00 2001
From: kiwiidb <kwinten.debacker@hotmail.com>
Date: Fri, 4 Feb 2022 12:13:40 +0100
Subject: [PATCH] chore: wrap LND client

---
 lib/service/invoicesubscription.go |  3 ++-
 lib/service/service.go             |  4 ++--
 lnd/interface.go                   | 20 ++++++++++++++++++++
 lnd/lnd.go                         | 30 ++++++++++++++++++++++++++++--
 4 files changed, 52 insertions(+), 5 deletions(-)
 create mode 100644 lnd/interface.go

diff --git a/lib/service/invoicesubscription.go b/lib/service/invoicesubscription.go
index 2cc06f1..cdf62f6 100644
--- a/lib/service/invoicesubscription.go
+++ b/lib/service/invoicesubscription.go
@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/getAlby/lndhub.go/db/models"
+	"github.com/getAlby/lndhub.go/lnd"
 	"github.com/getsentry/sentry-go"
 	"github.com/lightningnetwork/lnd/lnrpc"
 	"github.com/uptrace/bun"
@@ -95,7 +96,7 @@ func (svc *LndhubService) ProcessInvoiceUpdate(ctx context.Context, rawInvoice *
 	return nil
 }
 
-func (svc *LndhubService) ConnectInvoiceSubscription(ctx context.Context) (lnrpc.Lightning_SubscribeInvoicesClient, error) {
+func (svc *LndhubService) ConnectInvoiceSubscription(ctx context.Context) (lnd.SubscribeInvoicesWrapper, error) {
 	var invoice models.Invoice
 	invoiceSubscriptionOptions := lnrpc.InvoiceSubscription{}
 	// Find the oldest NOT settled invoice with an add_index
diff --git a/lib/service/service.go b/lib/service/service.go
index 4af9a3d..33dd0fe 100644
--- a/lib/service/service.go
+++ b/lib/service/service.go
@@ -8,8 +8,8 @@ import (
 	"github.com/btcsuite/btcd/btcec"
 	"github.com/getAlby/lndhub.go/db/models"
 	"github.com/getAlby/lndhub.go/lib/tokens"
+	"github.com/getAlby/lndhub.go/lnd"
 	"github.com/labstack/gommon/random"
-	"github.com/lightningnetwork/lnd/lnrpc"
 	"github.com/uptrace/bun"
 	"github.com/ziflex/lecho/v3"
 	"golang.org/x/crypto/bcrypt"
@@ -20,7 +20,7 @@ const alphaNumBytes = random.Alphanumeric
 type LndhubService struct {
 	Config         *Config
 	DB             *bun.DB
-	LndClient      lnrpc.LightningClient
+	LndClient      lnd.LightningClientWrapper
 	Logger         *lecho.Logger
 	IdentityPubkey *btcec.PublicKey
 }
diff --git a/lnd/interface.go b/lnd/interface.go
new file mode 100644
index 0000000..c5a39a7
--- /dev/null
+++ b/lnd/interface.go
@@ -0,0 +1,20 @@
+package lnd
+
+import (
+	"context"
+
+	"github.com/lightningnetwork/lnd/lnrpc"
+	"google.golang.org/grpc"
+)
+
+type LightningClientWrapper interface {
+	ListChannels(ctx context.Context, req *lnrpc.ListChannelsRequest, options ...grpc.CallOption) (*lnrpc.ListChannelsResponse, error)
+	SendPaymentSync(ctx context.Context, req *lnrpc.SendRequest, options ...grpc.CallOption) (*lnrpc.SendResponse, error)
+	AddInvoice(ctx context.Context, req *lnrpc.Invoice, options ...grpc.CallOption) (*lnrpc.AddInvoiceResponse, error)
+	SubscribeInvoices(ctx context.Context, req *lnrpc.InvoiceSubscription, options ...grpc.CallOption) (SubscribeInvoicesWrapper, error)
+	GetInfo(ctx context.Context, req *lnrpc.GetInfoRequest, options ...grpc.CallOption) (*lnrpc.GetInfoResponse, error)
+}
+
+type SubscribeInvoicesWrapper interface {
+	Recv() (*lnrpc.Invoice, error)
+}
diff --git a/lnd/lnd.go b/lnd/lnd.go
index 2b271ad..d62bd6e 100644
--- a/lnd/lnd.go
+++ b/lnd/lnd.go
@@ -1,6 +1,7 @@
 package lnd
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"encoding/hex"
@@ -23,8 +24,11 @@ type LNDoptions struct {
 	MacaroonHex  string
 }
 
-func NewLNDclient(lndOptions LNDoptions) (lnrpc.LightningClient, error) {
+type LNDWrapper struct {
+	client lnrpc.LightningClient
+}
 
+func NewLNDclient(lndOptions LNDoptions) (result *LNDWrapper, err error) {
 	// Get credentials either from a hex string, a file or the system's certificate store
 	var creds credentials.TransportCredentials
 	// if a hex string is provided
@@ -82,5 +86,27 @@ func NewLNDclient(lndOptions LNDoptions) (lnrpc.LightningClient, error) {
 		return nil, err
 	}
 
-	return lnrpc.NewLightningClient(conn), nil
+	return &LNDWrapper{
+		client: lnrpc.NewLightningClient(conn),
+	}, nil
+}
+
+func (wrapper *LNDWrapper) ListChannels(ctx context.Context, req *lnrpc.ListChannelsRequest, options ...grpc.CallOption) (*lnrpc.ListChannelsResponse, error) {
+	return wrapper.client.ListChannels(ctx, req, options...)
+}
+
+func (wrapper *LNDWrapper) SendPaymentSync(ctx context.Context, req *lnrpc.SendRequest, options ...grpc.CallOption) (*lnrpc.SendResponse, error) {
+	return wrapper.client.SendPaymentSync(ctx, req, options...)
+}
+
+func (wrapper *LNDWrapper) AddInvoice(ctx context.Context, req *lnrpc.Invoice, options ...grpc.CallOption) (*lnrpc.AddInvoiceResponse, error) {
+	return wrapper.client.AddInvoice(ctx, req, options...)
+}
+
+func (wrapper *LNDWrapper) SubscribeInvoices(ctx context.Context, req *lnrpc.InvoiceSubscription, options ...grpc.CallOption) (SubscribeInvoicesWrapper, error) {
+	return wrapper.client.SubscribeInvoices(ctx, req, options...)
+}
+
+func (wrapper *LNDWrapper) GetInfo(ctx context.Context, req *lnrpc.GetInfoRequest, options ...grpc.CallOption) (*lnrpc.GetInfoResponse, error) {
+	return wrapper.client.GetInfo(ctx, req, options...)
 }