pulse/pkg/licensing/persistence_test.go

package licensing

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"io"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func TestPersistence(t *testing.T) {
	// Create a temporary directory for config
	tmpDir, err := os.MkdirTemp("", "pulse-license-test-*")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	p, err := NewPersistence(tmpDir)
	if err != nil {
		t.Fatalf("Failed to create persistence: %v", err)
	}

	testLicenseKey := "test-license-key-123"

	t.Run("Save and Load", func(t *testing.T) {
		err := p.Save(testLicenseKey)
		if err != nil {
			t.Fatalf("Failed to save license: %v", err)
		}

		if !p.Exists() {
			t.Error("License file should exist")
		}

		loadedKey, err := p.Load()
		if err != nil {
			t.Fatalf("Failed to load license: %v", err)
		}

		if loadedKey != testLicenseKey {
			t.Errorf("Expected license key %s, got %s", testLicenseKey, loadedKey)
		}
	})

	t.Run("Save and Load with Grace Period", func(t *testing.T) {
		gracePeriodEnd := time.Now().Add(7 * 24 * time.Hour).Unix()
		err := p.SaveWithGracePeriod(testLicenseKey, &gracePeriodEnd)
		if err != nil {
			t.Fatalf("Failed to save license with grace period: %v", err)
		}

		persisted, err := p.LoadWithMetadata()
		if err != nil {
			t.Fatalf("Failed to load license with metadata: %v", err)
		}

		if persisted.LicenseKey != testLicenseKey {
			t.Errorf("Expected license key %s, got %s", testLicenseKey, persisted.LicenseKey)
		}

		if persisted.GracePeriodEnd == nil || *persisted.GracePeriodEnd != gracePeriodEnd {
			t.Errorf("Expected grace period end %v, got %v", gracePeriodEnd, persisted.GracePeriodEnd)
		}
	})

	t.Run("Load non-existent", func(t *testing.T) {
		tmpDirEmpty, _ := os.MkdirTemp("", "pulse-license-test-empty-*")
		defer os.RemoveAll(tmpDirEmpty)

		pEmpty, _ := NewPersistence(tmpDirEmpty)
		key, err := pEmpty.Load()
		if err != nil {
			t.Fatalf("Expected no error for non-existent license, got %v", err)
		}
		if key != "" {
			t.Errorf("Expected empty key, got %s", key)
		}
	})

	t.Run("Delete", func(t *testing.T) {
		err := p.Save(testLicenseKey)
		if err != nil {
			t.Fatalf("Failed to save license: %v", err)
		}

		if !p.Exists() {
			t.Fatal("License should exist before delete")
		}

		err = p.Delete()
		if err != nil {
			t.Fatalf("Failed to delete license: %v", err)
		}

		if p.Exists() {
			t.Error("License should not exist after delete")
		}
	})

	t.Run("Encryption check", func(t *testing.T) {
		err := p.Save(testLicenseKey)
		if err != nil {
			t.Fatalf("Failed to save license: %v", err)
		}

		// Read the file directly - it should be base64 encoded and encrypted
		licensePath := filepath.Join(tmpDir, LicenseFileName)
		data, err := os.ReadFile(licensePath)
		if err != nil {
			t.Fatalf("Failed to read license file: %v", err)
		}

		if string(data) == testLicenseKey {
			t.Error("License file should be encrypted, not raw text")
		}

		// Ensure it's not JSON either in raw form
		if data[0] == '{' {
			t.Error("License file should be encrypted, not raw JSON")
		}
	})

	t.Run("Decrypt with wrong key material", func(t *testing.T) {
		err := p.Save(testLicenseKey)
		if err != nil {
			t.Fatalf("Failed to save license: %v", err)
		}

		// Create a new persistence with different encryption key
		pWrong := &Persistence{
			configDir:     tmpDir,
			encryptionKey: "different-encryption-key",
			machineID:     "different-machine-id",
		}

		_, err = pWrong.Load()
		if err == nil {
			t.Error("Expected error when decrypting with wrong key material")
		}
	})

	t.Run("Backwards compatibility with machine-id", func(t *testing.T) {
		// This tests the scenario where a license was saved with machine-id
		// (old behavior before persistent key feature) and we're now loading it
		// with a different primary key but same machine-id as fallback

		tmpDirCompat, _ := os.MkdirTemp("", "pulse-license-compat-*")
		defer os.RemoveAll(tmpDirCompat)

		testKey := "compat-test-key"
		machineID := "test-machine-id-12345"

		// Simulate old behavior: directly encrypt with machine-id (without calling Save
		// which would create a persistent key). This is what old installations have.
		pOld := &Persistence{
			configDir:     tmpDirCompat,
			encryptionKey: machineID, // Old behavior used machine-id as encryption key
			machineID:     machineID,
		}

		// Manually encrypt and save without creating persistent key
		persisted := PersistedLicense{LicenseKey: testKey}
		jsonData, _ := json.Marshal(persisted)
		encrypted, err := encryptWithKey(jsonData, pOld.deriveLegacyKeyFrom(machineID))
		if err != nil {
			t.Fatalf("Failed to encrypt license: %v", err)
		}

		// Write directly to simulate old installation
		_ = os.MkdirAll(tmpDirCompat, 0700)
		licensePath := filepath.Join(tmpDirCompat, LicenseFileName)
		encoded := base64.StdEncoding.EncodeToString(encrypted)
		_ = os.WriteFile(licensePath, []byte(encoded), 0600)

		// Verify no persistent key file exists (simulating old installation)
		keyPath := filepath.Join(tmpDirCompat, PersistentKeyFileName)
		if _, err := os.Stat(keyPath); err == nil {
			t.Fatal("Persistent key should not exist for this test")
		}

		// Now try to load with a new primary key but same machine-id as fallback
		// This simulates what happens when a Docker user upgrades and gets a
		// persistent key file, but their old license was encrypted with machine-id
		pNew := &Persistence{
			configDir:     tmpDirCompat,
			encryptionKey: "new-persistent-key", // Different primary key (simulating new container)
			machineID:     machineID,            // Same machine-id for fallback
		}

		loaded, err := pNew.LoadWithMetadata()
		if err != nil {
			t.Fatalf("Failed to load license with machine-id fallback: %v\n"+
				"This means backwards compatibility is broken for existing Docker users", err)
		}

		if loaded.LicenseKey != testKey {
			t.Errorf("Expected license key %s, got %s", testKey, loaded.LicenseKey)
		}
	})

	t.Run("Backwards compatibility with v5.0.0 raw key materials", func(t *testing.T) {
		// v5.0.0 sealed license.enc with sha256("pulse-license-" + material)
		// where material was the raw machine-id file content (trailing
		// newline included), the hostname fallback, or the fixed dev
		// fallback. Licenses activated on v5.0.0 were never re-saved, so v6
		// must be able to open files sealed under each of those materials.
		rawMachineID := "0123456789abcdef0123456789abcdef\n"
		hostname := "pulse-host-01"

		for name, material := range map[string]string{
			"raw machine-id with trailing newline": rawMachineID,
			"hostname fallback":                    hostname,
			"dev fallback":                         legacyV5MachineIDFallback,
		} {
			t.Run(name, func(t *testing.T) {
				tmpDirCompat := t.TempDir()
				testKey := "v5-lifetime-key"

				sealer := &Persistence{configDir: tmpDirCompat}
				persisted := PersistedLicense{LicenseKey: testKey}
				jsonData, _ := json.Marshal(persisted)
				encrypted, err := encryptWithKey(jsonData, sealer.deriveLegacyKeyFrom(material))
				if err != nil {
					t.Fatalf("Failed to encrypt license: %v", err)
				}
				licensePath := filepath.Join(tmpDirCompat, LicenseFileName)
				encoded := base64.StdEncoding.EncodeToString(encrypted)
				if err := os.WriteFile(licensePath, []byte(encoded), 0600); err != nil {
					t.Fatalf("Failed to write license file: %v", err)
				}

				// v6 trims machine-id and generates a fresh persistent key,
				// so only the raw legacy materials can open the v5.0.0 file.
				pNew := &Persistence{
					configDir:          tmpDirCompat,
					encryptionKey:      "new-persistent-key",
					machineID:          strings.TrimSpace(rawMachineID),
					legacyKeyMaterials: []string{rawMachineID, hostname, legacyV5MachineIDFallback},
				}

				loaded, err := pNew.LoadWithMetadata()
				if err != nil {
					t.Fatalf("Failed to load v5.0.0-era license with raw key material %q: %v\n"+
						"This breaks migration for licenses activated on v5.0.0 before .license-key existed", name, err)
				}
				if loaded.LicenseKey != testKey {
					t.Errorf("Expected license key %s, got %s", testKey, loaded.LicenseKey)
				}
			})
		}
	})

	t.Run("NewPersistence collects v5.0.0 legacy key materials", func(t *testing.T) {
		p, err := NewPersistence(t.TempDir())
		if err != nil {
			t.Fatalf("Failed to create persistence: %v", err)
		}

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

		if !contains(p.legacyKeyMaterials, legacyV5MachineIDFallback) {
			t.Errorf("legacyKeyMaterials missing dev fallback %q: %v", legacyV5MachineIDFallback, p.legacyKeyMaterials)
		}
		if hostname, err := os.Hostname(); err == nil && hostname != "" {
			if !contains(p.legacyKeyMaterials, hostname) {
				t.Errorf("legacyKeyMaterials missing hostname %q: %v", hostname, p.legacyKeyMaterials)
			}
		}
		for _, path := range machineIDPaths {
			data, err := os.ReadFile(path)
			if err != nil || len(data) == 0 {
				continue
			}
			if !contains(p.legacyKeyMaterials, string(data)) {
				t.Errorf("legacyKeyMaterials missing raw content of %s: %v", path, p.legacyKeyMaterials)
			}
		}
	})

	t.Run("Backwards compatibility with legacy persistent-key derivation", func(t *testing.T) {
		tmpDirCompat, _ := os.MkdirTemp("", "pulse-license-persistent-compat-*")
		defer os.RemoveAll(tmpDirCompat)

		testKey := "compat-persistent-key"
		persistentKey := "persisted-random-key-material"

		if err := os.WriteFile(filepath.Join(tmpDirCompat, PersistentKeyFileName), []byte(persistentKey), 0600); err != nil {
			t.Fatalf("Failed to write persistent key file: %v", err)
		}

		pOld := &Persistence{
			configDir:     tmpDirCompat,
			encryptionKey: persistentKey,
		}
		persisted := PersistedLicense{LicenseKey: testKey}
		jsonData, _ := json.Marshal(persisted)
		encrypted, err := encryptWithKey(jsonData, pOld.deriveLegacyKeyFrom(persistentKey))
		if err != nil {
			t.Fatalf("Failed to encrypt license: %v", err)
		}

		licensePath := filepath.Join(tmpDirCompat, LicenseFileName)
		encoded := base64.StdEncoding.EncodeToString(encrypted)
		if err := os.WriteFile(licensePath, []byte(encoded), 0600); err != nil {
			t.Fatalf("Failed to write legacy persistent-key license: %v", err)
		}

		pNew, err := NewPersistence(tmpDirCompat)
		if err != nil {
			t.Fatalf("Failed to create persistence: %v", err)
		}

		loaded, err := pNew.LoadWithMetadata()
		if err != nil {
			t.Fatalf("Failed to load license with legacy persistent-key derivation: %v", err)
		}
		if loaded.LicenseKey != testKey {
			t.Errorf("Expected license key %s, got %s", testKey, loaded.LicenseKey)
		}
	})

	t.Run("plaintext license file rewrites encrypted storage on load", func(t *testing.T) {
		plainDir, err := os.MkdirTemp("", "pulse-license-plaintext-*")
		if err != nil {
			t.Fatalf("Failed to create temp dir: %v", err)
		}
		defer os.RemoveAll(plainDir)

		pPlain, err := NewPersistence(plainDir)
		if err != nil {
			t.Fatalf("Failed to create persistence: %v", err)
		}

		gracePeriodEnd := time.Now().Add(24 * time.Hour).Unix()
		plaintext := PersistedLicense{
			LicenseKey:     "pit-license-plaintext-test",
			GracePeriodEnd: &gracePeriodEnd,
		}

		raw, err := json.Marshal(plaintext)
		if err != nil {
			t.Fatalf("Failed to marshal plaintext license: %v", err)
		}

		licensePath := filepath.Join(plainDir, LicenseFileName)
		if err := os.WriteFile(licensePath, raw, 0600); err != nil {
			t.Fatalf("Failed to write plaintext license file: %v", err)
		}

		loaded, err := pPlain.LoadWithMetadata()
		if err != nil {
			t.Fatalf("LoadWithMetadata plaintext migration failed: %v", err)
		}

		if loaded.LicenseKey != plaintext.LicenseKey {
			t.Fatalf("LicenseKey = %q, want %q", loaded.LicenseKey, plaintext.LicenseKey)
		}
		if loaded.GracePeriodEnd == nil || *loaded.GracePeriodEnd != gracePeriodEnd {
			t.Fatalf("GracePeriodEnd = %v, want %v", loaded.GracePeriodEnd, gracePeriodEnd)
		}

		rewritten, err := os.ReadFile(licensePath)
		if err != nil {
			t.Fatalf("Failed to read rewritten license file: %v", err)
		}
		if bytes.Equal(rewritten, raw) {
			t.Fatal("expected plaintext license file to be rewritten encrypted")
		}
		if bytes.Contains(rewritten, []byte(plaintext.LicenseKey)) {
			t.Fatal("license key should not remain in plaintext after migration rewrite")
		}
	})
}

func encryptWithKey(plaintext []byte, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}

	return gcm.Seal(nonce, nonce, plaintext, nil), nil
}

func TestPersistenceEnforcesOwnerOnlyPermissions(t *testing.T) {
	tmpDir := t.TempDir()
	if err := os.Chmod(tmpDir, 0755); err != nil {
		t.Fatalf("failed to relax temp dir perms: %v", err)
	}

	p, err := NewPersistence(tmpDir)
	if err != nil {
		t.Fatalf("Failed to create persistence: %v", err)
	}
	if err := p.Save("owner-only-perm-test"); err != nil {
		t.Fatalf("Save failed: %v", err)
	}

	dirInfo, err := os.Stat(tmpDir)
	if err != nil {
		t.Fatalf("failed to stat config dir: %v", err)
	}
	if got := dirInfo.Mode().Perm(); got != 0700 {
		t.Fatalf("config dir perms = %o, want 700", got)
	}

	for _, file := range []string{LicenseFileName, PersistentKeyFileName} {
		path := filepath.Join(tmpDir, file)
		info, err := os.Stat(path)
		if err != nil {
			t.Fatalf("failed to stat %s: %v", file, err)
		}
		if got := info.Mode().Perm(); got != 0600 {
			t.Fatalf("%s perms = %o, want 600", file, got)
		}
	}
}

func TestNewPersistence_HardensExistingPersistentKeyFileOnLoad(t *testing.T) {
	tmpDir := t.TempDir()
	keyPath := filepath.Join(tmpDir, PersistentKeyFileName)
	if err := os.WriteFile(keyPath, []byte("persistent-key-material"), 0o644); err != nil {
		t.Fatalf("failed to write persistent key: %v", err)
	}
	if err := os.Chmod(keyPath, 0o644); err != nil {
		t.Fatalf("failed to loosen persistent key perms: %v", err)
	}

	p, err := NewPersistence(tmpDir)
	if err != nil {
		t.Fatalf("NewPersistence() error: %v", err)
	}
	if p.encryptionKey != "persistent-key-material" {
		t.Fatalf("encryptionKey = %q, want persistent key contents", p.encryptionKey)
	}

	info, err := os.Stat(keyPath)
	if err != nil {
		t.Fatalf("stat %s: %v", keyPath, err)
	}
	if got := info.Mode().Perm(); got != 0o600 {
		t.Fatalf("%s perms = %o, want 600", filepath.Base(keyPath), got)
	}
}

func TestResolvePersistencePathCanonicalizesConfigDir(t *testing.T) {
	baseDir := t.TempDir()
	inputDir := "  " + filepath.Join(baseDir, "nested", "..", "licensing") + string(os.PathSeparator) + ".  "

	path, err := resolvePersistencePath(inputDir, LicenseFileName)
	if err != nil {
		t.Fatalf("resolvePersistencePath() error = %v", err)
	}

	wantDir := filepath.Clean(filepath.Join(baseDir, "nested", "..", "licensing") + string(os.PathSeparator) + ".")
	if path != filepath.Join(wantDir, LicenseFileName) {
		t.Fatalf("resolvePersistencePath() = %q, want %q", path, filepath.Join(wantDir, LicenseFileName))
	}
}

func TestResolvePersistencePathRejectsBlankConfigDir(t *testing.T) {
	if _, err := resolvePersistencePath(" \t ", LicenseFileName); err == nil {
		t.Fatal("expected blank config dir to be rejected")
	}
}

func TestNewPersistenceRejectsSymlinkPersistentKey(t *testing.T) {
	tmpDir := t.TempDir()
	target := filepath.Join(tmpDir, "key-target")
	if err := os.WriteFile(target, []byte("test-key"), 0600); err != nil {
		t.Fatalf("failed to write target file: %v", err)
	}

	keyPath := filepath.Join(tmpDir, PersistentKeyFileName)
	if err := os.Symlink(target, keyPath); err != nil {
		t.Skipf("symlink unsupported on this platform: %v", err)
	}

	_, err := NewPersistence(tmpDir)
	if err == nil {
		t.Fatal("expected error for symlink persistent key path")
	}
	if !strings.Contains(err.Error(), "symlink") {
		t.Fatalf("expected symlink error, got: %v", err)
	}
}