From 9138d58152d870549fb04391e6ea9891bff2f51b Mon Sep 17 00:00:00 2001
From: Richard Ramos <info@richardramos.me>
Date: Wed, 28 Apr 2021 16:23:03 -0400
Subject: [PATCH] Add option to enable lightpush protocol

---
 waku/v2/node/wakunode2.go                    |  8 ++++++++
 waku/v2/node/wakuoptions.go                  | 10 ++++++++++
 waku/v2/protocol/lightpush/waku_lightpush.go |  9 ++++-----
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/waku/v2/node/wakunode2.go b/waku/v2/node/wakunode2.go
index 76007477..c19e8391 100644
--- a/waku/v2/node/wakunode2.go
+++ b/waku/v2/node/wakunode2.go
@@ -95,6 +95,10 @@ func New(ctx context.Context, opts ...WakuNodeOption) (*WakuNode, error) {
 		}
 	}
 
+	if params.enableLightPush {
+		w.mountLightPush()
+	}
+
 	for _, addr := range w.ListenAddresses() {
 		log.Info("Listening on ", addr)
 	}
@@ -147,6 +151,10 @@ func (w *WakuNode) mountRelay(opts ...wakurelay.Option) error {
 	return err
 }
 
+func (w *WakuNode) mountLightPush() {
+	w.lightPush = lightpush.NewWakuLightPush(w.ctx, w.host, w.relay)
+}
+
 func (w *WakuNode) startStore() error {
 	_, err := w.Subscribe(nil)
 	if err != nil {
diff --git a/waku/v2/node/wakuoptions.go b/waku/v2/node/wakuoptions.go
index 5fe3b9d7..f951e106 100644
--- a/waku/v2/node/wakuoptions.go
+++ b/waku/v2/node/wakuoptions.go
@@ -24,6 +24,8 @@ type WakuNodeParameters struct {
 	enableStore bool
 	storeMsgs   bool
 	store       *store.WakuStore
+
+	enableLightPush bool
 }
 
 type WakuNodeOption func(*WakuNodeParameters) error
@@ -107,6 +109,14 @@ func WithMessageProvider(s store.MessageProvider) WakuNodeOption {
 	}
 }
 
+// WithLightPush is a WakuNodeOption that enables the lightpush protocol
+func WithLightPush() WakuNodeOption {
+	return func(params *WakuNodeParameters) error {
+		params.enableLightPush = true
+		return nil
+	}
+}
+
 // Default options used in the libp2p node
 var DefaultLibP2POptions = []libp2p.Option{
 	libp2p.DefaultTransports,
diff --git a/waku/v2/protocol/lightpush/waku_lightpush.go b/waku/v2/protocol/lightpush/waku_lightpush.go
index 298eaf20..eabb4259 100644
--- a/waku/v2/protocol/lightpush/waku_lightpush.go
+++ b/waku/v2/protocol/lightpush/waku_lightpush.go
@@ -36,17 +36,16 @@ type WakuLightPush struct {
 	ctx   context.Context
 }
 
-func NewWakuLightPush(ctx context.Context, relay *relay.WakuRelay) *WakuLightPush {
+func NewWakuLightPush(ctx context.Context, h host.Host, relay *relay.WakuRelay) *WakuLightPush {
 	wakuLP := new(WakuLightPush)
 	wakuLP.relay = relay
 	wakuLP.ctx = ctx
-	return wakuLP
-}
-
-func (wakuLP *WakuLightPush) Start(h host.Host) {
 	wakuLP.h = h
+
 	wakuLP.h.SetStreamHandler(WakuLightPushProtocolId, wakuLP.onRequest)
 	log.Info("Light Push protocol started")
+
+	return wakuLP
 }
 
 func (wakuLP *WakuLightPush) onRequest(s network.Stream) {