bench: allow custom certs with -ca-cert, -client-cert, -client-key

twmb · Oct 17, 2021 · d368d11 · d368d11
1 parent 54f3092
commit d368d11
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 2 deletions.
diff --git a/examples/bench/README.md b/examples/bench/README.md
@@ -88,6 +88,12 @@ a consumer group.
 
 `-tls` if true, sets the benchmark to dial over tls
 
+`-ca-cert` specifies a custom CA to use when dialing (implies `-tls`)
+
+`-client-cert` specifies a client cert to use when dialing (implies `-tls`, requires `-client-key`)
+
+`-client-key` specifies a client key to use when dialing (implies `-tls`, requires `-client-cert`)
+
 `-sasl-method` specifies a SASL method to use when connecting. This supports
 `PLAIN`, `SCRAM-SHA-256`, `SCRAM-SHA-512`, or `AWS_MSK_IAM` (any casing, with
 or without dashes or underscores).

diff --git a/examples/bench/go.mod b/examples/bench/go.mod
@@ -5,6 +5,7 @@ go 1.16
 require (
  github.com/twmb/franz-go v1.0.0
  github.com/twmb/franz-go/plugin/kprom v0.3.0
+ github.com/twmb/tlscfg v1.2.0
 )
 
 replace (

diff --git a/examples/bench/go.sum b/examples/bench/go.sum
@@ -199,6 +199,8 @@ github.com/twmb/franz-go/pkg/kmsg v0.0.0-20210914042331-106aef61b693 h1:5O4u9Lc6
 github.com/twmb/franz-go/pkg/kmsg v0.0.0-20210914042331-106aef61b693/go.mod h1:SxG/xJKhgPu25SamAq0rrucfp7lbzCpEXOC+vH/ELrY=
 github.com/twmb/go-rbtree v1.0.0 h1:KxN7dXJ8XaZ4cvmHV1qqXTshxX3EBvX/toG5+UR49Mg=
 github.com/twmb/go-rbtree v1.0.0/go.mod h1:UlIAI8gu3KRPkXSobZnmJfVwCJgEhD/liWzT5ppzIyc=
+github.com/twmb/tlscfg v1.2.0 h1:WCzLHtmnVJ94+veAO4TLTB1ENx7TPYLkTl4Q6WFF4Vo=
+github.com/twmb/tlscfg v1.2.0/go.mod h1:GameEQddljI+8Es373JfQEBvtI4dCTLKWGJbqT2kErs=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

diff --git a/examples/bench/main.go b/examples/bench/main.go
@@ -15,6 +15,7 @@ import (
  "time"
 
  "github.com/twmb/franz-go/plugin/kprom"
+ "github.com/twmb/tlscfg"
 
  "github.com/twmb/franz-go/pkg/kgo"
  "github.com/twmb/franz-go/pkg/sasl/aws"
@@ -43,7 +44,10 @@ var (
  consume = flag.Bool("consume", false, "if true, consume rather than produce")
  group = flag.String("group", "", "if non-empty, group to use for consuming rather than direct partition consuming (consuming)")
 
- dialTLS = flag.Bool("tls", false, "if true, use tls for connecting")
+ dialTLS = flag.Bool("tls", false, "if true, use tls for connecting (if using well-known TLS certs)")
+ caFile = flag.String("ca-cert", "", "if non-empty, path to CA cert to use for TLS (implies -tls)")
+ certFile = flag.String("client-cert", "", "if non-empty, path to client cert to use for TLS (requires -client-key, implies -tls)")
+ keyFile = flag.String("client-key", "", "if non-empty, path to client key to use for TLS (requires -client-cert, implies -tls)")
 
  saslMethod = flag.String("sasl-method", "", "if non-empty, sasl method to use (must specify all options; supports plain, scram-sha-256, scram-sha-512, aws_msk_iam)")
  saslUser = flag.String("sasl-user", "", "if non-empty, username to use for sasl (must specify all options)")
@@ -75,6 +79,12 @@ func chk(err error, msg string, args ...interface{}) {
 func main() {
  flag.Parse()
 
+ var customTLS bool
+ if *caFile != "" || *certFile != "" || *keyFile != "" {
+ *dialTLS = true
+ customTLS = true
+ }
+
  if *recordBytes <= 0 {
  die("record bytes must be larger than zero")
  }
@@ -147,7 +157,18 @@ func main() {
  }
 
  if *dialTLS {
- opts = append(opts, kgo.Dialer((new(tls.Dialer)).DialContext))
+ if customTLS {
+ tc, err := tlscfg.New(
+ tlscfg.MaybeWithDiskCA(*caFile, tlscfg.ForClient),
+ tlscfg.MaybeWithDiskKeyPair(*certFile, *keyFile),
+ )
+ if err != nil {
+ die("unable to create tls config: %v", err)
+ }
+ opts = append(opts, kgo.DialTLSConfig(tc))
+ } else {
+ opts = append(opts, kgo.DialTLSConfig(new(tls.Config)))
+ }
  }
 
  if *saslMethod != "" || *saslUser != "" || *saslPass != "" {