flnxi-Wu/flnxi-agent-go/internal/api/handlers_test.go

package api

import (
	"context"
	"errors"
	"strings"
	"testing"
	"time"

	"flnxi-agent/internal/docker"

	"github.com/docker/docker/api/types/container"
	"github.com/docker/go-connections/nat"
)

func TestResolveUpdateImageRejectsDirectImageWithoutFixedTag(t *testing.T) {
	tests := []struct {
		name    string
		image   string
		wantErr string
	}{
		{name: "no tag", image: "registry.local:5000/material/flnxi-material", wantErr: "fixed version tag"},
		{name: "latest", image: "registry.local:5000/material/flnxi-material:latest", wantErr: "latest is not allowed"},
		{name: "empty tag", image: "registry.local:5000/material/flnxi-material:", wantErr: "fixed version tag"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, err := resolveUpdateImage(UpdateRequest{Image: tt.image})
			if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
				t.Fatalf("resolveUpdateImage() error = %v, want contains %q", err, tt.wantErr)
			}
		})
	}
}

func TestResolveUpdateImageAcceptsDirectImageWithFixedTag(t *testing.T) {
	image := "registry.local:5000/material/flnxi-material:1.2.3"
	got, err := resolveUpdateImage(UpdateRequest{Image: image})
	if err != nil {
		t.Fatalf("resolveUpdateImage() unexpected error: %v", err)
	}
	if got != image {
		t.Fatalf("resolveUpdateImage() = %q, want %q", got, image)
	}
}

func TestResolveUpdateImageRejectsLatestVersion(t *testing.T) {
	t.Setenv(envMaterialImageRepository, "registry.local:5000/material/flnxi-material")

	_, err := resolveUpdateImage(UpdateRequest{Version: "LATEST"})
	if err == nil || !strings.Contains(err.Error(), "latest is not allowed") {
		t.Fatalf("resolveUpdateImage() error = %v, want latest rejection", err)
	}
}

func TestResolveInstallImagesRequiresFixedMaterialTag(t *testing.T) {
	tests := []struct {
		name    string
		image   string
		wantErr string
	}{
		{name: "missing", image: "", wantErr: envMaterialImage},
		{name: "no tag", image: "registry.local:5000/material/flnxi-material", wantErr: "fixed version tag"},
		{name: "latest", image: "registry.local:5000/material/flnxi-material:latest", wantErr: "latest is not allowed"},
		{name: "fixed tag", image: "registry.local:5000/material/flnxi-material:1.0.1"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, _, _, err := resolveInstallImages(InstallRequest{MaterialImage: tt.image})
			if tt.wantErr != "" {
				if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
					t.Fatalf("resolveInstallImages() error = %v, want contains %q", err, tt.wantErr)
				}
				return
			}
			if err != nil {
				t.Fatalf("resolveInstallImages() unexpected error: %v", err)
			}
			if got != tt.image {
				t.Fatalf("resolveInstallImages() material = %q, want %q", got, tt.image)
			}
		})
	}
}

func TestResolveInstallLicenseConfigRequiresMaterialRuntimeSecrets(t *testing.T) {
	secrets := InstallSecrets{LocalEncryptionKey: strings.Repeat("b", 32)}

	_, err := resolveInstallLicenseConfig(InstallRequest{}, secrets)
	if err == nil {
		t.Fatal("resolveInstallLicenseConfig() expected error")
	}
	for _, want := range []string{
		envLicenseServerURL,
		envLicenseProjectCode,
		envLicenseRequestSigningSecret,
		envLicensePublicKey,
	} {
		if !strings.Contains(err.Error(), want) {
			t.Fatalf("resolveInstallLicenseConfig() error = %q, want %q", err.Error(), want)
		}
	}
}

func TestMaterialAppEnvIncludesLicenseRuntimeConfig(t *testing.T) {
	req := InstallRequest{LicenseKey: "FLNXI-TEST"}
	secrets := InstallSecrets{
		DbUser:             "flnxi",
		DbPassword:         "db-pass",
		DbName:             "flnxi",
		MinioRootUser:      "minio",
		MinioRootPassword:  "minio-pass",
		S3Bucket:           "materials",
		NextAuthSecret:     "next-secret",
		LocalEncryptionKey: strings.Repeat("b", 32),
	}
	licenseConfig := InstallLicenseConfig{
		ServerURL:            "https://license.example.test",
		ProjectCode:          "flnxi-sucai",
		RequestSigningSecret: strings.Repeat("a", 32),
		PublicKey:            "-----BEGIN PUBLIC KEY-----\\nabc\\n-----END PUBLIC KEY-----",
		LocalEncryptionKey:   secrets.LocalEncryptionKey,
	}

	env := envSliceToMap(materialAppEnv(req, secrets, licenseConfig, "192.168.1.20", 4320))
	for _, want := range []string{
		envLicenseServerURL,
		envLicenseProjectCode,
		envLicenseRequestSigningSecret,
		envLicensePublicKey,
		envLocalEncryptionKey,
		"DATABASE_URL",
		"S3_PUBLIC_DOMAIN",
		"NEXTAUTH_URL",
		"APP_PUBLIC_URL",
	} {
		if env[want] == "" {
			t.Fatalf("materialAppEnv() missing %s in %v", want, env)
		}
	}
	if env["APP_PUBLIC_URL"] != "http://192.168.1.20:4320" {
		t.Fatalf("APP_PUBLIC_URL = %q, want host URL", env["APP_PUBLIC_URL"])
	}
}

func TestRunInstallStopsBeforePullWhenLicenseConfigMissing(t *testing.T) {
	resetInstallStatusForTest()

	fake := &fakeInstallDocker{}
	runInstall(fake, InstallRequest{MaterialImage: "registry.local:5000/material/flnxi-material:1.0.1"})

	statusMutex.RLock()
	status := *installStatus
	statusMutex.RUnlock()

	if status.Status != "failed" {
		t.Fatalf("install status = %q, want failed", status.Status)
	}
	if !strings.Contains(status.Error, "许可证配置缺失") {
		t.Fatalf("install error = %q, want license config error", status.Error)
	}
	if len(fake.pulledImages) != 0 {
		t.Fatalf("pulled images = %v, want none before license config is valid", fake.pulledImages)
	}
	if len(fake.createdContainers) != 0 {
		t.Fatalf("created containers = %v, want none before license config is valid", fake.createdContainers)
	}
}

func TestRunUpdateRestoresOldContainerAfterHealthCheckFailure(t *testing.T) {
	resetUpdateStatusForTest()

	old := &docker.ContainerInfo{
		ID:      "old-container",
		ImageID: "sha256:old",
		State:   "running",
		Config: &container.Config{
			Image:        "registry.local:5000/material/flnxi-material:1.0.0",
			User:         "1000",
			Env:          []string{"A=B"},
			ExposedPorts: nat.PortSet{"3000/tcp": struct{}{}},
		},
		HostConfig: &container.HostConfig{
			Binds:         []string{"/data/material:/app/data"},
			NetworkMode:   "flnxi-network",
			RestartPolicy: container.RestartPolicy{Name: "unless-stopped"},
			PortBindings:  nat.PortMap{"3000/tcp": []nat.PortBinding{{HostIP: "0.0.0.0", HostPort: "3000"}}},
		},
	}
	fake := &fakeUpdateDocker{container: old, healthErr: errors.New("not healthy")}

	runUpdate(fake, UpdateRequest{Image: "registry.local:5000/material/flnxi-material:2.0.0"})

	statusMutex.RLock()
	status := *updateStatus
	statusMutex.RUnlock()
	if status.Status != "failed" {
		t.Fatalf("update status = %q, want failed", status.Status)
	}
	if !strings.Contains(status.Error, "已恢复旧版本容器") {
		t.Fatalf("update error = %q, want rollback success message", status.Error)
	}
	if fake.container == nil {
		t.Fatal("container was not restored")
	}
	if fake.container.Config.Image != old.Config.Image {
		t.Fatalf("restored image = %q, want %q", fake.container.Config.Image, old.Config.Image)
	}
	if fake.container.Config.User != old.Config.User {
		t.Fatalf("restored user = %q, want %q", fake.container.Config.User, old.Config.User)
	}
	if len(fake.createdImages) != 2 {
		t.Fatalf("created images = %v, want new image then old image", fake.createdImages)
	}
	if fake.createdImages[0] != "registry.local:5000/material/flnxi-material:2.0.0" {
		t.Fatalf("first created image = %q, want new image", fake.createdImages[0])
	}
	if fake.createdImages[1] != old.Config.Image {
		t.Fatalf("second created image = %q, want old image", fake.createdImages[1])
	}
	if !contains(fake.removedIDs, "old-container") || !contains(fake.removedIDs, "new-container") {
		t.Fatalf("removed IDs = %v, want old and failed new containers removed", fake.removedIDs)
	}
}

func TestValidateModuleSyncRequestRejectsExpiredPullToken(t *testing.T) {
	req := ModuleSyncRequest{
		ModuleID: "intranet-penetration-core",
		Image:    "registry.local:5000/flnxi/premium/intranet-penetration-core:1.0.0",
		Digest:   testModuleDigest,
		PullCredential: ModulePullCredential{
			RegistryHost: "registry.local:5000",
			Repository:   "flnxi/premium/intranet-penetration-core",
			Tag:          "1.0.0",
			Digest:       testModuleDigest,
			Username:     "flnxi-module",
			PullToken:    "token",
			ExpiresAt:    time.Now().Add(-time.Minute).Format(time.RFC3339),
		},
	}

	err := validateModuleSyncRequest(req)
	if err == nil || !strings.Contains(err.Error(), "已过期") {
		t.Fatalf("validateModuleSyncRequest() error = %v, want expired token rejection", err)
	}
}

func TestValidateModuleSyncRequestRequiresDigest(t *testing.T) {
	req := validModuleSyncRequest()
	req.Digest = ""
	req.PullCredential.Digest = ""

	err := validateModuleSyncRequest(req)
	if err == nil || !strings.Contains(err.Error(), "digest 不能为空") {
		t.Fatalf("validateModuleSyncRequest() error = %v, want missing digest rejection", err)
	}
}

func TestValidateModuleSyncRequestRejectsInvalidDigest(t *testing.T) {
	req := validModuleSyncRequest()
	req.Digest = "sha256:abc"
	req.PullCredential.Digest = "sha256:abc"

	err := validateModuleSyncRequest(req)
	if err == nil || !strings.Contains(err.Error(), "sha256 格式") {
		t.Fatalf("validateModuleSyncRequest() error = %v, want invalid digest rejection", err)
	}
}

func TestModuleSyncLockRejectsConcurrentSameModule(t *testing.T) {
	resetModuleStatusesForTest()

	key := moduleContainerName(validModuleSyncRequest())
	if !tryBeginModuleSync(key) {
		t.Fatal("first module sync lock acquisition should succeed")
	}
	if tryBeginModuleSync(key) {
		t.Fatal("second module sync lock acquisition should fail")
	}
	finishModuleSync(key)
	if !tryBeginModuleSync(key) {
		t.Fatal("module sync lock should be released")
	}
	finishModuleSync(key)
}

func TestRunModuleSyncStartsContainerAfterDigestMatch(t *testing.T) {
	resetModuleStatusesForTest()

	req := validModuleSyncRequest()
	fake := &fakeModuleDocker{
		repoDigests: []string{"registry.local:5000/flnxi/premium/intranet-penetration-core@" + testModuleDigest},
	}

	runModuleSync(fake, req)

	moduleStatusesMutex.RLock()
	status := moduleStatuses[req.ModuleID]
	moduleStatusesMutex.RUnlock()

	if status.Status != "running" {
		t.Fatalf("module status = %q, want running: %+v", status.Status, status)
	}
	if len(fake.pulledImages) != 1 || fake.pulledImages[0] != req.Image {
		t.Fatalf("pulled images = %v, want %q", fake.pulledImages, req.Image)
	}
	if fake.lastAuth == "" {
		t.Fatal("expected registry auth for module pull")
	}
	if len(fake.createdImages) != 1 || fake.createdImages[0] != req.Image {
		t.Fatalf("created images = %v, want module image", fake.createdImages)
	}
	if fake.lastHostConfig == nil || !fake.lastHostConfig.ReadonlyRootfs {
		t.Fatal("module container should use a read-only root filesystem by default")
	}
	if !contains(fake.lastHostConfig.CapDrop, "ALL") {
		t.Fatalf("module container CapDrop = %v, want ALL", fake.lastHostConfig.CapDrop)
	}
	if !contains(fake.lastHostConfig.SecurityOpt, "no-new-privileges:true") {
		t.Fatalf("module container SecurityOpt = %v, want no-new-privileges", fake.lastHostConfig.SecurityOpt)
	}
}

func TestRunModuleSyncRestoresOldContainerAfterHealthFailure(t *testing.T) {
	resetModuleStatusesForTest()

	req := validModuleSyncRequest()
	old := &docker.ContainerInfo{
		ID:      "old-module",
		ImageID: "sha256:old-module",
		State:   "running",
		Config: &container.Config{
			Image: "registry.local:5000/flnxi/premium/intranet-penetration-core:0.9.0",
			Env:   []string{"OLD=1"},
		},
		HostConfig: &container.HostConfig{
			NetworkMode:   "flnxi-network",
			RestartPolicy: container.RestartPolicy{Name: "unless-stopped"},
		},
	}
	fake := &fakeModuleDocker{
		container:   old,
		repoDigests: []string{"registry.local:5000/flnxi/premium/intranet-penetration-core@" + testModuleDigest},
		healthErr:   errors.New("not healthy"),
	}

	runModuleSync(fake, req)

	moduleStatusesMutex.RLock()
	status := moduleStatuses[req.ModuleID]
	moduleStatusesMutex.RUnlock()

	if status.Status != "unhealthy" {
		t.Fatalf("module status = %q, want unhealthy", status.Status)
	}
	if fake.container == nil || fake.container.Config.Image != old.Config.Image {
		t.Fatalf("restored container = %+v, want old image %q", fake.container, old.Config.Image)
	}
	if !contains(fake.removedIDs, "old-module") || !contains(fake.removedIDs, "new-module") {
		t.Fatalf("removed IDs = %v, want old and failed new module", fake.removedIDs)
	}
}

const testModuleDigest = "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"

func validModuleSyncRequest() ModuleSyncRequest {
	return ModuleSyncRequest{
		ModuleID: "intranet-penetration-core",
		Image:    "registry.local:5000/flnxi/premium/intranet-penetration-core:1.0.0",
		Digest:   testModuleDigest,
		PullCredential: ModulePullCredential{
			RegistryHost: "registry.local:5000",
			Repository:   "flnxi/premium/intranet-penetration-core",
			Tag:          "1.0.0",
			Digest:       testModuleDigest,
			Username:     "flnxi-module",
			PullToken:    "token",
			ExpiresAt:    time.Now().Add(time.Hour).Format(time.RFC3339),
		},
	}
}

func resetUpdateStatusForTest() {
	statusMutex.Lock()
	updateStatus = &UpdateStatus{Status: "running", Step: "初始化", Progress: 0}
	statusMutex.Unlock()
}

func resetInstallStatusForTest() {
	statusMutex.Lock()
	installStatus = &InstallStatus{Status: "running", Step: "初始化", Progress: 0}
	statusMutex.Unlock()
}

func resetModuleStatusesForTest() {
	moduleStatusesMutex.Lock()
	moduleStatuses = map[string]ModuleRuntimeStatus{}
	moduleStatusesMutex.Unlock()

	moduleSyncLocksMutex.Lock()
	moduleSyncLocks = map[string]bool{}
	moduleSyncLocksMutex.Unlock()
}

func envSliceToMap(values []string) map[string]string {
	result := make(map[string]string, len(values))
	for _, value := range values {
		parts := strings.SplitN(value, "=", 2)
		if len(parts) != 2 {
			continue
		}
		result[parts[0]] = parts[1]
	}
	return result
}

type fakeInstallDocker struct {
	pulledImages      []string
	createdContainers []string
}

func (f *fakeInstallDocker) PullImageWithProgress(ctx context.Context, image string, authStr string, onProgress func(current, total int64, status string)) error {
	f.pulledImages = append(f.pulledImages, image)
	if onProgress != nil {
		onProgress(1, 1, "done")
	}
	return nil
}

func (f *fakeInstallDocker) ExecOnHost(ctx context.Context, cmd string) error {
	return nil
}

func (f *fakeInstallDocker) CreateNetworkIfNotExists(ctx context.Context, name string) error {
	return nil
}

func (f *fakeInstallDocker) CreateAndStartContainer(ctx context.Context, config *container.Config, hostConfig *container.HostConfig, name string) (string, error) {
	f.createdContainers = append(f.createdContainers, name)
	return name + "-id", nil
}

type fakeUpdateDocker struct {
	container     *docker.ContainerInfo
	healthErr     error
	pulledImages  []string
	createdImages []string
	removedIDs    []string
	stoppedIDs    []string
	startedIDs    []string
}

type fakeModuleDocker struct {
	container      *docker.ContainerInfo
	healthErr      error
	repoDigests    []string
	pulledImages   []string
	createdImages  []string
	removedIDs     []string
	stoppedIDs     []string
	startedIDs     []string
	lastAuth       string
	lastHostConfig *container.HostConfig
}

func (f *fakeModuleDocker) GetContainer(ctx context.Context, name string) (*docker.ContainerInfo, error) {
	return f.container, nil
}

func (f *fakeModuleDocker) PullImage(ctx context.Context, image string, authStr string) error {
	f.pulledImages = append(f.pulledImages, image)
	f.lastAuth = authStr
	return nil
}

func (f *fakeModuleDocker) GetImageRepoDigests(ctx context.Context, image string) ([]string, error) {
	return f.repoDigests, nil
}

func (f *fakeModuleDocker) StopContainer(ctx context.Context, containerID string) error {
	f.stoppedIDs = append(f.stoppedIDs, containerID)
	if f.container != nil && f.container.ID == containerID {
		f.container.State = "exited"
	}
	return nil
}

func (f *fakeModuleDocker) RemoveContainer(ctx context.Context, containerID string) error {
	f.removedIDs = append(f.removedIDs, containerID)
	if f.container != nil && f.container.ID == containerID {
		f.container = nil
	}
	return nil
}

func (f *fakeModuleDocker) StartContainer(ctx context.Context, containerID string) error {
	f.startedIDs = append(f.startedIDs, containerID)
	if f.container != nil && f.container.ID == containerID {
		f.container.State = "running"
	}
	return nil
}

func (f *fakeModuleDocker) CreateAndStartContainer(ctx context.Context, config *container.Config, hostConfig *container.HostConfig, name string) (string, error) {
	f.createdImages = append(f.createdImages, config.Image)
	f.lastHostConfig = hostConfig
	id := "new-module"
	if strings.Contains(config.Image, ":0.9.0") {
		id = "restored-module"
	}
	f.container = &docker.ContainerInfo{
		ID:         id,
		ImageID:    "sha256:" + id,
		State:      "running",
		Config:     config,
		HostConfig: hostConfig,
	}
	return id, nil
}

func (f *fakeModuleDocker) WaitForHealthy(ctx context.Context, containerID string, timeout time.Duration) error {
	return f.healthErr
}

func (f *fakeUpdateDocker) GetContainer(ctx context.Context, name string) (*docker.ContainerInfo, error) {
	return f.container, nil
}

func (f *fakeUpdateDocker) PullImage(ctx context.Context, image string, authStr string) error {
	f.pulledImages = append(f.pulledImages, image)
	return nil
}

func (f *fakeUpdateDocker) StopContainer(ctx context.Context, containerID string) error {
	f.stoppedIDs = append(f.stoppedIDs, containerID)
	if f.container != nil && f.container.ID == containerID {
		f.container.State = "exited"
	}
	return nil
}

func (f *fakeUpdateDocker) RemoveContainer(ctx context.Context, containerID string) error {
	f.removedIDs = append(f.removedIDs, containerID)
	if f.container != nil && f.container.ID == containerID {
		f.container = nil
	}
	return nil
}

func (f *fakeUpdateDocker) StartContainer(ctx context.Context, containerID string) error {
	f.startedIDs = append(f.startedIDs, containerID)
	if f.container != nil && f.container.ID == containerID {
		f.container.State = "running"
	}
	return nil
}

func (f *fakeUpdateDocker) CreateAndStartContainer(ctx context.Context, config *container.Config, hostConfig *container.HostConfig, name string) (string, error) {
	f.createdImages = append(f.createdImages, config.Image)
	id := "new-container"
	if strings.Contains(config.Image, ":1.0.0") {
		id = "restored-container"
	}
	f.container = &docker.ContainerInfo{
		ID:         id,
		ImageID:    "sha256:" + id,
		State:      "running",
		Config:     config,
		HostConfig: hostConfig,
	}
	return id, nil
}

func (f *fakeUpdateDocker) WaitForHealthy(ctx context.Context, containerID string, timeout time.Duration) error {
	return f.healthErr
}

func (f *fakeUpdateDocker) RemoveImage(ctx context.Context, imageID string) error {
	return nil
}

func (f *fakeUpdateDocker) PruneImages(ctx context.Context) error {
	return nil
}

func contains(values []string, want string) bool {
	for _, value := range values {
		if value == want {
			return true
		}
	}
	return false
}

var _ updateDockerClient = (*fakeUpdateDocker)(nil)
var _ installDockerClient = (*fakeInstallDocker)(nil)
var _ moduleDockerClient = (*fakeModuleDocker)(nil)