gpt-load/internal/control/bootstrap_test.go

package control

import (
	"bytes"
	"context"
	"encoding/hex"
	"encoding/json"
	"errors"
	"log"
	"strings"
	"testing"
	"time"

	"github.com/sirupsen/logrus"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"gpt-load/internal/channel"
	"gpt-load/internal/platform/encryption"
	"gpt-load/internal/state"
	"gpt-load/internal/storage/models"
)

const bootstrapMarkerForTest = models.InternalSystemSettingPrefix + "bootstrap.default_access_key.v1"

var errBootstrapEncryption = errors.New("forced bootstrap encryption failure")

func TestEnsureInitialStateCreatesDefaultAccessKeyAndMarker(t *testing.T) {
	t.Parallel()
	fixture := newServiceFixture(t)
	fixture.service.random = bytes.NewReader([]byte{
		0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
		0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
	})
	beforeRevision := fixture.manager.Current().Revision

	if err := fixture.service.EnsureInitialState(context.Background()); err != nil {
		t.Fatalf("EnsureInitialState() error = %v", err)
	}

	var row models.AccessKey
	if err := fixture.db.First(&row).Error; err != nil {
		t.Fatalf("query default AccessKey: %v", err)
	}
	plaintext, err := fixture.encryption.Decrypt(row.KeyValue)
	if err != nil {
		t.Fatalf("Decrypt(default key) error = %v", err)
	}
	if !strings.HasPrefix(plaintext, "sk-gl-") || len(plaintext) != len("sk-gl-")+32 {
		t.Fatal("default plaintext shape is invalid")
	}
	if _, err := hex.DecodeString(strings.TrimPrefix(plaintext, "sk-gl-")); err != nil {
		t.Fatalf("default plaintext suffix is not hex: %v", err)
	}
	if row.KeyHash != fixture.encryption.Hash(plaintext) {
		t.Fatal("default access key hash does not match plaintext")
	}
	if row.Name != "Default" || row.Status != string(state.AccessKeyStatusActive) ||
		row.RPMLimit != 0 {
		t.Fatalf("default row = %#v", row)
	}
	var filters AccessKeyFilters
	if err := json.Unmarshal(row.Filters, &filters); err != nil {
		t.Fatalf("decode default filters: %v", err)
	}
	if filters.Groups == nil || filters.Protocols == nil || filters.Models == nil ||
		len(filters.Groups) != 0 || len(filters.Protocols) != 0 || len(filters.Models) != 0 {
		t.Fatalf("default filters = %#v, want three non-nil empty slices", filters)
	}

	var marker models.SystemSetting
	if err := fixture.db.First(&marker, "key = ?", bootstrapMarkerForTest).Error; err != nil {
		t.Fatalf("query bootstrap marker: %v", err)
	}
	if marker.Key != bootstrapMarkerForTest || marker.Value != "true" {
		t.Fatalf("bootstrap marker = %#v, want exact JSON true", marker)
	}
	if got := fixture.manager.Current().Revision; got != beforeRevision {
		t.Fatalf("snapshot revision = %d, want unchanged %d", got, beforeRevision)
	}
	if len(fixture.manager.Current().AccessKeysByHash) != 0 {
		t.Fatal("EnsureInitialState() published a runtime snapshot")
	}
}

func TestEnsureInitialStateWithExistingAccessKeyOnlyWritesMarker(t *testing.T) {
	t.Parallel()
	fixture := newServiceFixture(t)
	const plaintext = "gl-existing-access-key"
	ciphertext, err := fixture.encryption.Encrypt(plaintext)
	if err != nil {
		t.Fatalf("Encrypt(existing key) error = %v", err)
	}
	existing := models.AccessKey{
		Name:      "Existing",
		KeyValue:  ciphertext,
		KeyHash:   fixture.encryption.Hash(plaintext),
		KeySuffix: "0000",
		Status:    string(state.AccessKeyStatusActive),
		Filters:   models.JSON(`{"groups":[],"protocols":[],"models":[]}`),
	}
	if err := fixture.db.Create(&existing).Error; err != nil {
		t.Fatalf("create existing AccessKey: %v", err)
	}

	if err := fixture.service.EnsureInitialState(context.Background()); err != nil {
		t.Fatalf("EnsureInitialState() error = %v", err)
	}

	var rows []models.AccessKey
	if err := fixture.db.Order("id ASC").Find(&rows).Error; err != nil {
		t.Fatalf("query AccessKeys: %v", err)
	}
	if len(rows) != 1 || rows[0].ID != existing.ID {
		t.Fatalf("AccessKeys = %#v, want only existing row", rows)
	}
	decrypted, err := fixture.encryption.Decrypt(rows[0].KeyValue)
	if err != nil || decrypted != plaintext {
		t.Fatalf("existing credential = %q, %v, want unchanged", decrypted, err)
	}
	assertBootstrapMarkerCount(t, fixture, 1)
}

func TestEnsureInitialStateDoesNotLogExpectedMarkerMissAsError(t *testing.T) {
	t.Parallel()
	fixture := newServiceFixture(t)
	var logs bytes.Buffer
	fixture.service.db = fixture.db.Session(&gorm.Session{
		Logger: logger.New(log.New(&logs, "\r\n", log.LstdFlags), logger.Config{
			SlowThreshold:        200 * time.Millisecond,
			LogLevel:             logger.Warn,
			ParameterizedQueries: true,
			Colorful:             true,
		}),
	})

	if err := fixture.service.EnsureInitialState(context.Background()); err != nil {
		t.Fatalf("EnsureInitialState() error = %v", err)
	}
	if strings.Contains(logs.String(), "record not found") {
		t.Fatalf("expected marker miss was logged as an error: %q", logs.String())
	}
}

func TestEnsureInitialStateIsIdempotentWhenMarkerExists(t *testing.T) {
	t.Parallel()
	fixture := newServiceFixture(t)
	if err := fixture.db.Create(&models.SystemSetting{
		Key: bootstrapMarkerForTest, Value: "true",
	}).Error; err != nil {
		t.Fatalf("create bootstrap marker: %v", err)
	}
	if err := fixture.db.Exec("DROP TABLE access_keys").Error; err != nil {
		t.Fatalf("drop AccessKey table: %v", err)
	}

	if err := fixture.service.EnsureInitialState(context.Background()); err != nil {
		t.Fatalf("EnsureInitialState() with marker error = %v", err)
	}
	assertBootstrapMarkerCount(t, fixture, 1)
}

func TestEnsureInitialStateDoesNotRecreateDeletedFinalKey(t *testing.T) {
	t.Parallel()
	fixture := newServiceFixture(t)
	fixture.service.random = bytes.NewReader(make([]byte, 16))
	if err := fixture.service.EnsureInitialState(context.Background()); err != nil {
		t.Fatalf("first EnsureInitialState() error = %v", err)
	}
	if err := fixture.db.Where("1 = 1").Delete(&models.AccessKey{}).Error; err != nil {
		t.Fatalf("delete final AccessKey: %v", err)
	}

	if err := fixture.service.EnsureInitialState(context.Background()); err != nil {
		t.Fatalf("second EnsureInitialState() error = %v", err)
	}
	assertAccessKeyCount(t, fixture, 0)
	assertBootstrapMarkerCount(t, fixture, 1)
}

func TestEnsureInitialStateRollsBackWhenEncryptionFails(t *testing.T) {
	t.Parallel()
	fixture := newServiceFixture(t)
	fixture.service.random = bytes.NewReader(make([]byte, 16))
	fixture.service.encryption = bootstrapFailingEncryptService{Service: fixture.encryption}

	err := fixture.service.EnsureInitialState(context.Background())
	if !errors.Is(err, errBootstrapEncryption) {
		t.Fatalf("EnsureInitialState() error = %v, want encryption failure", err)
	}
	assertAccessKeyCount(t, fixture, 0)
	assertBootstrapMarkerCount(t, fixture, 0)
}

func TestEnsureInitialStateRollsBackAccessKeyWhenMarkerWriteFails(t *testing.T) {
	t.Parallel()
	fixture := newServiceFixture(t)
	fixture.service.random = bytes.NewReader(make([]byte, 16))
	if err := fixture.db.Exec(`
		CREATE TRIGGER reject_bootstrap_marker
		BEFORE INSERT ON system_settings
		WHEN NEW.key = '_internal.bootstrap.default_access_key.v1'
		BEGIN
		  SELECT RAISE(ABORT, 'marker rejected');
		END;
	`).Error; err != nil {
		t.Fatalf("create marker rejection trigger: %v", err)
	}

	if err := fixture.service.EnsureInitialState(context.Background()); err == nil {
		t.Fatal("EnsureInitialState() error = nil, want marker rejection")
	}
	assertAccessKeyCount(t, fixture, 0)
	assertBootstrapMarkerCount(t, fixture, 0)
}

func TestEnsureInitialStateDoesNotLogPlaintext(t *testing.T) {
	// 不标记 t.Parallel()：本测试劫持了全局 logrus 输出/格式，与其他并行测试同时运行会互相覆盖断言。
	fixture := newServiceFixture(t)
	randomBytes := bytes.Repeat([]byte{0xab}, 16)
	fixture.service.random = bytes.NewReader(randomBytes)
	expectedPlaintext := "sk-gl-" + hex.EncodeToString(randomBytes)

	var logs bytes.Buffer
	logger := logrus.StandardLogger()
	previousOutput, previousFormatter, previousLevel := logger.Out, logger.Formatter, logger.GetLevel()
	logrus.SetOutput(&logs)
	logrus.SetFormatter(&logrus.JSONFormatter{DisableTimestamp: true})
	logrus.SetLevel(logrus.DebugLevel)
	t.Cleanup(func() {
		logrus.SetOutput(previousOutput)
		logrus.SetFormatter(previousFormatter)
		logrus.SetLevel(previousLevel)
	})

	if err := fixture.service.EnsureInitialState(context.Background()); err != nil {
		t.Fatalf("EnsureInitialState() error = %v", err)
	}
	if strings.Contains(logs.String(), expectedPlaintext) {
		t.Fatalf("bootstrap logs exposed plaintext: %q", logs.String())
	}
}

func TestEnsureInitialStateRecoversInterruptedSubscriptionAuth(t *testing.T) {
	t.Parallel()

	fixture := newServiceFixture(t)
	now := time.UnixMilli(1_800_000_000_000)
	fixture.service.now = func() time.Time { return now }
	group := validControlGroup("subscription-recovery")
	group.ChannelID = string(channel.Codex)
	group.ConnectionType = models.ConnectionTypeSubscription
	if err := fixture.db.Create(group).Error; err != nil {
		t.Fatal(err)
	}
	credential := models.Credential{
		GroupID: group.ID, Data: "cipher", Fingerprint: "secret", IdentityFingerprint: "identity",
		SecretVersion: 2, AuthState: models.CredentialAuthStateRefreshing,
		Status: models.CredentialStatusActive, CreatedAtMS: now.Add(-time.Hour).UnixMilli(), UpdatedAtMS: now.Add(-time.Hour).UnixMilli(),
	}
	if err := fixture.db.Create(&credential).Error; err != nil {
		t.Fatal(err)
	}
	stage := models.CredentialStage{
		ID: "00000000-0000-4000-8000-000000000901", ChannelID: "codex",
		ConnectionType: models.ConnectionTypeSubscription, AuthorizationMethod: "browser_oauth",
		Status: models.CredentialStageExchanging, EncryptedPayload: "encrypted-stage", PayloadSchemaVersion: 1,
		SafeSummaryJSON: models.JSON(`{}`), ExpiresAtMS: now.Add(time.Minute).UnixMilli(),
		CreatedAtMS: now.Add(-time.Minute).UnixMilli(), UpdatedAtMS: now.Add(-time.Minute).UnixMilli(),
	}
	if err := fixture.db.Create(&stage).Error; err != nil {
		t.Fatal(err)
	}
	if err := fixture.service.EnsureInitialState(t.Context()); err != nil {
		t.Fatal(err)
	}
	if err := fixture.db.Take(&credential, credential.ID).Error; err != nil {
		t.Fatal(err)
	}
	if credential.AuthState != models.CredentialAuthStateOutcomeUnknown || credential.AuthErrorCode != "refresh_interrupted" {
		t.Fatalf("credential = %#v", credential)
	}
	if err := fixture.db.Take(&stage, "id = ?", stage.ID).Error; err != nil {
		t.Fatal(err)
	}
	if stage.Status != models.CredentialStageOutcomeUnknown || stage.EncryptedPayload != "" || stage.ErrorCode != "authorization_exchange_interrupted" {
		t.Fatalf("stage = %#v", stage)
	}
}

type bootstrapFailingEncryptService struct {
	encryption.Service
}

func (bootstrapFailingEncryptService) Encrypt(string) (string, error) {
	return "", errBootstrapEncryption
}

func assertAccessKeyCount(t *testing.T, fixture serviceFixture, want int64) {
	t.Helper()
	var count int64
	if err := fixture.db.Model(&models.AccessKey{}).Count(&count).Error; err != nil {
		t.Fatalf("count AccessKeys: %v", err)
	}
	if count != want {
		t.Fatalf("AccessKey count = %d, want %d", count, want)
	}
}

func assertBootstrapMarkerCount(t *testing.T, fixture serviceFixture, want int64) {
	t.Helper()
	var count int64
	if err := fixture.db.Model(&models.SystemSetting{}).
		Where("key = ?", bootstrapMarkerForTest).Count(&count).Error; err != nil {
		t.Fatalf("count bootstrap marker: %v", err)
	}
	if count != want {
		t.Fatalf("bootstrap marker count = %d, want %d", count, want)
	}
}

func assertModelPriceCount(t *testing.T, fixture serviceFixture, want int64) {
	t.Helper()
	var count int64
	if err := fixture.db.Model(&models.ModelPrice{}).Count(&count).Error; err != nil {
		t.Fatalf("count ModelPrice rows: %v", err)
	}
	if count != want {
		t.Fatalf("ModelPrice count = %d, want %d", count, want)
	}
}