From 29845c64b8fadce7b4d49bacd2506f0973203c3a Mon Sep 17 00:00:00 2001
From: MGSousa <mikesilva.290@gmail.com>
Date: Tue, 19 Dec 2023 18:04:41 +0000
Subject: [PATCH] Allow to read from env a service-uri var, security
 improvements

---
 exporter/exporter.go |  1 -
 exporter/http.go     |  7 ++++++-
 main.go              | 36 ++++++++++++++++++++++++++----------
 3 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/exporter/exporter.go b/exporter/exporter.go
index 231803f..60b1132 100644
--- a/exporter/exporter.go
+++ b/exporter/exporter.go
@@ -106,7 +106,6 @@ func (e *Exporter) process(ch chan<- prometheus.Metric) {
 				value = 0
 			}
 
-			// check last filed for the specific ValueType
 			// TODO: custom filter by a map file
 			if len(m) > 3 {
 				mType = prometheus.UntypedValue
diff --git a/exporter/http.go b/exporter/http.go
index 450cc49..16549d6 100644
--- a/exporter/http.go
+++ b/exporter/http.go
@@ -47,7 +47,12 @@ func FetchMetrics(uri string) (any, error) {
 }
 
 func get(uri string) ([]byte, error) {
-	res, err := http.Get(uri)
+	req, err := http.NewRequest(http.MethodGet, uri, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	res, err := http.DefaultClient.Do(req)
 	if err != nil {
 		log.Errorf("Could not fetch metrics for endpoint of target: %s", uri)
 		return nil, err
diff --git a/main.go b/main.go
index 204d33c..8dfab79 100644
--- a/main.go
+++ b/main.go
@@ -16,12 +16,12 @@ import (
 )
 
 var (
-	internalVersion string
+	internalVersion, serviceEndpoint string
 
 	listenAddress        = flag.String("listen-address", ":19100", "Address to listen on for telemetry.")
 	metricsPath          = flag.String("telemetry-path", "/metrics", "Base path under which to expose metrics.")
 	serviceName          = flag.String("service-name", "", "Service name to reference")
-	serviceUri           = flag.String("service-uri", "http://localhost:5066", "HTTP address of the service.")
+	serviceUri           = flag.String("service-uri", "", "Endpoint address of the service.")
 	serviceMetricsPath   = flag.String("service-metrics-path", "metrics", "Service path to scrape metrics from.")
 	serviceVersionScrape = flag.Bool("service-version-scrape", false, "Enable whether the service will be internally scraped for fetching remote build version or not")
 	debugLevel           = flag.Bool("debug", false, "Enable debug mode")
@@ -32,18 +32,29 @@ func main() {
 
 	name := *serviceName
 	if strings.Trim(name, " ") == "" {
-		log.Fatalln("Service name not known! Specify by -service-name SERVICE")
+		log.Fatalln("Service name not set! Specify by -service-name SERVICE")
 	}
 
+	// check if service endpoint is set
+	// if not try to fetch from env
+	serviceEndpoint = *serviceUri
+	if serviceEndpoint == "" {
+		var exists bool
+		if serviceEndpoint, exists = os.LookupEnv("SERVICE_ENDPOINT"); !exists || serviceEndpoint == "" {
+			log.Fatalln("Service URI is not set! Specify either a '-service-uri' flag OR an environment variable 'SERVICE_ENDPOINT'")
+		}
+	}
+
+	// enable debug mode if required
 	if *debugLevel {
 		log.SetLevel(log.DebugLevel)
 	}
 
 	log.Info("Check if target is reachable...")
 	if *serviceVersionScrape {
-		internalVersion = checkEndpoint(*serviceUri)
+		internalVersion = checkEndpoint(serviceEndpoint)
 	} else {
-		checkEndpoint(*serviceUri)
+		checkEndpoint(serviceEndpoint)
 	}
 	log.Info("Target endpoint is reachable")
 
@@ -54,7 +65,7 @@ func main() {
 	registry.MustRegister(versionMetric)
 
 	// register remote service metrics
-	exporter := exporter.NewCollector(name, fmt.Sprintf("%s/%s", *serviceUri, *serviceMetricsPath), internalVersion)
+	exporter := exporter.NewCollector(name, fmt.Sprintf("%s/%s", serviceEndpoint, *serviceMetricsPath), internalVersion)
 	registry.MustRegister(exporter)
 
 	http.Handle(*metricsPath, promhttp.HandlerFor(
@@ -68,7 +79,7 @@ func main() {
 
 	srv := &http.Server{
 		Addr:              *listenAddress,
-		ReadHeaderTimeout: 5 * time.Second,
+		ReadHeaderTimeout: 30 * time.Second,
 	}
 	if err := srv.ListenAndServe(); err != nil {
 		log.WithFields(log.Fields{
@@ -77,10 +88,12 @@ func main() {
 	}
 }
 
-// TODO: improve this method
+// check if service endpoint is reachable
+// increases duration by 5secs for the next ticket if unavailable
 func checkEndpoint(endpoint string) string {
+	var duration = 5
 	stopCh := make(chan bool)
-	t := time.NewTicker(2 * time.Second)
+	t := time.NewTicker(time.Duration(duration) * time.Second)
 
 	stats := &exporter.HttpStats{}
 
@@ -91,7 +104,10 @@ discovery:
 			if stats = exporter.FetchStats(endpoint); stats != nil {
 				break discovery
 			}
-			log.Errorln("base endpoint not available, retrying in 2s")
+			log.Errorf("service endpoint not available, retrying in %ds", duration)
+
+			t = time.NewTicker(time.Duration(duration) * time.Second)
+			duration = duration + 5
 			continue
 
 		case <-stopCh: