fix: patch correctly config in talosctl upgrade-k8s

The current code was stipping non-`v1alpha1.Config` documents. Provide a
proper method in the config provider, and update places using it.

Signed-off-by: Andrey Smirnov <andrey.smirnov@siderolabs.com>

internal/integration/api/apply-config.go

@@ -25,7 +25,6 @@ import (
 	machineapi "github.com/siderolabs/talos/pkg/machinery/api/machine"
 	"github.com/siderolabs/talos/pkg/machinery/client"
 	"github.com/siderolabs/talos/pkg/machinery/config"
-	"github.com/siderolabs/talos/pkg/machinery/config/container"
 	"github.com/siderolabs/talos/pkg/machinery/config/machine"
 	"github.com/siderolabs/talos/pkg/machinery/config/types/v1alpha1"
 	"github.com/siderolabs/talos/pkg/machinery/constants"
@@ -285,10 +284,9 @@ func (suite *ApplyConfigSuite) TestApplyConfigRotateEncryptionSecrets() {
 	}
 
 	for _, keys := range keySets {
-		cfg.EncryptionKeys = keys
-
-		data, err := container.NewV1Alpha1(machineConfig).Bytes()
-		suite.Require().NoError(err)
+		data := suite.PatchV1Alpha1Config(provider, func(cfg *v1alpha1.Config) {
+			cfg.MachineConfig.MachineSystemDiskEncryption.EphemeralPartition.EncryptionKeys = keys
+		})
 
 		suite.AssertRebooted(
 			suite.ctx, node, func(nodeCtx context.Context) error {

internal/integration/base/api.go

@@ -21,7 +21,6 @@ import (
 
 	"github.com/cosi-project/runtime/pkg/safe"
 	"github.com/cosi-project/runtime/pkg/state"
-	"github.com/siderolabs/gen/xslices"
 	"github.com/siderolabs/go-retry/retry"
 	"github.com/stretchr/testify/suite"
 	"google.golang.org/grpc/backoff"
@@ -35,9 +34,7 @@ import (
 	"github.com/siderolabs/talos/pkg/machinery/client"
 	clientconfig "github.com/siderolabs/talos/pkg/machinery/client/config"
 	"github.com/siderolabs/talos/pkg/machinery/config"
-	configtypes "github.com/siderolabs/talos/pkg/machinery/config/config"
 	"github.com/siderolabs/talos/pkg/machinery/config/configloader"
-	"github.com/siderolabs/talos/pkg/machinery/config/container"
 	"github.com/siderolabs/talos/pkg/machinery/config/machine"
 	"github.com/siderolabs/talos/pkg/machinery/config/types/v1alpha1"
 	"github.com/siderolabs/talos/pkg/machinery/constants"
@@ -584,18 +581,11 @@ func (apiSuite *APISuite) AssertExpectedModules(ctx context.Context, node string
 
 // PatchV1Alpha1Config patches v1alpha1 config in the config provider.
 func (apiSuite *APISuite) PatchV1Alpha1Config(provider config.Provider, patch func(*v1alpha1.Config)) []byte {
-	cfg := provider.RawV1Alpha1()
-	apiSuite.Require().NotNil(cfg)
+	ctr, err := provider.PatchV1Alpha1(func(c *v1alpha1.Config) error {
+		patch(c)
 
-	patch(cfg)
-
-	otherDocs := xslices.Filter(provider.Documents(), func(doc configtypes.Document) bool {
-		_, ok := doc.(*v1alpha1.Config)
-
-		return !ok
+		return nil
 	})
-
-	ctr, err := container.New(append([]configtypes.Document{cfg}, otherDocs...)...)
 	apiSuite.Require().NoError(err)
 
 	bytes, err := ctr.Bytes()

pkg/cluster/kubernetes/patch.go

@@ -6,14 +6,12 @@ package kubernetes
 
 import (
 	"context"
-	"errors"
 	"fmt"
 
 	"github.com/cosi-project/runtime/pkg/safe"
 
 	"github.com/siderolabs/talos/pkg/machinery/api/machine"
 	"github.com/siderolabs/talos/pkg/machinery/client"
-	"github.com/siderolabs/talos/pkg/machinery/config/container"
 	"github.com/siderolabs/talos/pkg/machinery/config/encoder"
 	v1alpha1config "github.com/siderolabs/talos/pkg/machinery/config/types/v1alpha1"
 	"github.com/siderolabs/talos/pkg/machinery/resources/config"
@@ -33,16 +31,14 @@ func patchNodeConfig(ctx context.Context, cluster UpgradeProvider, node string,
 		return fmt.Errorf("error fetching config resource: %w", err)
 	}
 
-	cfg := mc.Container().RawV1Alpha1()
-	if cfg == nil {
-		return errors.New("config is not v1alpha1 config")
-	}
+	provider := mc.Provider()
 
-	if err = patchFunc(cfg); err != nil {
+	newProvider, err := provider.PatchV1Alpha1(patchFunc)
+	if err != nil {
 		return fmt.Errorf("error patching config: %w", err)
 	}
 
-	cfgBytes, err := container.NewV1Alpha1(cfg).EncodeBytes(encoderOpt)
+	cfgBytes, err := newProvider.EncodeBytes(encoderOpt)
 	if err != nil {
 		return fmt.Errorf("error serializing config: %w", err)
 	}

pkg/machinery/config/container/container.go

@@ -97,6 +97,28 @@ func (container *Container) Clone() coreconfig.Provider {
 	}
 }
 
+// PatchV1Alpha1 patches the container's v1alpha1.Config while preserving other config documents.
+func (container *Container) PatchV1Alpha1(patcher func(*v1alpha1.Config) error) (coreconfig.Provider, error) {
+	cfg := container.RawV1Alpha1()
+	if cfg == nil {
+		return nil, fmt.Errorf("v1alpha1.Config is not present in the container")
+	}
+
+	cfg = cfg.DeepCopy()
+
+	if err := patcher(cfg); err != nil {
+		return nil, err
+	}
+
+	otherDocs := xslices.Filter(container.Documents(), func(doc config.Document) bool {
+		_, ok := doc.(*v1alpha1.Config)
+
+		return !ok
+	})
+
+	return New(append([]config.Document{cfg}, otherDocs...)...)
+}
+
 // Readonly implements config.Container interface.
 func (container *Container) Readonly() bool {
 	return container.readonly

pkg/machinery/config/container/container_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/siderolabs/talos/pkg/machinery/config/config"
 	"github.com/siderolabs/talos/pkg/machinery/config/configloader"
 	"github.com/siderolabs/talos/pkg/machinery/config/container"
+	"github.com/siderolabs/talos/pkg/machinery/config/machine"
 	"github.com/siderolabs/talos/pkg/machinery/config/types/runtime/extensions"
 	"github.com/siderolabs/talos/pkg/machinery/config/types/siderolink"
 	"github.com/siderolabs/talos/pkg/machinery/config/types/v1alpha1"
@@ -89,6 +90,35 @@ func TestNewDuplicate(t *testing.T) {
 	assert.EqualError(t, err, "duplicate document: SideroLinkConfig/")
 }
 
+func TestPatchV1Alpha1(t *testing.T) {
+	t.Parallel()
+
+	v1alpha1Cfg := &v1alpha1.Config{
+		MachineConfig: &v1alpha1.MachineConfig{
+			MachineType: "worker",
+		},
+	}
+
+	sideroLinkCfg := siderolink.NewConfigV1Alpha1()
+	sideroLinkCfg.APIUrlConfig.URL = must.Value(url.Parse("https://siderolink.api/?jointoken=secret&user=alice"))(t)
+
+	cfg, err := container.New(v1alpha1Cfg, sideroLinkCfg)
+	require.NoError(t, err)
+
+	patchedCfg, err := cfg.PatchV1Alpha1(func(cfg *v1alpha1.Config) error {
+		cfg.MachineConfig.MachineType = "controlplane"
+
+		return nil
+	})
+	require.NoError(t, err)
+
+	assert.Equal(t, machine.TypeWorker, cfg.Machine().Type())
+	assert.Equal(t, machine.TypeControlPlane, patchedCfg.Machine().Type())
+
+	assert.Equal(t, "https://siderolink.api/?jointoken=secret&user=alice", cfg.SideroLink().APIUrl().String())
+	assert.Equal(t, "https://siderolink.api/?jointoken=secret&user=alice", patchedCfg.SideroLink().APIUrl().String())
+}
+
 func TestValidate(t *testing.T) {
 	t.Parallel()
 

pkg/machinery/config/provider.go

@@ -42,6 +42,9 @@ type Provider interface {
 	// Clone returns a copy of the Provider.
 	Clone() Provider
 
+	// PatchV1Alpha1 patches the container's v1alpha1.Config while preserving other config documents.
+	PatchV1Alpha1(patcher func(*v1alpha1.Config) error) (Provider, error)
+
 	// RedactSecrets returns a copy of the Provider with all secrets replaced with the given string.
 	RedactSecrets(string) Provider