From 36ab79eb01e0f2db9f707e25a6669f77867c96f1 Mon Sep 17 00:00:00 2001
From: Wyatt Daviau <wdaviau@cs.stanford.edu>
Date: Fri, 26 Jan 2018 18:21:30 -0500
Subject: [PATCH] upgrade skips upgrading from current state format

License: MIT
Signed-off-by: Wyatt Daviau <wdaviau@cs.stanford.edu>
---
 ipfs-cluster-service/state.go | 39 ++++++++++++++++++++++++++++-------
 state/mapstate/map_state.go   |  7 +++++++
 2 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/ipfs-cluster-service/state.go b/ipfs-cluster-service/state.go
index bfbbd8a9..331a0f6b 100644
--- a/ipfs-cluster-service/state.go
+++ b/ipfs-cluster-service/state.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"encoding/json"
 	"errors"
 	"io"
@@ -15,11 +16,16 @@ import (
 var errNoSnapshot = errors.New("no snapshot found")
 
 func upgrade() error {
-	newState, err := restoreStateFromDisk()
+	newState, current, err := restoreStateFromDisk()
 	if err != nil {
 		return err
 	}
 
+	if current {
+		logger.Warning("Skipping migration of up-to-date state")
+		return nil
+	}
+
 	cfg, clusterCfg, _, _, consensusCfg, _, _, _, _ := makeConfigs()
 
 	err = cfg.LoadJSONFromFile(configPath)
@@ -32,7 +38,7 @@ func upgrade() error {
 }
 
 func export(w io.Writer) error {
-	stateToExport, err := restoreStateFromDisk()
+	stateToExport, _, err := restoreStateFromDisk()
 	if err != nil {
 		return err
 	}
@@ -40,12 +46,15 @@ func export(w io.Writer) error {
 	return exportState(stateToExport, w)
 }
 
-func restoreStateFromDisk() (*mapstate.MapState, error) {
+// restoreStateFromDisk returns a mapstate containing the latest
+// snapshot, a flag set to true when the state format has the
+// current version and an error
+func restoreStateFromDisk() (*mapstate.MapState, bool, error) {
 	cfg, _, _, _, consensusCfg, _, _, _, _ := makeConfigs()
 
 	err := cfg.LoadJSONFromFile(configPath)
 	if err != nil {
-		return nil, err
+		return nil, false, err
 	}
 
 	r, snapExists, err := raft.LastStateRaw(consensusCfg)
@@ -53,17 +62,31 @@ func restoreStateFromDisk() (*mapstate.MapState, error) {
 		err = errNoSnapshot
 	}
 	if err != nil {
-		return nil, err
+		return nil, false, err
 	}
 
 	stateFromSnap := mapstate.NewMapState()
-	err = stateFromSnap.Migrate(r)
+	// duplicate reader to both check version and migrate
+	var buf bytes.Buffer
+	r2 := io.TeeReader(r, &buf)
+	raw, err := ioutil.ReadAll(r2)
 	if err != nil {
-		return nil, err
+		return nil, false, err
+	}
+	err = stateFromSnap.Unmarshal(raw)
+	if err != nil {
+		return nil, false, err
+	}
+	if stateFromSnap.GetVersion() == mapstate.Version {
+		return stateFromSnap, true, nil
 	}
 
-	return stateFromSnap, nil
+	err = stateFromSnap.Migrate(&buf)
+	if err != nil {
+		return nil, false, err
+	}
 
+	return stateFromSnap, false, nil
 }
 
 func stateImport(r io.Reader) error {
diff --git a/state/mapstate/map_state.go b/state/mapstate/map_state.go
index 1301b9e2..624eb664 100644
--- a/state/mapstate/map_state.go
+++ b/state/mapstate/map_state.go
@@ -4,6 +4,7 @@ package mapstate
 
 import (
 	"bytes"
+	"errors"
 	"io"
 	"io/ioutil"
 	"sync"
@@ -99,6 +100,9 @@ func (st *MapState) Migrate(r io.Reader) error {
 		return err
 	}
 	err = st.Unmarshal(bs)
+	if err != nil {
+		return err
+	}
 	if st.Version == Version { // Unmarshal restored for us
 		return nil
 	}
@@ -142,6 +146,9 @@ func (st *MapState) Marshal() ([]byte, error) {
 func (st *MapState) Unmarshal(bs []byte) error {
 	// Check version byte
 	// logger.Debugf("The incoming bytes to unmarshal: %x", bs)
+	if len(bs) < 1 {
+		return errors.New("cannot unmarshal from empty bytes")
+	}
 	v := int(bs[0])
 	logger.Debugf("The interpreted version: %d", v)
 	if v != Version { // snapshot is out of date