From fa21601d070d633c9b609da09ecf93eb2402afb7 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Tue, 9 Jul 2019 14:04:51 +0200
Subject: [PATCH] lnd+cmd/lnd/main: add ListenerCfg to Main

ListenerCfg allows passing custom listeners to the main method, to be
used for the wallet unlocker and rpc server. If these are set these will
be used instead of the regular RPC listeners.
---
 cmd/lnd/main.go |  2 +-
 lnd.go          | 36 +++++++++++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/cmd/lnd/main.go b/cmd/lnd/main.go
index 177fb67d..3448b77a 100644
--- a/cmd/lnd/main.go
+++ b/cmd/lnd/main.go
@@ -11,7 +11,7 @@ import (
 func main() {
 	// Call the "real" main in a nested manner so the defers will properly
 	// be executed in the case of a graceful shutdown.
-	if err := lnd.Main(); err != nil {
+	if err := lnd.Main(lnd.ListenerCfg{}); err != nil {
 		if e, ok := err.(*flags.Error); ok && e.Type == flags.ErrHelp {
 		} else {
 			fmt.Fprintln(os.Stderr, err)
diff --git a/lnd.go b/lnd.go
index d550407a..982e672b 100644
--- a/lnd.go
+++ b/lnd.go
@@ -93,6 +93,18 @@ var (
 	}
 )
 
+// ListenerCfg is a wrapper around custom listeners that can be passed to lnd
+// when calling its main method.
+type ListenerCfg struct {
+	// WalletUnlocker can be set to the listener to use for the wallet
+	// unlocker. If nil a regular network listener will be created.
+	WalletUnlocker net.Listener
+
+	// RPCListener can be set to the listener to use for the RPC server. If
+	// nil a regular network listener will be created.
+	RPCListener net.Listener
+}
+
 // rpcListeners is a function type used for closures that fetches a set of RPC
 // listeners for the current configuration, and the GRPC server options to use
 // with these listeners. If no custom listeners are present, this should return
@@ -103,7 +115,7 @@ type rpcListeners func() ([]net.Listener, func(), []grpc.ServerOption, error)
 // Main is the true entry point for lnd. This function is required since defers
 // created in the top-level scope of a main method aren't executed if os.Exit()
 // is called.
-func Main() error {
+func Main(lisCfg ListenerCfg) error {
 	// Load the configuration, and parse any command line options. This
 	// function will also set up logging properly.
 	loadedConfig, err := loadConfig()
@@ -280,6 +292,17 @@ func Main() error {
 	walletUnlockerListeners := func() ([]net.Listener, func(),
 		[]grpc.ServerOption, error) {
 
+		// If we have chosen to start with a dedicated listener for the
+		// wallet unlocker, we return it directly, and empty server
+		// options to deactivate TLS.
+		// TODO(halseth): any point in adding TLS support for custom
+		// listeners?
+		if lisCfg.WalletUnlocker != nil {
+			return []net.Listener{lisCfg.WalletUnlocker}, func() {},
+				[]grpc.ServerOption{}, nil
+		}
+
+		// Otherwise we'll return the regular listeners.
 		return getListeners()
 	}
 
@@ -505,6 +528,17 @@ func Main() error {
 	rpcListeners := func() ([]net.Listener, func(), []grpc.ServerOption,
 		error) {
 
+		// If we have chosen to start with a dedicated listener for the
+		// rpc server, we return it directly, and empty server options
+		// to deactivate TLS.
+		// TODO(halseth): any point in adding TLS support for custom
+		// listeners?
+		if lisCfg.RPCListener != nil {
+			return []net.Listener{lisCfg.RPCListener}, func() {},
+				[]grpc.ServerOption{}, nil
+		}
+
+		// Otherwise we'll return the regular listeners.
 		return getListeners()
 	}