Allowed passing listener

wzshiming · Nov 22, 2022 · 60389ba · 60389ba
1 parent 9389638
commit 60389ba
Showing 1 changed file with 16 additions and 12 deletions.
diff --git a/simple_server.go b/simple_server.go
@@ -194,25 +194,29 @@ func serverConfig(addr string) (host, user, pwd string, config *ssh.ServerConfig
 // Run the server
 func (s *SimpleServer) Run(ctx context.Context) error {
 	var listenConfig net.ListenConfig
-	listener, err := listenConfig.Listen(ctx, s.Network, s.Address)
-	if err != nil {
-		return err
+	if s.Listener == nil {
+		listener, err := listenConfig.Listen(ctx, s.Network, s.Address)
+		if err != nil {
+			return err
+		}
+		s.Listener = listener
 	}
-	s.Listener = listener
-	s.Address = listener.Addr().String()
-	return s.Serve(listener)
+	s.Address = s.Listener.Addr().String()
+	return s.Serve(s.Listener)
 }
 
 // Start the server
 func (s *SimpleServer) Start(ctx context.Context) error {
 	var listenConfig net.ListenConfig
-	listener, err := listenConfig.Listen(ctx, s.Network, s.Address)
-	if err != nil {
-		return err
+	if s.Listener == nil {
+		listener, err := listenConfig.Listen(ctx, s.Network, s.Address)
+		if err != nil {
+			return err
+		}
+		s.Listener = listener
 	}
-	s.Listener = listener
-	s.Address = listener.Addr().String()
-	go s.Serve(listener)
+	s.Address = s.Listener.Addr().String()
+	go s.Serve(s.Listener)
 	return nil
 }