《游戏服务端编程实践》4.3.1 微信 / Steam / Google OAuth 登录

// file: ch04/oauth/provider.go
// tier: L1
// verify: go test ./ch04/oauth/
package oauth

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"strings"
	"sync"
	"time"
)

type ProviderKind string

const (
	WeChat ProviderKind = "wechat"
	Steam  ProviderKind = "steam"
	Google ProviderKind = "google"
)

var (
	ErrStateMismatch = errors.New("oauth: state mismatch")
	ErrStateExpired  = errors.New("oauth: state expired")
	ErrCodeInvalid   = errors.New("oauth: authorization code not found")
	ErrCodeExpired   = errors.New("oauth: authorization code expired")
	ErrCodeReplayed  = errors.New("oauth: authorization code already used")
	ErrMissingOpenID = errors.New("oauth: provider returned empty external id")
)

type ProviderMeta struct {
	Kind                 ProviderKind
	CodeTTL              time.Duration
	ClientSecretRequired bool
	RefreshToken         bool
	IDField              string
}

var providerMeta = map[ProviderKind]ProviderMeta{
	WeChat: {Kind: WeChat, CodeTTL: 5 * time.Minute, ClientSecretRequired: true, RefreshToken: true, IDField: "openid"},
	Steam:  {Kind: Steam, ClientSecretRequired: false, IDField: "steamid"},
	Google: {Kind: Google, CodeTTL: 10 * time.Minute, RefreshToken: true, IDField: "sub"},
}

func MetaOf(kind ProviderKind) (ProviderMeta, bool) {
	m, ok := providerMeta[kind]
	return m, ok
}

type Identity struct {
	Provider   ProviderKind
	ExternalID string
	Nickname   string
	Claims     map[string]string
}

type Provider interface {
	Kind() ProviderKind
	Exchange(ctx context.Context, code string) (Identity, error)
}

type StateStore struct {
	mu     sync.Mutex
	ttl    time.Duration
	now    func() time.Time
	states map[string]time.Time
}

func NewStateStore(ttl time.Duration) *StateStore {
	return &StateStore{ttl: ttl, now: time.Now, states: make(map[string]time.Time)}
}

func (s *StateStore) Issue() (string, error) {
	buf := make([]byte, 16)
	if _, err := rand.Read(buf); err != nil {
		return "", err
	}
	state := hex.EncodeToString(buf)
	s.mu.Lock()
	defer s.mu.Unlock()
	s.states[state] = s.now().Add(s.ttl)
	return state, nil
}

func (s *StateStore) Consume(state string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	exp, ok := s.states[state]
	if !ok || state == "" {
		return ErrStateMismatch
	}
	delete(s.states, state)
	if s.now().After(exp) {
		return ErrStateExpired
	}
	return nil
}

type Authenticator struct {
	provider Provider
	states   *StateStore
}

func NewAuthenticator(p Provider, stateTTL time.Duration) *Authenticator {
	return &Authenticator{provider: p, states: NewStateStore(stateTTL)}
}

func (a *Authenticator) Begin() (string, error) { return a.states.Issue() }

func (a *Authenticator) Finish(ctx context.Context, code, state string) (Identity, error) {
	if err := a.states.Consume(state); err != nil {
		return Identity{}, err
	}
	return a.provider.Exchange(ctx, code)
}

type fakeCode struct {
	externalID string
	expiresAt  time.Time
	used       bool
}

type FakeProvider struct {
	kind  ProviderKind
	mu    sync.Mutex
	now   func() time.Time
	codes map[string]fakeCode
}

func newFake(kind ProviderKind) *FakeProvider {
	return &FakeProvider{kind: kind, now: time.Now, codes: make(map[string]fakeCode)}
}

func NewWeChatFake() *FakeProvider { return newFake(WeChat) }
func NewSteamFake() *FakeProvider  { return newFake(Steam) }
func NewGoogleFake() *FakeProvider { return newFake(Google) }

func (f *FakeProvider) Kind() ProviderKind { return f.kind }

func (f *FakeProvider) IssueCode(code, externalID string, ttl time.Duration) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.codes[code] = fakeCode{externalID: externalID, expiresAt: f.now().Add(ttl)}
}

func (f *FakeProvider) Exchange(_ context.Context, code string) (Identity, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	c, ok := f.codes[code]
	if !ok {
		return Identity{}, ErrCodeInvalid
	}
	if c.used {
		return Identity{}, ErrCodeReplayed
	}
	if f.now().After(c.expiresAt) {
		return Identity{}, ErrCodeExpired
	}
	c.used = true
	f.codes[code] = c
	if strings.TrimSpace(c.externalID) == "" {
		return Identity{}, ErrMissingOpenID
	}
	return Identity{Provider: f.kind, ExternalID: c.externalID, Claims: map[string]string{"code": code}}, nil
}

const steamClaimPrefix = "https://steamcommunity.com/openid/id/"

func SteamIDFromClaimedID(claimedID string) (string, error) {
	if !strings.HasPrefix(claimedID, steamClaimPrefix) {
		return "", ErrMissingOpenID
	}
	id := strings.TrimPrefix(claimedID, steamClaimPrefix)
	if id == "" || strings.ContainsAny(id, "/?#&") {
		return "", ErrMissingOpenID
	}
	return id, nil
}

// file: ch04/oauth/provider_test.go
// tier: L1
// verify: go test ./ch04/oauth/
package oauth

import (
	"context"
	"errors"
	"testing"
	"time"
)

func TestFinishRejectsStateMismatch(t *testing.T) {
	fake := NewWeChatFake()
	auth := NewAuthenticator(fake, time.Minute)
	state, err := auth.Begin()
	if err != nil {
		t.Fatalf("begin: %v", err)
	}
	fake.IssueCode("code-1", "openid-1", time.Minute)

	if _, err := auth.Finish(context.Background(), "code-1", state+"tampered"); !errors.Is(err, ErrStateMismatch) {
		t.Fatalf("tampered state: got %v, want %v", err, ErrStateMismatch)
	}
	if _, err := auth.Finish(context.Background(), "code-1", state); err != nil {
		t.Fatalf("original state should remain valid: %v", err)
	}
	if _, err := auth.Finish(context.Background(), "code-1", state); !errors.Is(err, ErrStateMismatch) {
		t.Fatalf("state reuse: got %v, want %v", err, ErrStateMismatch)
	}
}

func TestExchangeRejectsCodeReplay(t *testing.T) {
	fake := NewGoogleFake()
	fake.IssueCode("code-2", "sub-2", time.Minute)
	ctx := context.Background()
	if _, err := fake.Exchange(ctx, "code-2"); err != nil {
		t.Fatalf("first exchange: %v", err)
	}
	if _, err := fake.Exchange(ctx, "code-2"); !errors.Is(err, ErrCodeReplayed) {
		t.Fatalf("replay: got %v, want %v", err, ErrCodeReplayed)
	}
}

func TestExchangeRejectsMissingOpenID(t *testing.T) {
	fake := NewWeChatFake()
	fake.IssueCode("code-3", "   ", time.Minute)
	if _, err := fake.Exchange(context.Background(), "code-3"); !errors.Is(err, ErrMissingOpenID) {
		t.Fatalf("got %v, want %v", err, ErrMissingOpenID)
	}
}

func TestExchangeRejectsExpiredCode(t *testing.T) {
	fake := NewWeChatFake()
	fake.IssueCode("code-4", "openid-4", -time.Second)
	if _, err := fake.Exchange(context.Background(), "code-4"); !errors.Is(err, ErrCodeExpired) {
		t.Fatalf("got %v, want %v", err, ErrCodeExpired)
	}
}

func TestSteamIDFromClaimedID(t *testing.T) {
	got, err := SteamIDFromClaimedID("https://steamcommunity.com/openid/id/76561198000000000")
	if err != nil || got != "76561198000000000" {
		t.Fatalf("got %q, %v", got, err)
	}
	if _, err := SteamIDFromClaimedID("https://evil.example/id/76561198000000000"); !errors.Is(err, ErrMissingOpenID) {
		t.Fatalf("unexpected claimed_id accepted: %v", err)
	}
}