Make a little more resistant to exit when running as WS server

check.go

@@ -3,6 +3,7 @@ package main
 import (
 	"bufio"
 	"encoding/json"
+	"errors"
 	"log"
 	"os/exec"
 	"strings"
@@ -106,17 +107,20 @@ func parseupdatedata(uptext []string) *UpdateInfo {
 	return &details
 }
 
-func updatedryrun(sendupdate bool) (*UpdateInfo, bool) {
+func updatedryrun(sendupdate bool) (*UpdateInfo, bool, error) {
+	details := UpdateInfo{ }
+	updetails := &details
+
 	cmd := exec.Command(PKGBIN, "-C", localpkgconf, "upgrade", "-n")
 	sendinfomsg("Checking system for updates")
 	stdout, err := cmd.StdoutPipe()
 	if err != nil {
-		log.Println("Failed dry run")
-		log.Fatal(err)
+		sendfatalmsg("Failed dry run of pkg upgrade")
+		return updetails, false, errors.New("ERROR")
 	}
 	if err := cmd.Start(); err != nil {
-		log.Println("Failed starting dry run")
-		log.Fatal(err)
+		sendfatalmsg("Failed dry run of pkg upgrade")
+		return updetails, false, errors.New("ERROR")
 	}
 	buff := bufio.NewScanner(stdout)
 
@@ -132,13 +136,11 @@ func updatedryrun(sendupdate bool) (*UpdateInfo, bool) {
 	//}
 
 	haveupdates := ! strings.Contains(strings.Join((allText), "\n"), "Your packages are up to date")
-	details := UpdateInfo{ }
-	updetails := &details
 	if ( haveupdates ) {
 		updetails = parseupdatedata(allText)
 	}
 
-	return updetails, haveupdates
+	return updetails, haveupdates, nil
 }
 
 func sendupdatedetails(haveupdates bool, updetails *UpdateInfo) {
@@ -165,12 +167,14 @@ func sendupdatedetails(haveupdates bool, updetails *UpdateInfo) {
 func checkforupdates() {
 	preparepkgconfig()
 	updatepkgdb()
-	updetails, haveupdates := updatedryrun(true)
+	updetails, haveupdates, uerr:= updatedryrun(true)
+	if ( uerr != nil ) {
+                destroymddev()
+		return
+	}
 
         // If we are using standalone update, cleanup
-        if ( updatefileflag != "" ) {
-                destroymddev()
-        }
+	destroymddev()
 
 	sendupdatedetails(haveupdates, updetails)
 }

doupdate.go

@@ -79,7 +79,10 @@ func doupdate(message []byte) {
 
 	// Check that updates are available
 	logtofile("Checking for updates")
-	details, haveupdates := updatedryrun(false)
+	details, haveupdates, uerr := updatedryrun(false)
+	if ( uerr != nil ) {
+		return
+	}
 	if ( ! haveupdates ) {
 		sendfatalmsg("ERROR: No updates to install!")
 		return
@@ -110,7 +113,6 @@ func doupdate(message []byte) {
 	// Start the upgrade with bool passed if doing kernel update
 	startupgrade(twostageupdate)
 
-	//os.Exit(0)
 }
 
 func cleanupbe() {
@@ -337,7 +339,7 @@ func startupgrade(twostage bool) {
 		cmd := exec.Command("umount", "-f", STAGEDIR + localimgmnt)
 		err := cmd.Run()
 		if ( err != nil ) {
-			log.Fatal(err)
+			log.Println("WARNING: Failed to umount " + STAGEDIR + localimgmnt)
 		}
 	}
 
@@ -350,9 +352,7 @@ func startupgrade(twostage bool) {
 
 	// If we are using standalone update, cleanup
 	destroymddev()
-	sendinfomsg("Success! Reboot your system to continue the update process.")
-	time.Sleep(500 * time.Millisecond)
-	os.Exit(0)
+	sendshutdownmsg("Success! Reboot your system to continue the update process.")
 }
 
 func renamebe() {

main.go

@@ -65,7 +65,11 @@ func connectws() {
 			connected = true
 			break
 		}
-		time.Sleep(5 * time.Millisecond);
+		time.Sleep(100 * time.Millisecond);
+		// If we fail first connect, lets fire up the internal server
+		if ( attempt == 1 ) {
+			go startws()
+		}
 	}
 	if (!connected) {
 		log.Fatal("Failed connecting to websocket server", err)
@@ -84,6 +88,7 @@ func closews() {
 		log.Println("write close:", err)
 		return
 	}
+	time.Sleep(100 * time.Millisecond);
 }
 
 func checkuid() {
@@ -115,31 +120,27 @@ func main() {
 	loadconfig()
 
 	if ( listtrainflag ) {
-		go startws()
 		connectws()
 		listtrains()
 		closews()
 		os.Exit(0)
 	}
 
 	if ( changetrainflag != "" ) {
-		go startws()
 		connectws()
 		settrain()
 		closews()
 		os.Exit(0)
 	}
 
 	if ( checkflag ) {
-		go startws()
 		connectws()
 		startcheck()
 		closews()
 		os.Exit(0)
 	}
 
 	if ( updateflag ) {
-		go startws()
 		connectws()
 		startupdate()
 		closews()

ws.go

@@ -31,7 +31,9 @@ func readws(w http.ResponseWriter, r *http.Request) {
 		// Start decoding the incoming JSON
 	        var env Envelope
 		if err := json.Unmarshal(message, &env); err != nil {
-			log.Fatal(err)
+			sendinfomsg("Invalid JSON received")
+			log.Println("Warning: Invalid JSON message received")
+			log.Println(err)
 	        }
 	        switch env.Method {
 	        case "check":