From dab6a6933f6d848025478794a95ad1a0d6842a01 Mon Sep 17 00:00:00 2001
From: fate0 <fate0@fatezero.org>
Date: Thu, 28 May 2020 22:46:05 +0800
Subject: [PATCH] add support for custom tls config

---
 cmd/cassowary/cli.go | 49 ++++++++++++++++++++++++++++++++++++++++++++
 pkg/client/load.go   |  1 +
 pkg/client/types.go  |  2 ++
 3 files changed, 52 insertions(+)

diff --git a/cmd/cassowary/cli.go b/cmd/cassowary/cli.go
index cfc7cae..fa9c661 100644
--- a/cmd/cassowary/cli.go
+++ b/cmd/cassowary/cli.go
@@ -1,9 +1,12 @@
 package main
 
 import (
+	"crypto/tls"
+	"crypto/x509"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io/ioutil"
 	"os"
 	"strconv"
 
@@ -151,6 +154,27 @@ func validateCLI(c *cli.Context) error {
 		httpMethod = "GET"
 	}
 
+	tlsConfig := new(tls.Config)
+	if c.String("ca") != "" {
+		pemCerts, err := ioutil.ReadFile(c.String("ca"))
+		if err != nil {
+			return err
+		}
+		ca := x509.NewCertPool()
+		if !ca.AppendCertsFromPEM(pemCerts) {
+			return fmt.Errorf("failed to read CA from PEM")
+		}
+		tlsConfig.RootCAs = ca
+	}
+
+	if c.String("cert") != "" && c.String("key") != "" {
+		cert, err := tls.LoadX509KeyPair(c.String("cert"), c.String("key"))
+		if err != nil {
+			return err
+		}
+		tlsConfig.Certificates = []tls.Certificate{cert}
+	}
+
 	cass := &client.Cassowary{
 		FileMode:          false,
 		BaseURL:           c.String("url"),
@@ -159,6 +183,7 @@ func validateCLI(c *cli.Context) error {
 		RequestHeader:     header,
 		Duration:          duration,
 		PromExport:        prometheusEnabled,
+		TLSConfig:         tlsConfig,
 		PromURL:           c.String("prompushgwurl"),
 		Cloudwatch:        c.Bool("cloudwatch"),
 		ExportMetrics:     c.Bool("json-metrics"),
@@ -282,6 +307,18 @@ func runCLI(args []string) {
 					Name:  "disable-keep-alive",
 					Usage: "use this flag to disable http keep-alive",
 				},
+				cli.StringFlag{
+					Name:  "ca",
+					Usage: "certificate authority",
+				},
+				cli.StringFlag{
+					Name:  "cert",
+					Usage: "authentication certificate",
+				},
+				cli.StringFlag{
+					Name:  "key",
+					Usage: "authentication key",
+				},
 			},
 			Action: validateCLIFile,
 		},
@@ -345,6 +382,18 @@ func runCLI(args []string) {
 					Name:  "disable-keep-alive",
 					Usage: "use this flag to disable http keep-alive",
 				},
+				cli.StringFlag{
+					Name:  "ca",
+					Usage: "certificate authority",
+				},
+				cli.StringFlag{
+					Name:  "cert",
+					Usage: "authentication certificate",
+				},
+				cli.StringFlag{
+					Name:  "key",
+					Usage: "authentication key",
+				},
 			},
 			Action: validateCLI,
 		},
diff --git a/pkg/client/load.go b/pkg/client/load.go
index a4be154..649cd64 100644
--- a/pkg/client/load.go
+++ b/pkg/client/load.go
@@ -148,6 +148,7 @@ func (c *Cassowary) Coordinate() (ResultMetrics, error) {
 	c.Client = &http.Client{
 		Timeout: time.Second * time.Duration(c.Timeout),
 		Transport: &http.Transport{
+			TLSClientConfig:     c.TLSConfig,
 			MaxIdleConnsPerHost: 10000,
 			DisableCompression:  false,
 			DisableKeepAlives:   c.DisableKeepAlive,
diff --git a/pkg/client/types.go b/pkg/client/types.go
index 5cc94be..c0e9b2c 100644
--- a/pkg/client/types.go
+++ b/pkg/client/types.go
@@ -1,6 +1,7 @@
 package client
 
 import (
+	"crypto/tls"
 	"net/http"
 
 	"github.com/schollz/progressbar"
@@ -18,6 +19,7 @@ type Cassowary struct {
 	ExportMetricsFile     string
 	PromExport            bool
 	Cloudwatch            bool
+	TLSConfig             *tls.Config
 	PromURL               string
 	RequestHeader         []string
 	URLPaths              []string