whisper: refactoring go-routines

1 files changed, 32 insertions, 24 deletions

diff --git a/cmd/wnode/main.go b/cmd/wnode/main.go
index 00a854fe8..84bdfa4c3 100644
--- a/cmd/wnode/main.go
+++ b/cmd/wnode/main.go
@@ -110,6 +110,7 @@ func main() {
 	processArgs()
 	initialize()
 	run()
+	shutdown()
 }
 
 func processArgs() {
@@ -209,21 +210,6 @@ func initialize() {
 		MinimumAcceptedPOW: *argPoW,
 	}
 
-	if *mailServerMode {
-		if len(msPassword) == 0 {
-			msPassword, err = console.Stdin.PromptPassword("Please enter the Mail Server password: ")
-			if err != nil {
-				utils.Fatalf("Failed to read Mail Server password: %s", err)
-			}
-		}
-
-		shh = whisper.New(cfg)
-		shh.RegisterServer(&mailServer)
-		mailServer.Init(shh, *argDBPath, msPassword, *argServerPoW)
-	} else {
-		shh = whisper.New(cfg)
-	}
-
 	if *argPoW != whisper.DefaultMinimumPoW {
 		err := shh.SetMinimumPoW(*argPoW)
 		if err != nil {
@@ -265,6 +251,26 @@ func initialize() {
 		maxPeers = 800
 	}
 
+	_, err = crand.Read(entropy[:])
+	if err != nil {
+		utils.Fatalf("crypto/rand failed: %s", err)
+	}
+
+	if *mailServerMode {
+		if len(msPassword) == 0 {
+			msPassword, err = console.Stdin.PromptPassword("Please enter the Mail Server password: ")
+			if err != nil {
+				utils.Fatalf("Failed to read Mail Server password: %s", err)
+			}
+		}
+
+		shh = whisper.New(cfg)
+		shh.RegisterServer(&mailServer)
+		mailServer.Init(shh, *argDBPath, msPassword, *argServerPoW)
+	} else {
+		shh = whisper.New(cfg)
+	}
+
 	server = &p2p.Server{
 		Config: p2p.Config{
 			PrivateKey:     nodeid,
@@ -278,17 +284,13 @@ func initialize() {
 			TrustedNodes:   peers,
 		},
 	}
-
-	_, err = crand.Read(entropy[:])
-	if err != nil {
-		utils.Fatalf("crypto/rand failed: %s", err)
-	}
 }
 
-func startServer() {
+func startServer() error {
 	err := server.Start()
 	if err != nil {
-		utils.Fatalf("Failed to start Whisper peer: %s.", err)
+		fmt.Printf("Failed to start Whisper peer: %s.", err)
+		return err
 	}
 
 	fmt.Printf("my public key: %s \n", common.ToHex(crypto.FromECDSAPub(&asymKey.PublicKey)))
@@ -307,6 +309,7 @@ func startServer() {
 	if !*forwarderMode {
 		fmt.Printf("Please type the message. To quit type: '%s'\n", quitCommand)
 	}
+	return nil
 }
 
 func isKeyValid(k *ecdsa.PublicKey) bool {
@@ -420,8 +423,10 @@ func waitForConnection(timeout bool) {
 }
 
 func run() {
-	defer mailServer.Close()
-	startServer()
+	err := startServer()
+	if err != nil {
+		return
+	}
 	defer server.Stop()
 	shh.Start(nil)
 	defer shh.Stop()
@@ -439,8 +444,11 @@ func run() {
 	} else {
 		sendLoop()
 	}
+}
 
+func shutdown() {
 	close(done)
+	mailServer.Close()
 }
 
 func sendLoop() {