Golang State.Set方法代码示例

// patch1 adds the snap type and the current revision to the snap state.
func patch1(s *state.State) error {
	var stateMap map[string]*snapstate.SnapState

	err := s.Get("snaps", &stateMap)
	if err == state.ErrNoState {
		return nil
	}
	if err != nil {
		return err
	}

	for snapName, snapState := range stateMap {
		seq := snapState.Sequence
		if len(seq) == 0 {
			continue
		}
		typ := snap.TypeApp
		snapInfo, err := readInfo(snapName, seq[len(seq)-1])
		if err != nil {
			logger.Noticef("Recording type for snap %q: cannot retrieve info, assuming it's a app: %v", snapName, err)
		} else {
			logger.Noticef("Recording type for snap %q: setting to %q", snapName, snapInfo.Type)
			typ = snapInfo.Type
		}
		snapState.SetType(typ)
		snapState.Current = seq[len(seq)-1].Revision
	}

	s.Set("snaps", stateMap)
	return nil
}

// NewUser tracks a new authenticated user and saves its details in the state
func NewUser(st *state.State, username, macaroon string, discharges []string) (*UserState, error) {
	var authStateData AuthState

	err := st.Get("auth", &authStateData)
	if err == state.ErrNoState {
		authStateData = AuthState{}
	} else if err != nil {
		return nil, err
	}

	sort.Strings(discharges)
	authStateData.LastID++
	authenticatedUser := UserState{
		ID:              authStateData.LastID,
		Username:        username,
		Macaroon:        macaroon,
		Discharges:      discharges,
		StoreMacaroon:   macaroon,
		StoreDischarges: discharges,
	}
	authStateData.Users = append(authStateData.Users, authenticatedUser)

	st.Set("auth", authStateData)

	return &authenticatedUser, nil
}

func newRunnerManager(s *state.State) *runnerManager {
	rm := &runnerManager{
		runner: state.NewTaskRunner(s),
	}

	rm.runner.AddHandler("runMgr1", func(t *state.Task, _ *tomb.Tomb) error {
		s := t.State()
		s.Lock()
		defer s.Unlock()
		s.Set("runMgr1Mark", 1)
		return nil
	}, nil)
	rm.runner.AddHandler("runMgr2", func(t *state.Task, _ *tomb.Tomb) error {
		s := t.State()
		s.Lock()
		defer s.Unlock()
		s.Set("runMgr2Mark", 1)
		return nil
	}, nil)
	rm.runner.AddHandler("runMgrEnsureBefore", func(t *state.Task, _ *tomb.Tomb) error {
		s := t.State()
		s.Lock()
		defer s.Unlock()
		s.EnsureBefore(20 * time.Millisecond)
		return nil
	}, nil)

	return rm
}

// Init initializes an empty state to the current implemented patch level.
func Init(s *state.State) {
	s.Lock()
	defer s.Unlock()
	if s.Get("patch-level", new(int)) != state.ErrNoState {
		panic("internal error: expected empty state, attempting to override patch-level without actual patching")
	}
	s.Set("patch-level", Level)
}

// patch2:
// - migrates SnapSetup.Name to SnapSetup.SideInfo.RealName
// - backfills SnapState.{Sequence,Candidate}.RealName if its missing
func patch2(s *state.State) error {

	var stateMap map[string]*OldSnapState
	err := s.Get("snaps", &stateMap)
	if err == state.ErrNoState {
		return nil
	}
	if err != nil {
		return err
	}

	// migrate SnapSetup in all tasks:
	//  - the new SnapSetup uses SideInfo, backfil from Candidate
	//  - also move SnapSetup.{Name,Revision} into SnapSetup.SideInfo.{RealName,Revision}
	var oldSS OldSnapSetup
	var newSS snapstate.SnapSetup
	for _, t := range s.Tasks() {
		err := t.Get("snap-setup", &oldSS)
		if err == state.ErrNoState {
			continue
		}
		if err != nil && err != state.ErrNoState {
			return err
		}
		// some things stay the same
		newSS.Channel = oldSS.Channel
		newSS.Flags = oldSS.Flags
		newSS.SnapPath = oldSS.SnapPath
		newSS.DownloadInfo = oldSS.DownloadInfo
		newSS.SideInfo = oldSS.SideInfo
		// ... and some change
		if newSS.SideInfo == nil {
			newSS.SideInfo = &snap.SideInfo{}
			if snapst, ok := stateMap[oldSS.Name]; ok && snapst.Candidate != nil {
				newSS.SideInfo = snapst.Candidate
			}
		}
		if newSS.SideInfo.RealName == "" {
			newSS.SideInfo.RealName = oldSS.Name
		}
		if newSS.SideInfo.Revision.Unset() {
			newSS.SideInfo.Revision = oldSS.Revision
		}
		t.Set("snap-setup", &newSS)
	}

	// backfill snapstate.SnapState.{Sequence,Candidate} with RealName
	// (if that is missing, was missing for e.g. sideloaded snaps)
	for snapName, snapState := range stateMap {
		for _, si := range snapState.Sequence {
			setRealName(si, snapName)
		}
	}
	s.Set("snaps", stateMap)

	return nil
}

func applyOne(patch func(s *state.State) error, s *state.State, level int) error {
	s.Lock()
	defer s.Unlock()

	err := patch(s)
	if err != nil {
		return err
	}

	s.Set("patch-level", level+1)
	return nil
}

// patch2:
// - migrates SnapSetup.Name to SnapSetup.SideInfo.RealName
// - backfills SnapState.{Sequence,Candidate}.RealName if its missing
func patch2(s *state.State) error {

	var oldStateMap map[string]*patch1SnapState
	err := s.Get("snaps", &oldStateMap)
	if err == state.ErrNoState {
		return nil
	}
	if err != nil {
		return err
	}
	newStateMap := make(map[string]*patch2SnapState, len(oldStateMap))

	for key, oldSnapState := range oldStateMap {
		newStateMap[key] = patch2SnapStateFromPatch1(oldSnapState, key)
	}

	// migrate SnapSetup in all tasks:
	//  - the new SnapSetup uses SideInfo, backfil from Candidate
	//  - also move SnapSetup.{Name,Revision} into SnapSetup.SideInfo.{RealName,Revision}
	var oldSS patch1SnapSetup
	for _, t := range s.Tasks() {
		var newSS patch2SnapSetup
		err := t.Get("snap-setup", &oldSS)
		if err == state.ErrNoState {
			continue
		}
		if err != nil && err != state.ErrNoState {
			return err
		}
		// some things stay the same
		newSS.Channel = oldSS.Channel
		newSS.Flags = patch2Flags(oldSS.Flags)
		newSS.SnapPath = oldSS.SnapPath
		// ... and some change
		newSS.SideInfo = &patch2SideInfo{}
		if snapst, ok := oldStateMap[oldSS.Name]; ok && snapst.Candidate != nil {
			newSS.SideInfo = patch2SideInfoFromPatch1(snapst.Candidate, oldSS.Name)
		}
		if newSS.SideInfo.RealName == "" {
			newSS.SideInfo.RealName = oldSS.Name
		}
		if newSS.SideInfo.Revision.Unset() {
			newSS.SideInfo.Revision = oldSS.Revision
		}
		t.Set("snap-setup", &newSS)
	}

	s.Set("snaps", newStateMap)

	return nil
}

// SetDevice updates the device details in the state.
func SetDevice(st *state.State, device *DeviceState) error {
	var authStateData AuthState

	err := st.Get("auth", &authStateData)
	if err == state.ErrNoState {
		authStateData = AuthState{}
	} else if err != nil {
		return err
	}

	authStateData.Device = device
	st.Set("auth", authStateData)

	return nil
}

// patch6:
//  - move from a flags-are-ints world to a flags-are-struct-of-bools world
func patch6(st *state.State) error {
	var oldStateMap map[string]*patch4SnapState
	err := st.Get("snaps", &oldStateMap)
	if err == state.ErrNoState {
		return nil
	}
	if err != nil {
		return err
	}
	newStateMap := make(map[string]*patch6SnapState, len(oldStateMap))

	for key, old := range oldStateMap {
		newStateMap[key] = &patch6SnapState{
			SnapType:    old.SnapType,
			Sequence:    old.Sequence,
			Active:      old.Active,
			Current:     old.Current,
			Channel:     old.Channel,
			patch6Flags: patch6FlagsFromPatch4(old.Flags),
		}
	}

	for _, task := range st.Tasks() {
		var old patch4SnapSetup
		err := task.Get("snap-setup", &old)
		if err == state.ErrNoState {
			continue
		}
		if err != nil && err != state.ErrNoState {
			return err
		}

		task.Set("snap-setup", &patch6SnapSetup{
			Channel:      old.Channel,
			UserID:       old.UserID,
			SnapPath:     old.SnapPath,
			DownloadInfo: old.DownloadInfo,
			SideInfo:     old.SideInfo,
			patch6Flags:  patch6FlagsFromPatch4(old.Flags),
		})
	}

	st.Set("snaps", newStateMap)

	return nil
}

// UpdateUser updates user in state
func UpdateUser(st *state.State, user *UserState) error {
	var authStateData AuthState

	err := st.Get("auth", &authStateData)
	if err != nil {
		return err
	}

	for i := range authStateData.Users {
		if authStateData.Users[i].ID == user.ID {
			authStateData.Users[i] = *user
			st.Set("auth", authStateData)
			return nil
		}
	}

	return fmt.Errorf("invalid user")
}

// Set sets the SnapState of the given snap, overwriting any earlier state.
func Set(s *state.State, name string, snapst *SnapState) {
	var snaps map[string]*json.RawMessage
	err := s.Get("snaps", &snaps)
	if err != nil && err != state.ErrNoState {
		panic("internal error: cannot unmarshal snaps state: " + err.Error())
	}
	if snaps == nil {
		snaps = make(map[string]*json.RawMessage)
	}
	if snapst == nil || (len(snapst.Sequence) == 0) {
		delete(snaps, name)
	} else {
		data, err := json.Marshal(snapst)
		if err != nil {
			panic("internal error: cannot marshal snap state: " + err.Error())
		}
		raw := json.RawMessage(data)
		snaps[name] = &raw
	}
	s.Set("snaps", snaps)
}

func setAliases(st *state.State, snapName string, aliases map[string]string) {
	var allAliases map[string]*json.RawMessage
	err := st.Get("aliases", &allAliases)
	if err != nil && err != state.ErrNoState {
		panic("internal error: cannot unmarshal snap aliases state: " + err.Error())
	}
	if allAliases == nil {
		allAliases = make(map[string]*json.RawMessage)
	}
	if len(aliases) == 0 {
		delete(allAliases, snapName)
	} else {
		data, err := json.Marshal(aliases)
		if err != nil {
			panic("internal error: cannot marshal snap aliases state: " + err.Error())
		}
		raw := json.RawMessage(data)
		allAliases[snapName] = &raw
	}
	st.Set("aliases", allAliases)
}

// NewUser tracks a new authenticated user and saves its details in the state
func NewUser(st *state.State, username, email, macaroon string, discharges []string) (*UserState, error) {
	var authStateData AuthState

	err := st.Get("auth", &authStateData)
	if err == state.ErrNoState {
		authStateData = AuthState{}
	} else if err != nil {
		return nil, err
	}

	if authStateData.MacaroonKey == nil {
		authStateData.MacaroonKey, err = generateMacaroonKey()
		if err != nil {
			return nil, err
		}
	}

	authStateData.LastID++

	localMacaroon, err := newUserMacaroon(authStateData.MacaroonKey, authStateData.LastID)
	if err != nil {
		return nil, err
	}

	sort.Strings(discharges)
	authenticatedUser := UserState{
		ID:              authStateData.LastID,
		Username:        username,
		Email:           email,
		Macaroon:        localMacaroon,
		Discharges:      nil,
		StoreMacaroon:   macaroon,
		StoreDischarges: discharges,
	}
	authStateData.Users = append(authStateData.Users, authenticatedUser)

	st.Set("auth", authStateData)

	return &authenticatedUser, nil
}

// RemoveUser removes a user from the state given its ID
func RemoveUser(st *state.State, userID int) error {
	var authStateData AuthState

	err := st.Get("auth", &authStateData)
	if err != nil {
		return err
	}

	for i := range authStateData.Users {
		if authStateData.Users[i].ID == userID {
			// delete without preserving order
			n := len(authStateData.Users) - 1
			authStateData.Users[i] = authStateData.Users[n]
			authStateData.Users[n] = UserState{}
			authStateData.Users = authStateData.Users[:n]
			st.Set("auth", authStateData)
			return nil
		}
	}

	return fmt.Errorf("invalid user")
}

func setConns(st *state.State, conns map[string]connState) {
	st.Set("conns", conns)
}