Fix upgrade bug. Add time remaining in upgrading. Close #23

cf.go

@@ -12,7 +12,7 @@ import (
 	docopt "github.com/docopt/docopt-go"
 )
 
-const version = "v0.7.0"
+const version = "v0.7.1"
 
 func main() {
 	usage := `Codeforces Tool $%version%$ (cf). https://github.com/xalanq/cf-tool

cmd/upgrade.go

@@ -97,7 +97,8 @@ type WriteCounter struct {
 
 // Print print progress
 func (w *WriteCounter) Print() {
-	fmt.Printf("\rProgress: %v/%v KB  Speed: %v KB/s", w.Count/1024, w.Total/1024, (w.Count-w.last)/1024)
+	fmt.Printf("\rProgress: %v/%v KB  Speed: %v KB/s  Remain: %.0f s           ",
+		w.Count/1024, w.Total/1024, (w.Count-w.last)/1024, float64(w.Total-w.Count)/float64(w.Count-w.last))
 	w.last = w.Count
 }
 
@@ -107,7 +108,28 @@ func (w *WriteCounter) Write(p []byte) (int, error) {
 	return n, nil
 }
 
-func upgrade(url, exe string, size uint) (err error) {
+func upgrade(url, exePath string, size uint) (err error) {
+	updateDir := filepath.Dir(exePath)
+
+	oldPath := filepath.Join(updateDir, fmt.Sprintf(".%s.old", filepath.Base(exePath)))
+	color.Cyan("Move the old one to %v", oldPath)
+	if err = os.Rename(exePath, oldPath); err != nil {
+		return
+	}
+	defer func() {
+		if err != nil {
+			color.Cyan("Move the old one back")
+			if e := os.Rename(oldPath, exePath); e != nil {
+				color.Red(e.Error())
+			}
+		} else {
+			color.Cyan("Remove the old one")
+			if e := os.Remove(oldPath); e != nil {
+				color.Red(e.Error() + "\nYou could remove it manually")
+			}
+		}
+	}()
+
 	color.Cyan("Download %v", url)
 	counter := &WriteCounter{Count: 0, Total: size, last: 0}
 	counter.Print()
@@ -126,56 +148,41 @@ func upgrade(url, exe string, size uint) (err error) {
 	defer resp.Body.Close()
 
 	data, err := ioutil.ReadAll(io.TeeReader(resp.Body, counter))
-	if err != nil {
-		ticker.Stop()
-		counter.Print()
-		fmt.Println()
-		return
-	}
 	ticker.Stop()
 	counter.Print()
 	fmt.Println()
-	reader, err := zip.NewReader(bytes.NewReader(data), int64(len(data)))
 	if err != nil {
 		return
 	}
-
-	rc, err := reader.File[0].Open()
+	reader, err := zip.NewReader(bytes.NewReader(data), int64(len(data)))
 	if err != nil {
 		return
 	}
-	defer rc.Close()
-
-	newPath := filepath.Join(os.TempDir(), fmt.Sprintf(".%s.new", filepath.Base(exe)))
-	oldPath := filepath.Join(os.TempDir(), fmt.Sprintf(".%s.old", filepath.Base(exe)))
 
-	file, err := os.OpenFile(newPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
+	rc, err := reader.File[0].Open()
 	if err != nil {
 		return
 	}
-	defer file.Close()
-
-	_, err = io.Copy(file, rc)
+	newData, err := ioutil.ReadAll(rc)
+	rc.Close()
 	if err != nil {
 		return
 	}
-	file.Close()
 
-	err = os.Rename(exe, oldPath)
-	if err != nil {
-		os.Remove(newPath)
+	newPath := filepath.Join(updateDir, fmt.Sprintf(".%s.new", filepath.Base(exePath)))
+	color.Cyan("Save the new one to %v", newPath)
+	if err = ioutil.WriteFile(newPath, newData, 0755); err != nil {
 		return
 	}
 
-	err = os.Rename(newPath, exe)
-	if err != nil {
-		os.Rename(oldPath, exe)
-		os.Remove(newPath)
-		return
+	if err = os.Rename(newPath, exePath); err != nil {
+		color.Cyan("Delete the new one %v", newPath)
+		if e := os.Remove(newPath); e != nil {
+			color.Red(e.Error())
+		}
 	}
 
-	os.Remove(oldPath)
-	return nil
+	return
 }
 
 // Upgrade itself
@@ -198,16 +205,16 @@ func Upgrade(version string) error {
 		return nil
 	}
 
-	exe, err := os.Executable()
+	exePath, err := os.Executable()
 	if err != nil {
 		return err
 	}
 
-	if exe, err = filepath.EvalSymlinks(exe); err != nil {
+	if exePath, err = filepath.EvalSymlinks(exePath); err != nil {
 		return err
 	}
 
-	if err = upgrade(url, exe, size); err != nil {
+	if err = upgrade(url, exePath, size); err != nil {
 		return err
 	}