From 9dc860d5fea70535323b2029822a06a34c63a2e9 Mon Sep 17 00:00:00 2001 From: Hugo Shaka Date: Thu, 9 Jan 2025 09:18:57 -0500 Subject: [PATCH] Fix proto resource 153 marshalling for autoupdate_* resources (#50688) * Fix proto resource 153 marshalling * Update tool/tctl/common/collection_test.go Co-authored-by: Alan Parra * Update tool/tctl/common/collection_test.go Co-authored-by: Alan Parra * Address feedback - Change from Resource153AdapterV2 to ProtoResource153Adapter - fix test failures and unmarshal proto resources properly - add a failing round-trip proto 153 test case - bonus: fix the table tesst reosurce create that did not support running a single row * Apply suggestions from code review Co-authored-by: Alan Parra * lint --------- Co-authored-by: Alan Parra --- api/types/resource_153.go | 39 +++++++++++ tool/tctl/common/collection.go | 6 +- tool/tctl/common/collection_test.go | 69 +++++++++++++++++++ tool/tctl/common/helpers_test.go | 9 +++ tool/tctl/common/resource_command_test.go | 82 ++++++++++++++++------- 5 files changed, 177 insertions(+), 28 deletions(-) diff --git a/api/types/resource_153.go b/api/types/resource_153.go index a09c39451cd3d..dbe69a1108466 100644 --- a/api/types/resource_153.go +++ b/api/types/resource_153.go @@ -18,6 +18,8 @@ import ( "encoding/json" "time" + "google.golang.org/protobuf/encoding/protojson" + "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/timestamppb" headerv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/header/v1" @@ -124,6 +126,10 @@ func (r *legacyToResource153Adapter) GetVersion() string { // [Resource] type. Implements [ResourceWithLabels] and CloneResource (where the) // wrapped resource supports cloning). // +// Resources153 implemented by proto-generated structs should use ProtoResource153ToLegacy +// instead as it will ensure the protobuf message is properly marshaled to JSON +// with protojson. +// // Note that CheckAndSetDefaults is a noop for the returned resource and // SetSubKind is not implemented and panics on use. func Resource153ToLegacy(r Resource153) Resource { @@ -348,3 +354,36 @@ func (r *resource153ToUnifiedResourceAdapter) CloneResource() ResourceWithLabels clone := r.inner.(ClonableResource153).CloneResource() return Resource153ToUnifiedResource(clone) } + +// ProtoResource153 is a Resource153 implemented by a protobuf-generated struct. +type ProtoResource153 interface { + Resource153 + proto.Message +} + +type protoResource153ToLegacyAdapter struct { + inner ProtoResource153 + resource153ToLegacyAdapter +} + +// MarshalJSON adds support for marshaling the wrapped resource (instead of +// marshaling the adapter itself). +func (r *protoResource153ToLegacyAdapter) MarshalJSON() ([]byte, error) { + return protojson.MarshalOptions{ + UseProtoNames: true, + }.Marshal(r.inner) +} + +// ProtoResource153ToLegacy transforms an RFD 153 style resource implemented by +// a proto-generated struct into a legacy [Resource] type. Implements +// [ResourceWithLabels] and CloneResource (where the wrapped resource supports +// cloning). +// +// Note that CheckAndSetDefaults is a noop for the returned resource and +// SetSubKind is not implemented and panics on use. +func ProtoResource153ToLegacy(r ProtoResource153) Resource { + return &protoResource153ToLegacyAdapter{ + r, + resource153ToLegacyAdapter{r}, + } +} diff --git a/tool/tctl/common/collection.go b/tool/tctl/common/collection.go index c31d2a25ed0bf..c1ea21addc2b6 100644 --- a/tool/tctl/common/collection.go +++ b/tool/tctl/common/collection.go @@ -1908,7 +1908,7 @@ type autoUpdateConfigCollection struct { } func (c *autoUpdateConfigCollection) resources() []types.Resource { - return []types.Resource{types.Resource153ToLegacy(c.config)} + return []types.Resource{types.ProtoResource153ToLegacy(c.config)} } func (c *autoUpdateConfigCollection) writeText(w io.Writer, verbose bool) error { @@ -1926,7 +1926,7 @@ type autoUpdateVersionCollection struct { } func (c *autoUpdateVersionCollection) resources() []types.Resource { - return []types.Resource{types.Resource153ToLegacy(c.version)} + return []types.Resource{types.ProtoResource153ToLegacy(c.version)} } func (c *autoUpdateVersionCollection) writeText(w io.Writer, verbose bool) error { @@ -1944,7 +1944,7 @@ type autoUpdateAgentRolloutCollection struct { } func (c *autoUpdateAgentRolloutCollection) resources() []types.Resource { - return []types.Resource{types.Resource153ToLegacy(c.rollout)} + return []types.Resource{types.ProtoResource153ToLegacy(c.rollout)} } func (c *autoUpdateAgentRolloutCollection) writeText(w io.Writer, verbose bool) error { diff --git a/tool/tctl/common/collection_test.go b/tool/tctl/common/collection_test.go index 166c5f6901599..f0679b9a65581 100644 --- a/tool/tctl/common/collection_test.go +++ b/tool/tctl/common/collection_test.go @@ -27,13 +27,19 @@ import ( "github.com/google/go-cmp/cmp" "github.com/google/uuid" "github.com/stretchr/testify/require" + "google.golang.org/protobuf/types/known/durationpb" + kyaml "k8s.io/apimachinery/pkg/util/yaml" "github.com/gravitational/teleport/api" + autoupdatev1pb "github.com/gravitational/teleport/api/gen/proto/go/teleport/autoupdate/v1" dbobjectv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/dbobject/v1" dbobjectimportrulev1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/dbobjectimportrule/v1" "github.com/gravitational/teleport/api/types" + "github.com/gravitational/teleport/api/types/autoupdate" "github.com/gravitational/teleport/api/types/label" "github.com/gravitational/teleport/lib/asciitable" + "github.com/gravitational/teleport/lib/defaults" + "github.com/gravitational/teleport/lib/services" "github.com/gravitational/teleport/lib/srv/db/common/databaseobject" "github.com/gravitational/teleport/lib/srv/db/common/databaseobjectimportrule" "github.com/gravitational/teleport/tool/common" @@ -431,3 +437,66 @@ func makeTestLabels(extraStaticLabels map[string]string) map[string]string { maps.Copy(labels, extraStaticLabels) return labels } + +// autoUpdateConfigBrokenCollection is an intentionally broken version of the +// autoUpdateConfigCollection that is not marshaling resources properly because +// it's doing json marshaling instead of protojson marshaling. +type autoUpdateConfigBrokenCollection struct { + autoUpdateConfigCollection +} + +func (c *autoUpdateConfigBrokenCollection) resources() []types.Resource { + // We use Resource153ToLegacy instead of ProtoResource153ToLegacy. + return []types.Resource{types.Resource153ToLegacy(c.config)} +} + +// This test makes sure we marshal and unmarshal proto-based Resource153 properly. +// We had a bug where types.Resource153 implemented by protobuf structs were not +// marshaled properly (they should be marshaled using protojson). This test +// checks we can do a round-trip with one of those proto-struct resource. +func TestRoundTripProtoResource153(t *testing.T) { + // Test setup: generate fixture. + initial, err := autoupdate.NewAutoUpdateConfig(&autoupdatev1pb.AutoUpdateConfigSpec{ + Agents: &autoupdatev1pb.AutoUpdateConfigSpecAgents{ + Mode: autoupdate.AgentsUpdateModeEnabled, + Strategy: autoupdate.AgentsStrategyTimeBased, + MaintenanceWindowDuration: durationpb.New(1 * time.Hour), + Schedules: &autoupdatev1pb.AgentAutoUpdateSchedules{ + Regular: []*autoupdatev1pb.AgentAutoUpdateGroup{ + { + Name: "group1", + Days: []string{types.Wildcard}, + }, + }, + }, + }, + }) + require.NoError(t, err) + + // Test execution: dump the resource into a YAML manifest. + collection := &autoUpdateConfigCollection{config: initial} + buf := &bytes.Buffer{} + require.NoError(t, writeYAML(collection, buf)) + + // Test execution: load the YAML manifest back. + decoder := kyaml.NewYAMLOrJSONDecoder(buf, defaults.LookaheadBufSize) + var raw services.UnknownResource + require.NoError(t, decoder.Decode(&raw)) + result, err := services.UnmarshalProtoResource[*autoupdatev1pb.AutoUpdateConfig](raw.Raw) + require.NoError(t, err) + + // Test validation: check that the loaded content matches what we had before. + require.Equal(t, result, initial) + + // Test execution: now dump the resource into a YAML manifest with a + // collection using types.Resource153ToLegacy instead of types.ProtoResource153ToLegacy + brokenCollection := &autoUpdateConfigBrokenCollection{autoUpdateConfigCollection{initial}} + buf = &bytes.Buffer{} + require.NoError(t, writeYAML(brokenCollection, buf)) + + // Test execution: load the YAML manifest back and see that we can't unmarshal it. + decoder = kyaml.NewYAMLOrJSONDecoder(buf, defaults.LookaheadBufSize) + require.NoError(t, decoder.Decode(&raw)) + _, err = services.UnmarshalProtoResource[*autoupdatev1pb.AutoUpdateConfig](raw.Raw) + require.Error(t, err) +} diff --git a/tool/tctl/common/helpers_test.go b/tool/tctl/common/helpers_test.go index 0cf773852c96f..b235a40e8b5e2 100644 --- a/tool/tctl/common/helpers_test.go +++ b/tool/tctl/common/helpers_test.go @@ -35,6 +35,7 @@ import ( "github.com/jonboulle/clockwork" "github.com/stretchr/testify/require" "gopkg.in/yaml.v2" + kyaml "k8s.io/apimachinery/pkg/util/yaml" "github.com/gravitational/teleport/api/breaker" apidefaults "github.com/gravitational/teleport/api/defaults" @@ -43,6 +44,7 @@ import ( "github.com/gravitational/teleport/lib/config" "github.com/gravitational/teleport/lib/service" "github.com/gravitational/teleport/lib/service/servicecfg" + "github.com/gravitational/teleport/lib/services" "github.com/gravitational/teleport/lib/utils" commonclient "github.com/gravitational/teleport/tool/tctl/common/client" tctlcfg "github.com/gravitational/teleport/tool/tctl/common/config" @@ -153,6 +155,13 @@ func mustDecodeJSON[T any](t *testing.T, r io.Reader) T { return out } +func mustTranscodeYAMLToJSON(t *testing.T, r io.Reader) []byte { + decoder := kyaml.NewYAMLToJSONDecoder(r) + var resource services.UnknownResource + require.NoError(t, decoder.Decode(&resource)) + return resource.Raw +} + func mustDecodeYAMLDocuments[T any](t *testing.T, r io.Reader, out *[]T) { t.Helper() decoder := yaml.NewDecoder(r) diff --git a/tool/tctl/common/resource_command_test.go b/tool/tctl/common/resource_command_test.go index 61b2c2650f53a..ff280c07a6f08 100644 --- a/tool/tctl/common/resource_command_test.go +++ b/tool/tctl/common/resource_command_test.go @@ -36,6 +36,7 @@ import ( "github.com/jonboulle/clockwork" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "google.golang.org/protobuf/encoding/protojson" "google.golang.org/protobuf/testing/protocmp" "k8s.io/apimachinery/pkg/util/yaml" @@ -1371,17 +1372,29 @@ func TestCreateResources(t *testing.T) { process := testenv.MakeTestServer(t, testenv.WithLogger(utils.NewSlogLoggerForTests())) rootClient := testenv.MakeDefaultAuthClient(t, process) + // tctlGetAllValidations allows tests to register post-test validations to validate + // that their resource is present in "tctl get all" output. + // This allows running test rows instead of the whole test table. + var tctlGetAllValidations []func(t *testing.T, out string) + tests := []struct { - kind string - create func(t *testing.T, clt *authclient.Client) + kind string + create func(t *testing.T, clt *authclient.Client) + getAllCheck func(t *testing.T, out string) }{ { kind: types.KindGithubConnector, create: testCreateGithubConnector, + getAllCheck: func(t *testing.T, s string) { + assert.Contains(t, s, "kind: github") + }, }, { kind: types.KindRole, create: testCreateRole, + getAllCheck: func(t *testing.T, s string) { + assert.Contains(t, s, "kind: role") + }, }, { kind: types.KindServerInfo, @@ -1390,6 +1403,9 @@ func TestCreateResources(t *testing.T) { { kind: types.KindUser, create: testCreateUser, + getAllCheck: func(t *testing.T, s string) { + assert.Contains(t, s, "kind: user") + }, }, { kind: types.KindDatabaseObjectImportRule, @@ -1402,10 +1418,16 @@ func TestCreateResources(t *testing.T) { { kind: types.KindClusterNetworkingConfig, create: testCreateClusterNetworkingConfig, + getAllCheck: func(t *testing.T, s string) { + assert.Contains(t, s, "kind: cluster_networking_config") + }, }, { kind: types.KindClusterAuthPreference, create: testCreateAuthPreference, + getAllCheck: func(t *testing.T, s string) { + assert.Contains(t, s, "kind: cluster_auth_preference") + }, }, { kind: types.KindSessionRecordingConfig, @@ -1440,6 +1462,9 @@ func TestCreateResources(t *testing.T) { for _, test := range tests { t.Run(test.kind, func(t *testing.T) { test.create(t, rootClient) + if test.getAllCheck != nil { + tctlGetAllValidations = append(tctlGetAllValidations, test.getAllCheck) + } }) } @@ -1447,12 +1472,9 @@ func TestCreateResources(t *testing.T) { out, err := runResourceCommand(t, rootClient, []string{"get", "all"}) require.NoError(t, err) s := out.String() - require.NotEmpty(t, s) - assert.Contains(t, s, "kind: github") - assert.Contains(t, s, "kind: cluster_auth_preference") - assert.Contains(t, s, "kind: cluster_networking_config") - assert.Contains(t, s, "kind: user") - assert.Contains(t, s, "kind: role") + for _, validateGetAll := range tctlGetAllValidations { + validateGetAll(t, s) + } } func testCreateGithubConnector(t *testing.T, clt *authclient.Client) { @@ -2326,18 +2348,21 @@ version: v1 _, err = runResourceCommand(t, clt, []string{"create", resourceYAMLPath}) require.NoError(t, err) - // Get the resource buf, err := runResourceCommand(t, clt, []string{"get", types.KindAutoUpdateConfig, "--format=json"}) require.NoError(t, err) - resources := mustDecodeJSON[[]*autoupdate.AutoUpdateConfig](t, buf) - require.Len(t, resources, 1) + + rawResources := mustDecodeJSON[[]services.UnknownResource](t, buf) + require.Len(t, rawResources, 1) + var resource autoupdate.AutoUpdateConfig + require.NoError(t, protojson.Unmarshal(rawResources[0].Raw, &resource)) var expected autoupdate.AutoUpdateConfig - require.NoError(t, yaml.Unmarshal([]byte(resourceYAML), &expected)) + expectedJSON := mustTranscodeYAMLToJSON(t, bytes.NewReader([]byte(resourceYAML))) + require.NoError(t, protojson.Unmarshal(expectedJSON, &expected)) require.Empty(t, cmp.Diff( - []*autoupdate.AutoUpdateConfig{&expected}, - resources, + &expected, + &resource, protocmp.IgnoreFields(&headerv1.Metadata{}, "revision"), protocmp.Transform(), )) @@ -2368,18 +2393,21 @@ version: v1 _, err = runResourceCommand(t, clt, []string{"create", resourceYAMLPath}) require.NoError(t, err) - // Get the resource buf, err := runResourceCommand(t, clt, []string{"get", types.KindAutoUpdateVersion, "--format=json"}) require.NoError(t, err) - resources := mustDecodeJSON[[]*autoupdate.AutoUpdateVersion](t, buf) - require.Len(t, resources, 1) + + rawResources := mustDecodeJSON[[]services.UnknownResource](t, buf) + require.Len(t, rawResources, 1) + var resource autoupdate.AutoUpdateVersion + require.NoError(t, protojson.Unmarshal(rawResources[0].Raw, &resource)) var expected autoupdate.AutoUpdateVersion - require.NoError(t, yaml.Unmarshal([]byte(resourceYAML), &expected)) + expectedJSON := mustTranscodeYAMLToJSON(t, bytes.NewReader([]byte(resourceYAML))) + require.NoError(t, protojson.Unmarshal(expectedJSON, &expected)) require.Empty(t, cmp.Diff( - []*autoupdate.AutoUpdateVersion{&expected}, - resources, + &expected, + &resource, protocmp.IgnoreFields(&headerv1.Metadata{}, "revision"), protocmp.Transform(), )) @@ -2423,15 +2451,19 @@ version: v1 // Get the resource buf, err := runResourceCommand(t, clt, []string{"get", types.KindAutoUpdateAgentRollout, "--format=json"}) require.NoError(t, err) - resources := mustDecodeJSON[[]*autoupdate.AutoUpdateAgentRollout](t, buf) - require.Len(t, resources, 1) + + rawResources := mustDecodeJSON[[]services.UnknownResource](t, buf) + require.Len(t, rawResources, 1) + var resource autoupdate.AutoUpdateAgentRollout + require.NoError(t, protojson.Unmarshal(rawResources[0].Raw, &resource)) var expected autoupdate.AutoUpdateAgentRollout - require.NoError(t, yaml.Unmarshal([]byte(resourceYAML), &expected)) + expectedJSON := mustTranscodeYAMLToJSON(t, bytes.NewReader([]byte(resourceYAML))) + require.NoError(t, protojson.Unmarshal(expectedJSON, &expected)) require.Empty(t, cmp.Diff( - []*autoupdate.AutoUpdateAgentRollout{&expected}, - resources, + &expected, + &resource, protocmp.IgnoreFields(&headerv1.Metadata{}, "revision"), protocmp.Transform(), ))