Modify all custom types to enable custom YAML marshalling and unmarshalling

authorized_keys.go

@@ -5,14 +5,102 @@ import (
 	"io/ioutil"
 
 	"golang.org/x/crypto/ssh"
+	"gopkg.in/yaml.v2"
 )
 
+type AuthorizedKeysMap map[string]AuthorizedKeys
+
+// Can be removed once flags are deprecated
+func (a AuthorizedKeysMap) Set(value string) error {
+	authorizedKeyPaths := make(map[string]string)
+	err := yaml.Unmarshal([]byte(value), &authorizedKeyPaths)
+	if err != nil {
+		return err
+	}
+
+	if authorizedKeyPaths != nil {
+		return a.set(authorizedKeyPaths)
+	}
+
+	return nil
+}
+
+// Can be removed once flags are deprecated
+func (a AuthorizedKeysMap) String() string {
+	if a == nil {
+		return ""
+	}
+
+	authorizedKeysString, _ := a.convertToString()
+	return authorizedKeysString
+}
+
+// Can be removed once flags are deprecated
+func (a AuthorizedKeysMap) Type() string {
+	return "AuthorizedKeysMap"
+}
+
+func (a AuthorizedKeysMap) convertToString() (string, error) {
+	authorizedKeysPaths := make(map[string]string)
+	for key, authorizedKey := range a {
+		authorizedKeysPaths[key] = authorizedKey.File
+	}
+
+	authorizedKeysString, err := yaml.Marshal(authorizedKeysPaths)
+	if err != nil {
+		return "", err
+	}
+
+	return string(authorizedKeysString), nil
+}
+
+func (a *AuthorizedKeysMap) set(value map[string]string) error {
+	authorizedKeysMap := make(AuthorizedKeysMap)
+	for key, authorizedKeyPath := range value {
+		var authorizedKeys AuthorizedKeys
+		err := authorizedKeys.Set(authorizedKeyPath)
+		if err != nil {
+			return err
+		}
+
+		authorizedKeysMap[key] = authorizedKeys
+	}
+
+	a = &authorizedKeysMap
+
+	return nil
+}
+
 type AuthorizedKeys struct {
 	File string
 	Keys []ssh.PublicKey
 }
 
-func (f *AuthorizedKeys) UnmarshalFlag(value string) error {
+// Reload reloads the value of the Keys
+func (a *AuthorizedKeys) Reload() error {
+	return a.Set(a.File)
+}
+
+func (a AuthorizedKeys) MarshalYAML() (interface{}, error) {
+	return a.File, nil
+}
+
+func (a *AuthorizedKeys) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var path string
+	err := unmarshal(&path)
+	if err != nil {
+		return err
+	}
+
+	if path != "" {
+		return a.Set(path)
+	}
+
+	return nil
+}
+
+// Can be removed once flags are deprecated
+func (a *AuthorizedKeys) Set(value string) error {
 	authorizedKeysBytes, err := ioutil.ReadFile(value)
 	if err != nil {
 		return fmt.Errorf("failed to read authorized keys: %s", err)
@@ -32,13 +120,22 @@ func (f *AuthorizedKeys) UnmarshalFlag(value string) error {
 		authorizedKeysBytes = rest
 	}
 
-	f.File = value
-	f.Keys = authorizedKeys
+	a.File = value
+	a.Keys = authorizedKeys
 
 	return nil
 }
 
-// Reload reloads the value of the Keys
-func (f *AuthorizedKeys) Reload() error {
-	return f.UnmarshalFlag(f.File)
+// Can be removed once flags are deprecated
+func (a *AuthorizedKeys) String() string {
+	if a == nil {
+		return ""
+	}
+
+	return a.File
+}
+
+// Can be removed once flags are deprecated
+func (a *AuthorizedKeys) Type() string {
+	return "AuthorizedKeys"
 }

cipher.go

@@ -8,18 +8,54 @@ import (
 
 type Cipher struct {
 	cipher.AEAD
+	originalCipher string
 }
 
-func (flag *Cipher) UnmarshalFlag(val string) error {
-	block, err := aes.NewCipher([]byte(val))
+func (c Cipher) MarshalYAML() (interface{}, error) {
+	return c.originalCipher, nil
+}
+
+func (c *Cipher) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var value string
+	err := unmarshal(&value)
+	if err != nil {
+		return err
+	}
+
+	if value != "" {
+		return c.Set(value)
+	}
+
+	return nil
+}
+
+// Can be removed once flags are deprecated
+func (c *Cipher) Set(value string) error {
+	block, err := aes.NewCipher([]byte(value))
 	if err != nil {
 		return fmt.Errorf("failed to construct AES cipher: %s", err)
 	}
 
-	flag.AEAD, err = cipher.NewGCM(block)
+	c.AEAD, err = cipher.NewGCM(block)
 	if err != nil {
 		return fmt.Errorf("failed to construct GCM: %s", err)
 	}
 
+	c.originalCipher = value
+
 	return nil
 }
+
+// Can be removed once flags are deprecated
+func (c *Cipher) String() string {
+	if c == nil {
+		return ""
+	}
+
+	return c.originalCipher
+}
+
+// Can be removed once flags are deprecated
+func (c *Cipher) Type() string {
+	return "AEAD"
+}

directory.go

@@ -8,7 +8,26 @@ import (
 
 type Dir string
 
-func (f *Dir) UnmarshalFlag(value string) error {
+func (d Dir) MarshalYAML() (interface{}, error) {
+	return string(d), nil
+}
+
+func (d *Dir) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var value string
+	err := unmarshal(&value)
+	if err != nil {
+		return err
+	}
+
+	if value != "" {
+		return d.Set(value)
+	}
+
+	return nil
+}
+
+// Can be removed once flags are deprecated
+func (d *Dir) Set(value string) error {
 	stat, err := os.Stat(value)
 	if err == nil {
 		if !stat.IsDir() {
@@ -21,11 +40,25 @@ func (f *Dir) UnmarshalFlag(value string) error {
 		return err
 	}
 
-	*f = Dir(abs)
+	*d = Dir(abs)
 
 	return nil
 }
 
-func (f Dir) Path() string {
-	return string(f)
+// Can be removed once flags are deprecated
+func (d *Dir) String() string {
+	if d == nil {
+		return ""
+	}
+
+	return string(*d)
+}
+
+// Can be removed once flags are deprecated
+func (d *Dir) Type() string {
+	return "Dir"
+}
+
+func (d *Dir) Path() string {
+	return string(*d)
 }

file.go

@@ -4,11 +4,65 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 )
 
+type Files []File
+
+// Can be removed once flags are deprecated
+func (f *Files) Set(value string) error {
+	unparsedFiles := strings.Split(value, ",")
+
+	var parsedFiles Files
+	for _, unparsedFile := range unparsedFiles {
+		var file File
+		err := file.Set(strings.TrimSpace(unparsedFile))
+		if err != nil {
+			return err
+		}
+
+		parsedFiles = append(parsedFiles, file)
+	}
+
+	return nil
+}
+
+// Can be removed once flags are deprecated
+func (f *Files) String() string {
+	if f == nil {
+		return ""
+	}
+
+	return fmt.Sprintf("%v", *f)
+}
+
+// Can be removed once flags are deprecated
+func (f *Files) Type() string {
+	return "Files"
+}
+
 type File string
 
-func (f *File) UnmarshalFlag(value string) error {
+func (f File) MarshalYAML() (interface{}, error) {
+	return string(f), nil
+}
+
+func (f *File) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var value string
+	err := unmarshal(&value)
+	if err != nil {
+		return err
+	}
+
+	if value != "" {
+		return f.Set(value)
+	}
+
+	return nil
+}
+
+// Can be removed once flags are deprecated
+func (f *File) Set(value string) error {
 	stat, err := os.Stat(value)
 	if err != nil {
 		return err
@@ -28,6 +82,20 @@ func (f *File) UnmarshalFlag(value string) error {
 	return nil
 }
 
-func (f File) Path() string {
-	return string(f)
+// Can be removed once flags are deprecated
+func (f *File) String() string {
+	if f == nil {
+		return ""
+	}
+
+	return string(*f)
+}
+
+// Can be removed once flags are deprecated
+func (f *File) Type() string {
+	return "File"
+}
+
+func (f *File) Path() string {
+	return string(*f)
 }

go.mod

@@ -0,0 +1,12 @@
+module github.com/concourse/flag
+
+go 1.16
+
+require (
+	code.cloudfoundry.org/lager v2.0.0+incompatible
+	github.com/dgrijalva/jwt-go v3.2.0+incompatible
+	github.com/onsi/ginkgo v1.15.2
+	github.com/onsi/gomega v1.11.0
+	golang.org/x/crypto v0.0.0-20210322153248-0c34fe9e7dc2
+	gopkg.in/yaml.v2 v2.4.0
+)