chore: import upstream snapshot with attribution

2026-07-13 12:08:39 +08:00
commit a0df89c693
5252 changed files with 523444 additions and 0 deletions
@@ -0,0 +1,334 @@
+package postinit
+
+import (
+	"cmp"
+	"context"
+	"fmt"
+	"slices"
+
+	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/client"
+	portainer "github.com/portainer/portainer/api"
+	"github.com/portainer/portainer/api/dataservices"
+	dockerClient "github.com/portainer/portainer/api/docker/client"
+	"github.com/portainer/portainer/api/internal/endpointutils"
+	"github.com/portainer/portainer/api/internal/registryutils"
+	"github.com/portainer/portainer/api/kubernetes/cli"
+	"github.com/portainer/portainer/api/logs"
+	"github.com/portainer/portainer/api/pendingactions/actions"
+	"github.com/portainer/portainer/pkg/endpoints"
+
+	"github.com/rs/zerolog/log"
+)
+
+type PostInitMigrator struct {
+	kubeFactory        *cli.ClientFactory
+	dockerFactory      *dockerClient.ClientFactory
+	dataStore          dataservices.DataStore
+	assetsPath         string
+	kubernetesDeployer portainer.KubernetesDeployer
+}
+
+func NewPostInitMigrator(
+	kubeFactory *cli.ClientFactory,
+	dockerFactory *dockerClient.ClientFactory,
+	dataStore dataservices.DataStore,
+	assetsPath string,
+	kubernetesDeployer portainer.KubernetesDeployer,
+) *PostInitMigrator {
+	return &PostInitMigrator{
+		kubeFactory:        kubeFactory,
+		dockerFactory:      dockerFactory,
+		dataStore:          dataStore,
+		assetsPath:         assetsPath,
+		kubernetesDeployer: kubernetesDeployer,
+	}
+}
+
+// PostInitMigrate will run all post-init migrations, which require docker/kube clients for all edge or non-edge environments
+func (postInitMigrator *PostInitMigrator) PostInitMigrate() error {
+	var environments []portainer.Endpoint
+
+	if err := postInitMigrator.dataStore.UpdateTx(func(tx dataservices.DataStoreTx) error {
+		var err error
+		if environments, err = tx.Endpoint().ReadAll(func(endpoint portainer.Endpoint) bool {
+			return endpoints.HasDirectConnectivity(&endpoint)
+		}); err != nil {
+			return fmt.Errorf("failed to retrieve environments: %w", err)
+		}
+
+		var pendingActions []portainer.PendingAction
+		if pendingActions, err = tx.PendingActions().ReadAll(func(action portainer.PendingAction) bool {
+			return action.Action == actions.PostInitMigrateEnvironment
+		}); err != nil {
+			return fmt.Errorf("failed to retrieve pending actions: %w", err)
+		}
+
+		// Sort for the binary search in createPostInitMigrationPendingAction()
+		slices.SortFunc(pendingActions, func(a, b portainer.PendingAction) int {
+			return cmp.Compare(a.EndpointID, b.EndpointID)
+		})
+
+		for _, environment := range environments {
+			if !endpoints.IsEdgeEndpoint(&environment) {
+				continue
+			}
+
+			// Edge environments will run after the server starts, in pending actions
+			log.Info().
+				Int("endpoint_id", int(environment.ID)).
+				Msg("adding pending action 'PostInitMigrateEnvironment' for environment")
+
+			if err := postInitMigrator.createPostInitMigrationPendingAction(tx, environment.ID, pendingActions); err != nil {
+				log.Error().
+					Err(err).
+					Int("endpoint_id", int(environment.ID)).
+					Msg("error creating pending action for environment")
+			}
+		}
+
+		return err
+	}); err != nil {
+		log.Error().Err(err).Msg("error running post-init migrations")
+
+		return err
+	}
+
+	for _, environment := range environments {
+		if endpoints.IsEdgeEndpoint(&environment) {
+			continue
+		}
+
+		// Non-edge environments will run before the server starts.
+		if err := postInitMigrator.MigrateEnvironment(&environment); err != nil {
+			log.Error().
+				Err(err).
+				Int("endpoint_id", int(environment.ID)).
+				Msg("error running post-init migrations for non-edge environment")
+		}
+	}
+
+	return nil
+}
+
+// try to create a post init migration pending action. If it already exists, do nothing
+// this function exists for readability, not reusability
+// pending actions must be passed in ascending order by endpoint ID
+func (postInitMigrator *PostInitMigrator) createPostInitMigrationPendingAction(tx dataservices.DataStoreTx, environmentID portainer.EndpointID, pendingActions []portainer.PendingAction) error {
+	action := portainer.PendingAction{
+		EndpointID: environmentID,
+		Action:     actions.PostInitMigrateEnvironment,
+	}
+
+	if _, found := slices.BinarySearchFunc(pendingActions, environmentID, func(e portainer.PendingAction, id portainer.EndpointID) int {
+		return cmp.Compare(e.EndpointID, id)
+	}); found {
+		log.Debug().
+			Str("action", action.Action).
+			Int("endpoint_id", int(action.EndpointID)).
+			Msg("pending action already exists for environment, skipping...")
+
+		return nil
+	}
+
+	return tx.PendingActions().Create(&action)
+}
+
+// MigrateEnvironment runs migrations on a single environment
+func (migrator *PostInitMigrator) MigrateEnvironment(environment *portainer.Endpoint) error {
+	log.Info().
+		Int("endpoint_id", int(environment.ID)).
+		Msg("executing post init migration for environment")
+
+	switch {
+	case endpointutils.IsKubernetesEndpoint(environment):
+		// get the kubeclient for the environment, and skip all kube migrations if there's an error
+		kubeclient, err := migrator.kubeFactory.GetPrivilegedKubeClient(environment)
+		if err != nil {
+			log.Error().
+				Err(err).
+				Int("endpoint_id", int(environment.ID)).
+				Msg("error creating kubeclient for environment")
+
+			return err
+		}
+
+		// If one environment fails, it is logged and the next migration runs. The error is returned at the end and handled by pending actions
+		var latestErr error
+		kubernetesMigrations := []func() error{
+			func() error { return migrator.MigrateIngresses(*environment, kubeclient) },
+			func() error { return migrator.MigrateRegistrySASecrets(*environment, kubeclient) },
+		}
+
+		for _, migration := range kubernetesMigrations {
+			if err := migration(); err != nil {
+				latestErr = err
+			}
+		}
+
+		return latestErr
+	case endpointutils.IsDockerEndpoint(environment):
+		// get the docker client for the environment, and skip all docker migrations if there's an error
+		dockerClient, err := migrator.dockerFactory.CreateClient(environment, "", nil)
+		if err != nil {
+			log.Error().
+				Err(err).
+				Int("endpoint_id", int(environment.ID)).
+				Msg("error creating docker client for environment")
+
+			return err
+		}
+		defer logs.CloseAndLogErr(dockerClient)
+
+		if err := migrator.MigrateGPUs(*environment, dockerClient); err != nil {
+			log.Error().
+				Err(err).
+				Int("endpoint_id", int(environment.ID)).
+				Msg("error migrating GPUs for environment")
+
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (migrator *PostInitMigrator) MigrateRegistrySASecrets(environment portainer.Endpoint, kubeclient *cli.KubeClient) error {
+	if !environment.PostInitMigrations.MigrateRegistrySASecrets {
+		return nil
+	}
+
+	log.Debug().
+		Int("endpoint_id", int(environment.ID)).
+		Msg("migrating registry SA secrets for environment")
+
+	return migrator.dataStore.UpdateTx(func(tx dataservices.DataStoreTx) error {
+		env, err := tx.Endpoint().Endpoint(environment.ID)
+		if err != nil {
+			return err
+		}
+
+		if !env.PostInitMigrations.MigrateRegistrySASecrets {
+			return nil
+		}
+
+		registries, err := tx.Registry().ReadAll()
+		if err != nil {
+			return err
+		}
+
+		for _, registry := range registries {
+			access, ok := registry.RegistryAccesses[env.ID]
+			if !ok || len(access.Namespaces) == 0 {
+				continue
+			}
+
+			secretName := registryutils.RegistrySecretName(registry.ID)
+			for _, namespace := range access.Namespaces {
+				if err := kubeclient.AddImagePullSecretToServiceAccount(namespace, "default", secretName); err != nil {
+					log.Warn().
+						Err(err).
+						Int("endpoint_id", int(env.ID)).
+						Str("namespace", namespace).
+						Str("secret", secretName).
+						Msg("failed to add imagePullSecret to service account during registry SA secret migration")
+				}
+			}
+		}
+
+		env.PostInitMigrations.MigrateRegistrySASecrets = false
+		return tx.Endpoint().UpdateEndpoint(env.ID, env)
+	})
+}
+
+func (migrator *PostInitMigrator) MigrateIngresses(environment portainer.Endpoint, kubeclient *cli.KubeClient) error {
+	// Early exit if we do not need to migrate!
+	if !environment.PostInitMigrations.MigrateIngresses {
+		return nil
+	}
+
+	log.Debug().
+		Int("endpoint_id", int(environment.ID)).
+		Msg("migrating ingresses for environment")
+
+	if err := migrator.kubeFactory.MigrateEndpointIngresses(&environment, migrator.dataStore, kubeclient); err != nil {
+		log.Error().
+			Err(err).
+			Int("endpoint_id", int(environment.ID)).
+			Msg("error migrating ingresses for environment")
+
+		return err
+	}
+
+	return nil
+}
+
+// MigrateGPUs will check all docker endpoints for containers with GPUs and set EnableGPUManagement to true if any are found
+// If there's an error getting the containers, we'll log it and move on
+func (migrator *PostInitMigrator) MigrateGPUs(e portainer.Endpoint, dockerClient *client.Client) error {
+	return migrator.dataStore.UpdateTx(func(tx dataservices.DataStoreTx) error {
+		environment, err := tx.Endpoint().Endpoint(e.ID)
+		if err != nil {
+			log.Error().
+				Err(err).
+				Int("endpoint_id", int(e.ID)).
+				Msg("error getting environment")
+
+			return err
+		}
+
+		// Early exit if we do not need to migrate!
+		if !environment.PostInitMigrations.MigrateGPUs {
+			return nil
+		}
+
+		log.Debug().
+			Int("endpoint_id", int(e.ID)).
+			Msg("migrating GPUs for environment")
+
+		// Get all containers
+		containers, err := dockerClient.ContainerList(context.Background(), container.ListOptions{All: true})
+		if err != nil {
+			log.Error().
+				Err(err).
+				Int("endpoint_id", int(environment.ID)).
+				Msg("failed to list containers for environment")
+
+			return err
+		}
+
+		// Check for a gpu on each container. If even one GPU is found, set EnableGPUManagement to true for the whole environment
+	containersLoop:
+		for _, container := range containers {
+			// https://www.sobyte.net/post/2022-10/go-docker/ has nice documentation on the docker client with GPUs
+			containerDetails, err := dockerClient.ContainerInspect(context.Background(), container.ID)
+			if err != nil {
+				log.Error().Err(err).Msg("failed to inspect container")
+
+				continue
+			}
+
+			deviceRequests := containerDetails.HostConfig.DeviceRequests
+			for _, deviceRequest := range deviceRequests {
+				if deviceRequest.Driver == "nvidia" {
+					environment.EnableGPUManagement = true
+
+					break containersLoop
+				}
+			}
+		}
+
+		// Set the MigrateGPUs flag to false so we don't run this again
+		environment.PostInitMigrations.MigrateGPUs = false
+		if err := tx.Endpoint().UpdateEndpoint(environment.ID, environment); err != nil {
+			log.Error().
+				Err(err).
+				Int("endpoint_id", int(environment.ID)).
+				Msg("error updating EnableGPUManagement flag for environment")
+
+			return err
+		}
+
+		return nil
+	})
+}
@@ -0,0 +1,176 @@
+package postinit
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	portainer "github.com/portainer/portainer/api"
+	"github.com/portainer/portainer/api/datastore"
+	"github.com/portainer/portainer/api/pendingactions/actions"
+
+	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/client"
+	"github.com/segmentio/encoding/json"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMigrateGPUs(t *testing.T) {
+	t.Parallel()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if strings.HasSuffix(r.URL.Path, "/containers/json") {
+			containerSummary := []container.Summary{{ID: "container1"}}
+
+			if err := json.NewEncoder(w).Encode(containerSummary); err != nil {
+				w.WriteHeader(http.StatusInternalServerError)
+			}
+
+			return
+		}
+
+		container := container.InspectResponse{
+			ContainerJSONBase: &container.ContainerJSONBase{
+				ID: "container1",
+				HostConfig: &container.HostConfig{
+					Resources: container.Resources{
+						DeviceRequests: []container.DeviceRequest{
+							{Driver: "nvidia"},
+						},
+					},
+				},
+			},
+		}
+
+		if err := json.NewEncoder(w).Encode(container); err != nil {
+			w.WriteHeader(http.StatusInternalServerError)
+		}
+	}))
+	defer srv.Close()
+
+	_, store := datastore.MustNewTestStore(t, true, false)
+
+	migrator := &PostInitMigrator{dataStore: store}
+
+	dockerCli, err := client.NewClientWithOpts(client.WithHost(srv.URL), client.WithHTTPClient(http.DefaultClient))
+	require.NoError(t, err)
+
+	// Nonexistent endpoint
+
+	err = migrator.MigrateGPUs(portainer.Endpoint{}, dockerCli)
+	require.Error(t, err)
+
+	// Valid endpoint
+
+	endpoint := portainer.Endpoint{ID: 1, PostInitMigrations: portainer.EndpointPostInitMigrations{MigrateGPUs: true}}
+
+	err = store.Endpoint().Create(&endpoint)
+	require.NoError(t, err)
+
+	err = migrator.MigrateGPUs(endpoint, dockerCli)
+	require.NoError(t, err)
+
+	migratedEndpoint, err := store.Endpoint().Endpoint(endpoint.ID)
+	require.NoError(t, err)
+
+	require.Equal(t, endpoint.ID, migratedEndpoint.ID)
+	require.False(t, migratedEndpoint.PostInitMigrations.MigrateGPUs)
+	require.True(t, migratedEndpoint.EnableGPUManagement)
+}
+
+func TestPostInitMigrate_PendingActionsCreated(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name                   string
+		existingPendingActions []*portainer.PendingAction
+		expectedPendingActions int
+		expectedAction         string
+	}{
+		{
+			name: "when existing non-matching action exists, should add migration action",
+			existingPendingActions: []*portainer.PendingAction{
+				{
+					EndpointID: 7,
+					Action:     "some-other-action",
+				},
+			},
+			expectedPendingActions: 2,
+			expectedAction:         actions.PostInitMigrateEnvironment,
+		},
+		{
+			name: "when matching action exists, should not add duplicate",
+			existingPendingActions: []*portainer.PendingAction{
+				{
+					EndpointID: 7,
+					Action:     actions.PostInitMigrateEnvironment,
+				},
+			},
+			expectedPendingActions: 1,
+			expectedAction:         actions.PostInitMigrateEnvironment,
+		},
+		{
+			name:                   "when no actions exist, should add migration action",
+			existingPendingActions: []*portainer.PendingAction{},
+			expectedPendingActions: 1,
+			expectedAction:         actions.PostInitMigrateEnvironment,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			is := assert.New(t)
+			_, store := datastore.MustNewTestStore(t, true, true)
+
+			// Create test endpoint
+			endpoint := &portainer.Endpoint{
+				ID:          7,
+				UserTrusted: true,
+				Type:        portainer.EdgeAgentOnDockerEnvironment,
+				Edge: portainer.EnvironmentEdgeSettings{
+					AsyncMode: false,
+				},
+				EdgeID: "edgeID",
+			}
+			err := store.Endpoint().Create(endpoint)
+			require.NoError(t, err, "error creating endpoint")
+
+			// Create any existing pending actions
+			for _, action := range tt.existingPendingActions {
+				err = store.PendingActions().Create(action)
+				require.NoError(t, err, "error creating pending action")
+			}
+
+			migrator := NewPostInitMigrator(
+				nil, // kubeFactory not needed for this test
+				nil, // dockerFactory not needed for this test
+				store,
+				"",  // assetsPath not needed for this test
+				nil, // kubernetesDeployer not needed for this test
+			)
+
+			err = migrator.PostInitMigrate()
+			require.NoError(t, err, "PostInitMigrate should not return error")
+
+			// Verify the results
+			pendingActions, err := store.PendingActions().ReadAll()
+			require.NoError(t, err, "error reading pending actions")
+			is.Len(pendingActions, tt.expectedPendingActions, "unexpected number of pending actions")
+
+			// If we expect any actions, verify at least one has the expected action type
+			if tt.expectedPendingActions > 0 {
+				hasExpectedAction := false
+				for _, action := range pendingActions {
+					if action.Action == tt.expectedAction {
+						hasExpectedAction = true
+						is.Equal(endpoint.ID, action.EndpointID, "action should reference correct endpoint")
+
+						break
+					}
+				}
+
+				is.True(hasExpectedAction, "should have found action of expected type")
+			}
+		})
+	}
+}