From dadbdcbbc025bd35edfcdeab43a0c86b46b62bfd Mon Sep 17 00:00:00 2001
From: Sir Darkrengarius <darkren92@gmail.com>
Date: Sun, 2 Jun 2019 20:03:54 +0300
Subject: [PATCH] Move `net.Listener` arg to `NewServer`, remove `addr` arg

---
 cmd/messaging-server/commands/root.go | 10 +++++---
 pkg/dmsg/server.go                    | 37 +++++++++++++--------------
 pkg/dmsg/server_test.go               | 19 +++++++++-----
 3 files changed, 37 insertions(+), 29 deletions(-)

diff --git a/cmd/messaging-server/commands/root.go b/cmd/messaging-server/commands/root.go
index b10d2b638..a84348a52 100644
--- a/cmd/messaging-server/commands/root.go
+++ b/cmd/messaging-server/commands/root.go
@@ -74,12 +74,16 @@ var rootCmd = &cobra.Command{
 
 		l, err := net.Listen("tcp", conf.LocalAddress)
 		if err != nil {
-			log.Fatal(err)
+			logger.Fatalf("Error listening on %s: %v", conf.LocalAddress, err)
 		}
 
 		// Start
-		srv := dmsg.NewServer(conf.PubKey, conf.SecKey, conf.PublicAddress, client.NewHTTP(conf.Discovery))
-		log.Fatal(srv.Serve(l))
+		srv, err := dmsg.NewServer(conf.PubKey, conf.SecKey, l, client.NewHTTP(conf.Discovery))
+		if err != nil {
+			logger.Fatalf("Error creating DMSG server instance: %v", err)
+		}
+
+		log.Fatal(srv.Serve())
 	},
 }
 
diff --git a/pkg/dmsg/server.go b/pkg/dmsg/server.go
index 11ac01573..f74f848dd 100644
--- a/pkg/dmsg/server.go
+++ b/pkg/dmsg/server.go
@@ -16,6 +16,8 @@ import (
 	"github.com/skycoin/skywire/pkg/messaging-discovery/client"
 )
 
+var ErrListenerAlreadyWrappedToNoise = errors.New("listener is already wrapped to *noise.Listener")
+
 // NextConn provides information on the next connection.
 type NextConn struct {
 	l  *ServerConn
@@ -201,10 +203,9 @@ func (c *ServerConn) handleRequest(ctx context.Context, getLink getConnFunc, id
 type Server struct {
 	log *logging.Logger
 
-	pk   cipher.PubKey
-	sk   cipher.SecKey
-	addr string
-	dc   client.APIClient
+	pk cipher.PubKey
+	sk cipher.SecKey
+	dc client.APIClient
 
 	lis   net.Listener
 	conns map[cipher.PubKey]*ServerConn
@@ -214,15 +215,19 @@ type Server struct {
 }
 
 // NewServer creates a new dms_server.
-func NewServer(pk cipher.PubKey, sk cipher.SecKey, addr string, dc client.APIClient) *Server {
+func NewServer(pk cipher.PubKey, sk cipher.SecKey, l net.Listener, dc client.APIClient) (*Server, error) {
+	if _, ok := l.(*noise.Listener); ok {
+		return nil, ErrListenerAlreadyWrappedToNoise
+	}
+
 	return &Server{
 		log:   logging.MustGetLogger("dms_server"),
 		pk:    pk,
 		sk:    sk,
-		addr:  addr,
+		lis:   noise.WrapListener(l, pk, sk, false, noise.HandshakeXK),
 		dc:    dc,
 		conns: make(map[cipher.PubKey]*ServerConn),
-	}
+	}, nil
 }
 
 // SetLogger set's the logger.
@@ -269,24 +274,18 @@ func (s *Server) Close() (err error) {
 }
 
 // ListenAndServe serves the dms_server.
-//func (s *Server) ListenAndServe(addr string) error {
-func (s *Server) Serve(l net.Listener) error {
+func (s *Server) Serve() error {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	/*lis, err := net.Listen("tcp", addr)
-	if err != nil {
-		return err
-	}*/
 	if err := s.updateDiscEntry(ctx); err != nil {
 		return fmt.Errorf("updating server's discovery entry failed with: %s", err)
 	}
 
-	s.log.Infof("serving: pk(%s) addr(%s)", s.pk, l.Addr())
-	l = noise.WrapListener(l, s.pk, s.sk, false, noise.HandshakeXK)
-	s.lis = l
+	s.log.Infof("serving: pk(%s) addr(%s)", s.pk, s.lis.Addr())
+
 	for {
-		rawConn, err := l.Accept()
+		rawConn, err := s.lis.Accept()
 		if err != nil {
 			if err == io.ErrUnexpectedEOF {
 				continue
@@ -311,12 +310,12 @@ func (s *Server) updateDiscEntry(ctx context.Context) error {
 	s.log.Info("updating server discovery entry...")
 	entry, err := s.dc.Entry(ctx, s.pk)
 	if err != nil {
-		entry = client.NewServerEntry(s.pk, 0, s.addr, 10)
+		entry = client.NewServerEntry(s.pk, 0, s.lis.Addr().String(), 10)
 		if err := entry.Sign(s.sk); err != nil {
 			return err
 		}
 		return s.dc.SetEntry(ctx, entry)
 	}
-	entry.Server.Address = s.addr
+	entry.Server.Address = s.lis.Addr().String()
 	return s.dc.UpdateEntry(ctx, s.sk, entry)
 }
diff --git a/pkg/dmsg/server_test.go b/pkg/dmsg/server_test.go
index f82f5a912..43a3d177b 100644
--- a/pkg/dmsg/server_test.go
+++ b/pkg/dmsg/server_test.go
@@ -26,8 +26,10 @@ func TestNewServer(t *testing.T) {
 	l, err := net.Listen("tcp", "")
 	require.NoError(t, err)
 
-	s := NewServer(sPK, sSK, "", dc)
-	go s.Serve(l) //nolint:errcheck
+	s, err := NewServer(sPK, sSK, l, dc)
+	require.NoError(t, err)
+
+	go s.Serve() //nolint:errcheck
 
 	time.Sleep(time.Second)
 
@@ -41,9 +43,10 @@ func TestServer_ListenAndServe(t *testing.T) {
 	l, err := net.Listen("tcp", "")
 	require.NoError(t, err)
 
-	s := NewServer(sPK, sSK, "", dc)
-	go s.Serve(l)
+	s, err := NewServer(sPK, sSK, l, dc)
+	require.NoError(t, err)
 
+	go s.Serve()
 }
 
 // Given two client instances (a & b) and a server instance (s),
@@ -61,11 +64,13 @@ func TestNewClient(t *testing.T) {
 
 	dc := client.NewMock()
 
-	l, err := net.Listen("tcp", "")
+	l, err := net.Listen("tcp", sAddr)
+	require.NoError(t, err)
+
+	s, err := NewServer(sPK, sSK, l, dc)
 	require.NoError(t, err)
 
-	s := NewServer(sPK, sSK, sAddr, dc)
-	go s.Serve(l) //nolint:errcheck
+	go s.Serve() //nolint:errcheck
 
 	a := NewClient(aPK, aSK, dc)
 	a.SetLogger(logging.MustGetLogger("A"))