From c1b32d87368a69036f20e3a5029b9e4556bb572b Mon Sep 17 00:00:00 2001
From: Weny Xu <wenymedia@gmail.com>
Date: Tue, 10 May 2022 22:35:36 +0900
Subject: [PATCH] feat(store): add header for snapshot (#55)

---
 pkg/store/fsm.go | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/pkg/store/fsm.go b/pkg/store/fsm.go
index e6d0e1b..f6426a4 100644
--- a/pkg/store/fsm.go
+++ b/pkg/store/fsm.go
@@ -19,12 +19,14 @@ package store
 
 import (
 	"bytes"
+	"encoding/binary"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"github.com/casbin/casbin-mesh/pkg/adapter"
 	"github.com/casbin/casbin-mesh/pkg/auth"
 	"io"
+	"io/ioutil"
 	"log"
 	"sync"
 	"time"
@@ -302,6 +304,15 @@ type persistData struct {
 	CredentialStore []byte
 }
 
+// SnapshotHdr is used to identify the snapshot protocol version.
+// length 8 bytes
+func SnapshotHdr() []byte {
+	hdr := [8]byte{}
+	// protocol version
+	binary.LittleEndian.PutUint16(hdr[0:], uint16(1))
+	return hdr[:]
+}
+
 // Persist implements persistence of states
 func (f fsmSnapshot) Persist(sink raft.SnapshotSink) error {
 	defer func() {
@@ -318,7 +329,7 @@ func (f fsmSnapshot) Persist(sink raft.SnapshotSink) error {
 			return err
 		}
 		// JSON the cluster Enforcers.
-		if _, err := sink.Write(data); err != nil {
+		if _, err := sink.Write(append(SnapshotHdr(), data...)); err != nil {
 			return err
 		}
 
@@ -383,7 +394,8 @@ func (s *Store) Snapshot() (raft.FSMSnapshot, error) {
 func (s *Store) Restore(closer io.ReadCloser) error {
 	var err error
 	var data persistData
-	err = json.NewDecoder(closer).Decode(&data)
+	snap, err := ioutil.ReadAll(closer)
+	err = json.Unmarshal(snap[8:], &data)
 	if err != nil {
 		s.logger.Println("failed to decode restore data", err)
 		return err