github.com/lazyledger/lazyledger-core@v0.35.0-dev.0.20210613111200-4c651f053571/test/e2e/app/state.go

//nolint: gosec
package main

import (
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"sort"
	"sync"
)

// State is the application state.
type State struct {
	sync.RWMutex
	Height uint64
	Values map[string]string
	Hash   []byte

	// private fields aren't marshalled to disk.
	file            string
	persistInterval uint64
	initialHeight   uint64
}

// NewState creates a new state.
func NewState(file string, persistInterval uint64) (*State, error) {
	state := &State{
		Values:          make(map[string]string),
		file:            file,
		persistInterval: persistInterval,
	}
	state.Hash = hashItems(state.Values)
	err := state.load()
	switch {
	case errors.Is(err, os.ErrNotExist):
	case err != nil:
		return nil, err
	}
	return state, nil
}

// load loads state from disk. It does not take out a lock, since it is called
// during construction.
func (s *State) load() error {
	bz, err := ioutil.ReadFile(s.file)
	if err != nil {
		return fmt.Errorf("failed to read state from %q: %w", s.file, err)
	}
	err = json.Unmarshal(bz, s)
	if err != nil {
		return fmt.Errorf("invalid state data in %q: %w", s.file, err)
	}
	return nil
}

// save saves the state to disk. It does not take out a lock since it is called
// internally by Commit which does lock.
func (s *State) save() error {
	bz, err := json.Marshal(s)
	if err != nil {
		return fmt.Errorf("failed to marshal state: %w", err)
	}
	// We write the state to a separate file and move it to the destination, to
	// make it atomic.
	newFile := fmt.Sprintf("%v.new", s.file)
	err = ioutil.WriteFile(newFile, bz, 0644)
	if err != nil {
		return fmt.Errorf("failed to write state to %q: %w", s.file, err)
	}
	return os.Rename(newFile, s.file)
}

// Export exports key/value pairs as JSON, used for state sync snapshots.
func (s *State) Export() ([]byte, error) {
	s.RLock()
	defer s.RUnlock()
	return json.Marshal(s.Values)
}

// Import imports key/value pairs from JSON bytes, used for InitChain.AppStateBytes and
// state sync snapshots. It also saves the state once imported.
func (s *State) Import(height uint64, jsonBytes []byte) error {
	s.Lock()
	defer s.Unlock()
	values := map[string]string{}
	err := json.Unmarshal(jsonBytes, &values)
	if err != nil {
		return fmt.Errorf("failed to decode imported JSON data: %w", err)
	}
	s.Height = height
	s.Values = values
	s.Hash = hashItems(values)
	return s.save()
}

// Get fetches a value. A missing value is returned as an empty string.
func (s *State) Get(key string) string {
	s.RLock()
	defer s.RUnlock()
	return s.Values[key]
}

// Set sets a value. Setting an empty value is equivalent to deleting it.
func (s *State) Set(key, value string) {
	s.Lock()
	defer s.Unlock()
	if value == "" {
		delete(s.Values, key)
	} else {
		s.Values[key] = value
	}
}

// Commit commits the current state.
func (s *State) Commit() (uint64, []byte, error) {
	s.Lock()
	defer s.Unlock()
	s.Hash = hashItems(s.Values)
	switch {
	case s.Height > 0:
		s.Height++
	case s.initialHeight > 0:
		s.Height = s.initialHeight
	default:
		s.Height = 1
	}
	if s.persistInterval > 0 && s.Height%s.persistInterval == 0 {
		err := s.save()
		if err != nil {
			return 0, nil, err
		}
	}
	return s.Height, s.Hash, nil
}

// hashItems hashes a set of key/value items.
func hashItems(items map[string]string) []byte {
	keys := make([]string, 0, len(items))
	for key := range items {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	hasher := sha256.New()
	for _, key := range keys {
		_, _ = hasher.Write([]byte(key))
		_, _ = hasher.Write([]byte{0})
		_, _ = hasher.Write([]byte(items[key]))
		_, _ = hasher.Write([]byte{0})
	}
	return hasher.Sum(nil)
}