From f6d09da0a3602b903beee9f437b3f99be1966f14 Mon Sep 17 00:00:00 2001
From: Nic Klaassen <nic@goteleport.com>
Date: Wed, 26 Jun 2024 08:40:11 -0700
Subject: [PATCH] use MarshalText and test YAML decoding

---
 api/types/signaturealgorithmsuite.go |  8 ++--
 lib/services/authentication_test.go  | 57 ++++++++++++++++++----------
 2 files changed, 40 insertions(+), 25 deletions(-)

diff --git a/api/types/signaturealgorithmsuite.go b/api/types/signaturealgorithmsuite.go
index 270e5f199fd7a..29b186923cd6c 100644
--- a/api/types/signaturealgorithmsuite.go
+++ b/api/types/signaturealgorithmsuite.go
@@ -57,10 +57,10 @@ func (s *SignatureAlgorithmSuite) fromString(str string) error {
 	return nil
 }
 
-// UnmarshalJSON marshals a SignatureAlgorithmSuite value to JSON.
-func (s *SignatureAlgorithmSuite) MarshalJSON() ([]byte, error) {
-	out, err := json.Marshal(s.toString())
-	return out, trace.Wrap(err)
+// MarshalText marshals a SignatureAlgorithmSuite value to text. This gets used
+// by json.Marshal.
+func (s *SignatureAlgorithmSuite) MarshalText() ([]byte, error) {
+	return []byte(s.toString()), nil
 }
 
 // UnmarshalJSON unmarshals a SignatureAlgorithmSuite and supports the custom
diff --git a/lib/services/authentication_test.go b/lib/services/authentication_test.go
index 442be44916469..2bcf36ad391d1 100644
--- a/lib/services/authentication_test.go
+++ b/lib/services/authentication_test.go
@@ -20,14 +20,17 @@ package services_test
 
 import (
 	"fmt"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/gravitational/trace"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	kyaml "k8s.io/apimachinery/pkg/util/yaml"
 
 	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/services"
 )
 
@@ -135,75 +138,87 @@ func TestSignatureAlgorithmSuiteRoundtrip(t *testing.T) {
 
 func TestParseSignatureAlgorithmSuite(t *testing.T) {
 	for _, tc := range []struct {
-		json     string
+		rawValue string
 		expected types.SignatureAlgorithmSuite
 	}{
 		{
-			json:     `""`,
+			rawValue: `""`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_UNSPECIFIED,
 		},
 		{
-			json:     `"SIGNATURE_ALGORITHM_SUITE_UNSPECIFIED"`,
+			rawValue: `"SIGNATURE_ALGORITHM_SUITE_UNSPECIFIED"`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_UNSPECIFIED,
 		},
 		{
-			json:     `"legacy"`,
+			rawValue: `"legacy"`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_LEGACY,
 		},
 		{
-			json:     `"SIGNATURE_ALGORITHM_SUITE_LEGACY"`,
+			rawValue: `"SIGNATURE_ALGORITHM_SUITE_LEGACY"`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_LEGACY,
 		},
 		{
-			json:     `1`,
+			rawValue: `1`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_LEGACY,
 		},
 		{
-			json:     `1.0`,
+			rawValue: `1.0`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_LEGACY,
 		},
 		{
-			json:     `"balanced-v1"`,
+			rawValue: `"balanced-v1"`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_BALANCED_V1,
 		},
 		{
-			json:     `"SIGNATURE_ALGORITHM_SUITE_BALANCED_V1"`,
+			rawValue: `"SIGNATURE_ALGORITHM_SUITE_BALANCED_V1"`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_BALANCED_V1,
 		},
 		{
-			json:     `2`,
+			rawValue: `2`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_BALANCED_V1,
 		},
 		{
-			json:     `"fips-v1"`,
+			rawValue: `"fips-v1"`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_FIPS_V1,
 		},
 		{
-			json:     `"SIGNATURE_ALGORITHM_SUITE_FIPS_V1"`,
+			rawValue: `"SIGNATURE_ALGORITHM_SUITE_FIPS_V1"`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_FIPS_V1,
 		},
 		{
-			json:     `3`,
+			rawValue: `3`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_FIPS_V1,
 		},
 		{
-			json:     `"hsm-v1"`,
+			rawValue: `"hsm-v1"`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_HSM_V1,
 		},
 		{
-			json:     `"SIGNATURE_ALGORITHM_SUITE_HSM_V1"`,
+			rawValue: `"SIGNATURE_ALGORITHM_SUITE_HSM_V1"`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_HSM_V1,
 		},
 		{
-			json:     `4`,
+			rawValue: `4`,
 			expected: types.SignatureAlgorithmSuite_SIGNATURE_ALGORITHM_SUITE_HSM_V1,
 		},
 	} {
-		t.Run(tc.json, func(t *testing.T) {
-			raw := fmt.Sprintf(`{"spec":{"signature_algorithm_suite":%s}}`, tc.json)
-			unmarshaled, err := services.UnmarshalAuthPreference([]byte(raw))
-			require.NoError(t, err)
-			require.Equal(t, tc.expected, unmarshaled.GetSignatureAlgorithmSuite())
+		t.Run(tc.rawValue, func(t *testing.T) {
+			t.Run("json", func(t *testing.T) {
+				rawJSON := fmt.Sprintf(`{"spec":{"signature_algorithm_suite":%s}}`, tc.rawValue)
+				unmarshaled, err := services.UnmarshalAuthPreference([]byte(rawJSON))
+				require.NoError(t, err)
+				require.Equal(t, tc.expected, unmarshaled.GetSignatureAlgorithmSuite())
+			})
+
+			t.Run("yaml", func(t *testing.T) {
+				rawYAML := fmt.Sprintf(`spec: { signature_algorithm_suite: %s }`, tc.rawValue)
+
+				decoder := kyaml.NewYAMLOrJSONDecoder(strings.NewReader(rawYAML), defaults.LookaheadBufSize)
+				var authPref types.AuthPreferenceV2
+				err := decoder.Decode(&authPref)
+				require.NoError(t, err)
+				require.Equal(t, tc.expected, authPref.GetSignatureAlgorithmSuite())
+			})
 		})
 	}
 }