talos

1
// This Source Code Form is subject to the terms of the Mozilla Public
2
// License, v. 2.0. If a copy of the MPL was not distributed with this
3
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4

5
package argsbuilder
6

7
import (
8
	"fmt"
9
	"sort"
10
	"strings"
11

12
	"github.com/siderolabs/gen/maps"
13
)
14

15
// Key represents an arg key.
16
type Key = string
17

18
// Value represents an arg value.
19
type Value = string
20

21
// Args represents a set of args.
22
type Args map[Key]Value
23

24
// MustMerge implements the ArgsBuilder interface.
25
func (a Args) MustMerge(args Args, setters ...MergeOption) {
26
	if err := a.Merge(args, setters...); err != nil {
27
		panic(err)
28
	}
29
}
30

31
// Merge implements the ArgsBuilder interface.
32
//
33
//nolint:gocyclo
34
func (a Args) Merge(args Args, setters ...MergeOption) error {
35
	var opts MergeOptions
36

37
	for _, s := range setters {
38
		s(&opts)
39
	}
40

41
	policies := opts.Policies
42
	if policies == nil {
43
		policies = MergePolicies{}
44
	}
45

46
	for key, val := range args {
47
		policy := policies[key]
48

49
		switch policy {
50
		case MergeDenied:
51
			return NewDenylistError(key)
52
		case MergeAdditive:
53
			values := strings.Split(a[key], ",")
54
			definedValues := map[string]struct{}{}
55

56
			i := 0
57

58
			for _, v := range values {
59
				definedValues[strings.TrimSpace(v)] = struct{}{}
60

61
				if v != "" {
62
					values[i] = v
63
					i++
64
				}
65
			}
66

67
			values = values[:i]
68

69
			for _, v := range strings.Split(val, ",") {
70
				v = strings.TrimSpace(v)
71
				if _, defined := definedValues[v]; !defined {
72
					values = append(values, v)
73
				}
74
			}
75

76
			a[key] = strings.Join(values, ",")
77
		case MergeOverwrite:
78
			a[key] = val
79
		}
80
	}
81

82
	return nil
83
}
84

85
// Set implements the ArgsBuilder interface.
86
func (a Args) Set(k, v Key) ArgsBuilder {
87
	a[k] = v
88

89
	return a
90
}
91

92
// Args implements the ArgsBuilder interface.
93
func (a Args) Args() []string {
94
	keys := maps.Keys(a)
95
	sort.Strings(keys)
96

97
	args := []string{}
98

99
	for _, key := range keys {
100
		args = append(args, fmt.Sprintf("--%s=%s", key, a[key]))
101
	}
102

103
	return args
104
}
105

106
// Get returns an args value.
107
func (a Args) Get(k Key) Value {
108
	return a[k]
109
}
110

111
// Contains checks if an arg key exists.
112
func (a Args) Contains(k Key) bool {
113
	_, ok := a[k]
114

115
	return ok
116
}
117

118
// DenyListError represents an error indicating that an argument was supplied
119
// that is not allowed.
120
type DenyListError struct {
121
	s string
122
}
123

124
// NewDenylistError returns a DenyListError.
125
func NewDenylistError(s string) error {
126
	return &DenyListError{s}
127
}
128

129
// Error implements the Error interface.
130
func (b *DenyListError) Error() string {
131
	return fmt.Sprintf("extra arg %q is not allowed", b.s)
132
}
133