From 0f1b7741d7758d9e014c7a627f65e82c1adc42af Mon Sep 17 00:00:00 2001
From: Darren Shepherd <darren@rancher.com>
Date: Fri, 28 Aug 2015 21:13:01 -0700
Subject: [PATCH] Fix setting up state partition

---
 init/root.go | 59 +++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 51 insertions(+), 8 deletions(-)

diff --git a/init/root.go b/init/root.go
index 86aad6fb..8815dd26 100644
--- a/init/root.go
+++ b/init/root.go
@@ -5,6 +5,7 @@ import (
 	"io/ioutil"
 	"os"
 	"path"
+	"strings"
 	"syscall"
 
 	log "github.com/Sirupsen/logrus"
@@ -13,25 +14,67 @@ import (
 	"github.com/rancherio/os/config"
 )
 
-func prepareRoot(rootfs string) error {
-	usr := path.Join(rootfs, "usr")
-	if err := os.Remove(usr); err != nil && !os.IsNotExist(err) {
-		log.Errorf("Failed to delete %s, possibly invalid RancherOS state partition: %v", usr, err)
-		return err
+func cleanupTarget(rootfs, targetUsr, usr, usrVer, tmpDir string) (bool, error) {
+	log.Debugf("Deleting %s", targetUsr)
+	if err := os.Remove(targetUsr); err != nil && !os.IsNotExist(err) {
+		log.Errorf("Failed to delete %s, possibly invalid RancherOS state partition: %v", targetUsr, err)
+		return false, err
 	}
 
-	return nil
+	if err := dockerlaunch.CreateSymlink(usrVer, path.Join(rootfs, "usr")); err != nil {
+		return false, err
+	}
+
+	log.Debugf("Deleting %s", tmpDir)
+	if err := os.RemoveAll(tmpDir); err != nil {
+		// Don't care if this fails
+		log.Errorf("Failed to cleanup temp directory %s: %v", tmpDir, err)
+	}
+
+	if strings.HasSuffix(usrVer, "dev") {
+		log.Debugf("Deleteing old usr: %s", usr)
+		if err := os.RemoveAll(usr); err != nil {
+			// Don't care if this fails
+			log.Errorf("Failed to remove %s: %v", usr, err)
+		}
+		return true, nil
+	}
+
+	if _, err := os.Stat(usrVer); os.IsNotExist(err) {
+		return false, nil
+	}
+
+	return true, nil
 }
 
 func copyMoveRoot(rootfs string) error {
 	usrVer := fmt.Sprintf("usr-%s", config.VERSION)
 	usr := path.Join(rootfs, usrVer)
+	targetUsr := path.Join(rootfs, "usr")
+	tmpDir := path.Join(rootfs, "tmp")
 
-	if err := archive.CopyWithTar("/usr", usr); err != nil {
+	if cont, err := cleanupTarget(rootfs, targetUsr, usr, usrVer, tmpDir); !cont {
 		return err
 	}
 
-	if err := dockerlaunch.CreateSymlink(usrVer, path.Join(rootfs, "usr")); err != nil {
+	log.Debugf("Creating temp dir directory %s", tmpDir)
+	if err := os.MkdirAll(tmpDir, 0755); err != nil {
+		return err
+	}
+
+	usrVerTmp, err := ioutil.TempDir(tmpDir, usrVer)
+	if err != nil {
+		return err
+	}
+
+	log.Debugf("Copying to temp dir %s", usrVerTmp)
+
+	if err := archive.CopyWithTar("/usr", usrVerTmp); err != nil {
+		return err
+	}
+
+	log.Debugf("Renaming %s => %s", usrVerTmp, usr)
+	if err := os.Rename(usrVerTmp, usr); err != nil {
 		return err
 	}