《游戏服务端编程实践》17.3 回滚与安全策略

// file: ch17/rollback/rollback.go
// tier: L1
// verify: go test ./ch17/rollback/

// Package rollback 管理带签名的版本保留窗口与回滚操作。
package rollback

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"sync"
	"time"
)

var (
	ErrEmptyKey          = errors.New("rollback: empty signing key")
	ErrInvalidRetention  = errors.New("rollback: retention must be positive")
	ErrInvalidVersion    = errors.New("rollback: invalid version")
	ErrDuplicateVersion  = errors.New("rollback: duplicate version")
	ErrVersionOrder      = errors.New("rollback: version is older than latest")
	ErrUnknownVersion    = errors.New("rollback: unknown version")
	ErrOutsideWindow     = errors.New("rollback: version outside retention window")
	ErrSignatureMismatch = errors.New("rollback: signature mismatch")
	ErrNoActiveVersion   = errors.New("rollback: no active version")
)

// Artifact 是一个已签名的发布包。Signature 对 Version 与 Payload
// 一起签名，单独替换其中任何一个都会验签失败。
type Artifact struct {
	Version     int       `json:"version"`
	Payload     []byte    `json:"payload"`
	Signature   []byte    `json:"signature"`
	PublishedAt time.Time `json:"published_at"`
}

// Sign 计算版本号与载荷的 HMAC-SHA256。版本号按大端序拼在载荷之前，
// 防止把 v1 的签名挪用到 v2。
func Sign(key []byte, version int, payload []byte) []byte {
	if len(key) == 0 || version < 1 {
		return nil
	}
	mac := hmac.New(sha256.New, key)
	var buf [8]byte
	binary.BigEndian.PutUint64(buf[:], uint64(version))
	mac.Write(buf[:])
	mac.Write(payload)
	return mac.Sum(nil)
}

// Verify 校验签名。空密钥或空签名一律视为失败，避免“未签名即通过”。
func Verify(key []byte, version int, payload, signature []byte) bool {
	if len(key) == 0 || len(signature) == 0 {
		return false
	}
	want := Sign(key, version, payload)
	return len(want) == len(signature) && hmac.Equal(want, signature)
}

// Store 保留最近若干个版本，并记录所有曾经出现过的版本号，以便
// 区分“从未发布”和“已滑出保留窗口”。
type Store struct {
	mu       sync.Mutex
	key      []byte
	retain   int
	now      func() time.Time
	known    map[int]struct{}
	retained []*Artifact
	active   int
}

// New 创建存储。key 不能为空，retain 至少为 1。
func New(key []byte, retain int) (*Store, error) {
	return NewWithClock(key, retain, time.Now)
}

// NewWithClock 允许注入时钟，便于测试稳定地构造发布时间。
func NewWithClock(key []byte, retain int, now func() time.Time) (*Store, error) {
	if len(key) == 0 {
		return nil, ErrEmptyKey
	}
	if retain < 1 {
		return nil, ErrInvalidRetention
	}
	if now == nil {
		now = time.Now
	}
	return &Store{
		key:    append([]byte(nil), key...),
		retain: retain,
		now:    now,
		known:  make(map[int]struct{}),
	}, nil
}

// Publish 使用存储内的密钥签名并发布一个新版本。发布成功后该版本
// 立即成为活动版本。
func (s *Store) Publish(version int, payload []byte) (Artifact, error) {
	if version < 1 {
		return Artifact{}, fmt.Errorf("%w: %d", ErrInvalidVersion, version)
	}
	a := Artifact{
		Version:     version,
		Payload:     append([]byte(nil), payload...),
		Signature:   Sign(s.key, version, payload),
		PublishedAt: s.now(),
	}
	return s.publish(a)
}

// PublishArtifact 接收外部签名好的发布包，先验签再保留。回滚包也
// 必须走这里或 Rollback，不能绕过签名检查。
func (s *Store) PublishArtifact(a Artifact) (Artifact, error) {
	if a.Version < 1 {
		return Artifact{}, fmt.Errorf("%w: %d", ErrInvalidVersion, a.Version)
	}
	return s.publish(a)
}

func (s *Store) publish(a Artifact) (Artifact, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !Verify(s.key, a.Version, a.Payload, a.Signature) {
		return Artifact{}, fmt.Errorf("%w: version=%d", ErrSignatureMismatch, a.Version)
	}
	if _, ok := s.known[a.Version]; ok {
		return Artifact{}, fmt.Errorf("%w: %d", ErrDuplicateVersion, a.Version)
	}
	if n := len(s.retained); n > 0 && a.Version <= s.retained[n-1].Version {
		return Artifact{}, fmt.Errorf("%w: latest=%d incoming=%d",
			ErrVersionOrder, s.retained[n-1].Version, a.Version)
	}
	if a.PublishedAt.IsZero() {
		a.PublishedAt = s.now()
	}
	stored := cloneArtifact(a)
	s.known[a.Version] = struct{}{}
	s.retained = append(s.retained, &stored)
	s.active = a.Version
	if len(s.retained) > s.retain {
		s.retained = s.retained[len(s.retained)-s.retain:]
	}
	return cloneArtifact(stored), nil
}

// Rollback 把活动版本切回指定版本。目标必须在保留窗口内，并且
// 在切换前重新验签；这样即使存储被篡改，也不会回滚到被改过的包。
func (s *Store) Rollback(version int) (Artifact, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.known[version]; !ok {
		return Artifact{}, fmt.Errorf("%w: %d", ErrUnknownVersion, version)
	}
	var target *Artifact
	for _, a := range s.retained {
		if a.Version == version {
			target = a
			break
		}
	}
	if target == nil {
		return Artifact{}, fmt.Errorf("%w: %d retained=%v",
			ErrOutsideWindow, version, s.retainedVersionsLocked())
	}
	if !Verify(s.key, target.Version, target.Payload, target.Signature) {
		return Artifact{}, fmt.Errorf("%w: version=%d", ErrSignatureMismatch, target.Version)
	}
	s.active = version
	return cloneArtifact(*target), nil
}

// Active 返回当前活动版本的副本。
func (s *Store) Active() (Artifact, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.active == 0 {
		return Artifact{}, false
	}
	for _, a := range s.retained {
		if a.Version == s.active {
			return cloneArtifact(*a), true
		}
	}
	return Artifact{}, false
}

// Retained 返回保留窗口内的版本号，按发布顺序排列。
func (s *Store) Retained() []int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.retainedVersionsLocked()
}

func (s *Store) retainedVersionsLocked() []int {
	out := make([]int, len(s.retained))
	for i, a := range s.retained {
		out[i] = a.Version
	}
	return out
}

// Known 表示该版本是否曾经发布过，无论是否仍在保留窗口内。
func (s *Store) Known(version int) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	_, ok := s.known[version]
	return ok
}

func cloneArtifact(a Artifact) Artifact {
	return Artifact{
		Version:     a.Version,
		Payload:     append([]byte(nil), a.Payload...),
		Signature:   append([]byte(nil), a.Signature...),
		PublishedAt: a.PublishedAt,
	}
}

cd .local/checks/gs-practice-lab
gofmt -l ch17/rollback
go vet ./ch17/rollback/
go test ./ch17/rollback/ -count=1