From 4abaaa653234d686d4bfe646f8b02508aee586b1 Mon Sep 17 00:00:00 2001
From: rodrigom <rodrigom@btcx.com>
Date: Mon, 15 Jun 2020 10:44:42 +1000
Subject: [PATCH] Add support for TLS

---
 Dockerfile   |  7 ++++++-
 main.go      | 17 +++++++++++++++++
 zookeeper.go | 19 ++++++++++++++++++-
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index bbea9a0..e071d05 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,3 +1,7 @@
+FROM alpine:3.10 AS certs
+RUN apk update \
+    && apk add ca-certificates
+
 FROM golang:1.10 AS builder
 WORKDIR /go/src/github.com/carlpett/zookeeper_exporter/
 COPY . .
@@ -6,5 +10,6 @@ RUN make build
 FROM scratch
 EXPOSE 9141
 USER 1000
-ENTRYPOINT ["/zookeeper_exporter"]
 COPY --from=builder /go/src/github.com/carlpett/zookeeper_exporter/zookeeper_exporter /zookeeper_exporter
+COPY --from=certs /etc/ssl/certs /etc/ssl/certs
+ENTRYPOINT ["/zookeeper_exporter"]
diff --git a/main.go b/main.go
index 026ce05..7bc07e4 100644
--- a/main.go
+++ b/main.go
@@ -21,12 +21,29 @@ func init() {
 	}
 	logLevel = parsedLevel
 
+	if *enableTLS && (*certPath == "" || *certKeyPath == "") {
+		log.Fatal("-enable-tls requires -cert and -cert-key")
+	}
+
+	if *logJSON {
+		log.SetFormatter(&log.JSONFormatter{})
+	} else {
+		log.SetFormatter(&log.TextFormatter{
+			DisableColors: true,
+			FullTimestamp: true,
+		})
+	}
+
 	prometheus.MustRegister(version.NewCollector("zookeeper_exporter"))
 }
 
 var (
 	logLevel      log.Level = log.InfoLevel
+	logJSON                 = flag.Bool("log-json", false, "Log output as JSON")
 	bindAddr                = flag.String("bind-addr", ":9141", "bind address for the metrics server")
+	enableTLS               = flag.Bool("enable-tls", false, "Connect to zookeeper using TLS. Requires -cert and -cert-key")
+	certPath                = flag.String("cert", "", "path to certificate including any intermediaries")
+	certKeyPath             = flag.String("cert-key", "", "path to certificate key")
 	metricsPath             = flag.String("metrics-path", "/metrics", "path to metrics endpoint")
 	zookeeperAddr           = flag.String("zookeeper", "localhost:2181", "host:port for zookeeper socket")
 	rawLevel                = flag.String("log-level", "info", "log level")
diff --git a/zookeeper.go b/zookeeper.go
index d8a6f1f..6815633 100644
--- a/zookeeper.go
+++ b/zookeeper.go
@@ -3,6 +3,7 @@ package main
 import (
 	"bufio"
 	"bytes"
+	"crypto/tls"
 	"net"
 	"strconv"
 	"strings"
@@ -190,7 +191,7 @@ const (
 func sendZkCommand(fourLetterWord string) (string, bool) {
 	log.Debugf("Connecting to Zookeeper at %s", *zookeeperAddr)
 
-	conn, err := net.Dial("tcp", *zookeeperAddr)
+	conn, err := zkConnect()
 	if err != nil {
 		log.WithFields(log.Fields{"error": err}).Error("Unable to open connection to Zookeeper")
 		return "", false
@@ -223,3 +224,19 @@ func sendZkCommand(fourLetterWord string) (string, bool) {
 
 	return buffer.String(), true
 }
+
+func zkConnect() (net.Conn, error) {
+	if *enableTLS {
+		log.Debugf("TLS certificate: %s key: %s", *certPath, *certKeyPath)
+		cert, err := tls.LoadX509KeyPair(*certPath, *certKeyPath)
+		if err != nil {
+			log.WithFields(log.Fields{"error": err}).Error("Unable to read TLS cert or key")
+			return nil, err
+		}
+		return tls.Dial("tcp", *zookeeperAddr, &tls.Config{
+			Certificates: []tls.Certificate{cert},
+		})
+	}
+
+	return net.Dial("tcp", *zookeeperAddr)
+}