add tls support to command line tool

chenlujjj · Oct 14, 2020 · 2604bb2 · 2604bb2
1 parent 942345e
commit 2604bb2
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 10 deletions.
diff --git a/tools/asynq/cmd/root.go b/tools/asynq/cmd/root.go
@@ -5,6 +5,7 @@
 package cmd
 
 import (
+	"crypto/tls"
 	"fmt"
 	"io"
 	"os"
@@ -31,6 +32,7 @@ var (
 
 	useRedisCluster bool
 	clusterAddrs    string
+	tlsServerName   string
 )
 
 // rootCmd represents the base command when called without any subcommands
@@ -74,12 +76,15 @@ func init() {
 	rootCmd.PersistentFlags().StringVar(&clusterAddrs, "cluster_addrs",
 		"127.0.0.1:7000,127.0.0.1:7001,127.0.0.1:7002,127.0.0.1:7003,127.0.0.1:7004,127.0.0.1:7005",
 		"list of comma-separated redis server addresses")
+	rootCmd.PersistentFlags().StringVar(&tlsServerName, "tls_server",
+		"", "server name for TLS validation")
 	// Bind flags with config.
 	viper.BindPFlag("uri", rootCmd.PersistentFlags().Lookup("uri"))
 	viper.BindPFlag("db", rootCmd.PersistentFlags().Lookup("db"))
 	viper.BindPFlag("password", rootCmd.PersistentFlags().Lookup("password"))
 	viper.BindPFlag("cluster", rootCmd.PersistentFlags().Lookup("cluster"))
 	viper.BindPFlag("cluster_addrs", rootCmd.PersistentFlags().Lookup("cluster_addrs"))
+	viper.BindPFlag("tls_server", rootCmd.PersistentFlags().Lookup("tls_server"))
 }
 
 // initConfig reads in config file and ENV variables if set.
@@ -114,14 +119,16 @@ func createRDB() *rdb.RDB {
 	if useRedisCluster {
 		addrs := strings.Split(viper.GetString("cluster_addrs"), ",")
 		c = redis.NewClusterClient(&redis.ClusterOptions{
-			Addrs:    addrs,
-			Password: viper.GetString("password"),
+			Addrs:     addrs,
+			Password:  viper.GetString("password"),
+			TLSConfig: getTLSConfig(),
 		})
 	} else {
 		c = redis.NewClient(&redis.Options{
-			Addr:     viper.GetString("uri"),
-			DB:       viper.GetInt("db"),
-			Password: viper.GetString("password"),
+			Addr:      viper.GetString("uri"),
+			DB:        viper.GetInt("db"),
+			Password:  viper.GetString("password"),
+			TLSConfig: getTLSConfig(),
 		})
 	}
 	return rdb.NewRDB(c)
@@ -133,19 +140,29 @@ func createInspector() *asynq.Inspector {
 	if useRedisCluster {
 		addrs := strings.Split(viper.GetString("cluster_addrs"), ",")
 		connOpt = asynq.RedisClusterClientOpt{
-			Addrs:    addrs,
-			Password: viper.GetString("password"),
+			Addrs:     addrs,
+			Password:  viper.GetString("password"),
+			TLSConfig: getTLSConfig(),
 		}
 	} else {
 		connOpt = asynq.RedisClientOpt{
-			Addr:     viper.GetString("uri"),
-			DB:       viper.GetInt("db"),
-			Password: viper.GetString("password"),
+			Addr:      viper.GetString("uri"),
+			DB:        viper.GetInt("db"),
+			Password:  viper.GetString("password"),
+			TLSConfig: getTLSConfig(),
 		}
 	}
 	return asynq.NewInspector(connOpt)
 }
 
+func getTLSConfig() *tls.Config {
+	tlsServer := viper.GetString("tls_server")
+	if tlsServer == "" {
+		return nil
+	}
+	return &tls.Config{ServerName: tlsServer}
+}
+
 // printTable is a helper function to print data in table format.
 //
 // cols is a list of headers and printRow specifies how to print rows.

diff --git a/tools/go.sum b/tools/go.sum
@@ -108,6 +108,8 @@ github.com/prometheus/common v0.4.0/go.mod h1:TNfzLD0ON7rHzMJeJkieUDPYmFC7Snx/y8
 github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
 github.com/prometheus/procfs v0.0.0-20190507164030-5867b95ac084/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsTZCD3I8kEA=
 github.com/prometheus/tsdb v0.7.1/go.mod h1:qhTCs0VvXwvX/y3TZrWD7rabWM+ijKTux40TwIPHuXU=
+github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs=
+github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro=
 github.com/rogpeppe/fastuuid v0.0.0-20150106093220-6724a57986af/go.mod h1:XWv6SoW27p1b0cqNHllgS5HIMJraePCO15w5zCzIWYg=
 github.com/russross/blackfriday v1.5.2/go.mod h1:JO/DiYxRf+HjHt06OyowR9PTA263kcR/rfWxYHBV53g=
 github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=