forked from fgeller/kt
-
Notifications
You must be signed in to change notification settings - Fork 0
/
common.go
298 lines (264 loc) · 7.28 KB
/
common.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
package main
import (
"bufio"
"crypto/tls"
"crypto/x509"
"encoding/base64"
"encoding/hex"
"encoding/json"
"flag"
"fmt"
"io"
"io/ioutil"
"os"
"os/signal"
"os/user"
"regexp"
"strings"
"sync"
"github.com/Shopify/sarama"
"golang.org/x/crypto/ssh/terminal"
)
var (
invalidClientIDCharactersRegExp = regexp.MustCompile(`[^a-zA-Z0-9_-]`)
)
func listenForInterrupt(q chan struct{}) {
signals := make(chan os.Signal, 1)
signal.Notify(signals, os.Interrupt)
sig := <-signals
fmt.Fprintf(os.Stderr, "received signal %s\n", sig)
close(q)
}
var defaultKafkaVersion = sarama.V2_0_0_0
type commonFlags struct {
verbose bool
version sarama.KafkaVersion
tlsRequested bool
tlsCA string
tlsCert string
tlsCertKey string
brokerStrs []string
}
func (f *commonFlags) brokers() []string {
brokers := append([]string{}, f.brokerStrs...)
for i, b := range f.brokerStrs {
if !strings.Contains(b, ":") {
brokers[i] = b + ":9092"
}
}
return brokers
}
func (f *commonFlags) addFlags(flags *flag.FlagSet) {
f.brokerStrs = []string{"localhost"}
f.version = defaultKafkaVersion
flags.Var(listFlag{&f.brokerStrs}, "brokers", "Comma-separated list of brokers. Each broker definition may optionally contain a port number. The port defaults to 9092 when omitted.")
flags.Var(kafkaVersionFlag{v: &f.version}, "version", "Kafka protocol version")
flags.BoolVar(&f.tlsRequested, "tls", false, "Request server-side TLS without client-side.")
flags.StringVar(&f.tlsCA, "tlsca", "", "Path to the TLS certificate authority file")
flags.StringVar(&f.tlsCert, "tlscert", "", "Path to the TLS client certificate file")
flags.StringVar(&f.tlsCertKey, "tlscertkey", "", "Path to the TLS client certificate key file")
flags.BoolVar(&f.verbose, "verbose", false, "More verbose logging to stderr.")
}
func (f *commonFlags) saramaConfig(name string) (*sarama.Config, error) {
cfg := sarama.NewConfig()
cfg.Version = f.version
usr, err := user.Current()
var username string
if err != nil {
warningf("failed to read current user name: %v", err)
username = "anon"
} else {
username = usr.Username
}
cfg.ClientID = "kt-" + name + "-" + sanitizeUsername(username)
tlsConfig, err := setUpCerts(f.tlsCert, f.tlsCA, f.tlsCertKey)
if err != nil {
return nil, fmt.Errorf("cannot set up certificates: %v", err)
}
if tlsConfig == nil && f.tlsRequested {
// client-side not configured, but server-side TLS requested
tlsConfig = &tls.Config{}
}
if tlsConfig != nil {
cfg.Net.TLS.Enable = true
cfg.Net.TLS.Config = tlsConfig
}
if f.verbose {
fmt.Fprintf(os.Stderr, "sarama client configuration %#v\n", cfg)
}
return cfg, nil
}
type listFlag struct {
v *[]string
}
func (v listFlag) String() string {
if v.v == nil {
return ""
}
return strings.Join(*v.v, ",")
}
func (v listFlag) Set(s string) error {
if s == "" {
*v.v = nil
} else {
*v.v = strings.Split(s, ",")
}
return nil
}
type kafkaVersionFlag struct {
v *sarama.KafkaVersion
}
func (v kafkaVersionFlag) String() string {
if v.v == nil {
return ""
}
return v.v.String()
}
func (v kafkaVersionFlag) Set(s string) error {
vers, err := sarama.ParseKafkaVersion(strings.TrimPrefix(s, "v"))
if err != nil {
return fmt.Errorf("invalid kafka version %q: %v", s, err)
}
*v.v = vers
return nil
}
func logClose(name string, c io.Closer) {
if err := c.Close(); err != nil {
warningf("failed to close %#v: %v", name, err)
}
}
type printer struct {
mu sync.Mutex
marshal func(interface{}) ([]byte, error)
}
func newPrinter(pretty bool) *printer {
marshal := json.Marshal
if pretty && terminal.IsTerminal(1) {
marshal = func(i interface{}) ([]byte, error) { return json.MarshalIndent(i, "", " ") }
}
return &printer{
marshal: marshal,
}
}
func (p *printer) print(val interface{}) {
p.mu.Lock()
defer p.mu.Unlock()
buf, err := p.marshal(val)
if err != nil {
warningf("failed to marshal output %#v: %v", val, err)
}
fmt.Println(string(buf))
}
func warningf(f string, a ...interface{}) {
fmt.Fprintf(os.Stderr, "hkt: warning: %s\n", fmt.Sprintf(f, a...))
}
func readStdinLines(max int, out chan string) {
scanner := bufio.NewScanner(os.Stdin)
scanner.Buffer(nil, max)
for scanner.Scan() {
out <- scanner.Text()
}
if err := scanner.Err(); err != nil {
warningf("error reading standard input: %v", err)
}
close(out)
}
func sanitizeUsername(u string) string {
// Windows user may have format "DOMAIN|MACHINE\username", remove domain/machine if present
s := strings.Split(u, "\\")
u = s[len(s)-1]
// Windows account can contain spaces or other special characters not supported
// in client ID. Keep the bare minimum and ditch the rest.
return invalidClientIDCharactersRegExp.ReplaceAllString(u, "")
}
// setUpCerts takes the paths to a tls certificate, CA, and certificate key in
// a PEM format and returns a constructed tls.Config object.
func setUpCerts(certPath, caPath, keyPath string) (*tls.Config, error) {
if certPath == "" && caPath == "" && keyPath == "" {
return nil, nil
}
if certPath == "" || caPath == "" || keyPath == "" {
return nil, fmt.Errorf("certificate, CA and key path are required - got cert=%#v ca=%#v key=%#v", certPath, caPath, keyPath)
}
caString, err := ioutil.ReadFile(caPath)
if err != nil {
return nil, err
}
caPool := x509.NewCertPool()
ok := caPool.AppendCertsFromPEM(caString)
if !ok {
return nil, fmt.Errorf("unable to add cert at %s to certificate pool", caPath)
}
clientCert, err := tls.LoadX509KeyPair(certPath, keyPath)
if err != nil {
return nil, err
}
bundle := &tls.Config{
RootCAs: caPool,
Certificates: []tls.Certificate{clientCert},
}
bundle.BuildNameToCertificate()
return bundle, nil
}
// setFlagsFromEnv sets unset flags in fs from environment
// variables as specified by the flags map, which maps
// from flag name to the environment variable for that name.
//
// If a flag f is part of fs but has not been explicitly set on the
// command line, and flags[f] exists, then it will
// be set from os.Getenv(flags[f]).
func setFlagsFromEnv(fs *flag.FlagSet, flags map[string]string) error {
set := make(map[string]bool)
fs.Visit(func(f *flag.Flag) {
set[f.Name] = true
})
for name, env := range flags {
f := fs.Lookup(name)
if f == nil {
panic(fmt.Errorf("flag %q ($%s) not found", name, env))
}
if set[name] {
continue
}
if v := os.Getenv(env); v != "" {
if err := f.Value.Set(v); err != nil {
return fmt.Errorf("cannot parse $%s as -%s flag value: %v", env, name, err)
}
}
}
return nil
}
func decoderForType(typ string) (func(m json.RawMessage) ([]byte, error), error) {
var dec func(s string) ([]byte, error)
switch typ {
case "json":
// Easy case - we already have the JSON-marshaled data.
return func(m json.RawMessage) ([]byte, error) {
return m, nil
}, nil
case "hex":
dec = hex.DecodeString
case "base64":
dec = base64.StdEncoding.DecodeString
case "string":
dec = func(s string) ([]byte, error) {
return []byte(s), nil
}
default:
return nil, fmt.Errorf(`unsupported decoder %#v, only json, string, hex and base64 are supported`, typ)
}
return func(m json.RawMessage) ([]byte, error) {
var s string
if err := json.Unmarshal(m, &s); err != nil {
return nil, err
}
return dec(s)
}, nil
}
var nullJSON = json.RawMessage("null")
func min(x, y int64) int64 {
if x < y {
return x
}
return y
}