From 7840f7855f75d7159686aa3f38ab0f5bcb502859 Mon Sep 17 00:00:00 2001
From: Conor Broderick <conor.broderick@robustperception.io>
Date: Tue, 11 Jul 2017 17:14:07 +0100
Subject: [PATCH] Expose HTTP version of probe (#189)

---
 CONFIGURATION.md |  3 +++
 config.go        |  1 +
 example.yml      |  1 +
 http.go          | 29 ++++++++++++++++++++++++++++-
 http_test.go     | 27 +++++++++++++++++++++++++++
 5 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/CONFIGURATION.md b/CONFIGURATION.md
index 20c665a..c036cdb 100644
--- a/CONFIGURATION.md
+++ b/CONFIGURATION.md
@@ -38,6 +38,9 @@ The other placeholders are specified separately.
   # Accepted status codes for this probe. Defaults to 2xx.
   [ valid_status_codes: <string>, ... | default = "2xx" ]
 
+  # Accepted HTTP versions for this probe.
+  [ valid_http_versions: <string>, ... ]
+
   # The HTTP method the probe will use.
   [ method: <string> | default = "GET" ]
 
diff --git a/config.go b/config.go
index 5e29788..91880da 100644
--- a/config.go
+++ b/config.go
@@ -36,6 +36,7 @@ type Module struct {
 type HTTPProbe struct {
 	// Defaults to 2xx.
 	ValidStatusCodes       []int                   `yaml:"valid_status_codes,omitempty"`
+	ValidHTTPVersions      []string                `yaml:"valid_http_versions,omitempty"`
 	PreferredIPProtocol    string                  `yaml:"preferred_ip_protocol,omitempty"`
 	NoFollowRedirects      bool                    `yaml:"no_follow_redirects,omitempty"`
 	FailIfSSL              bool                    `yaml:"fail_if_ssl,omitempty"`
diff --git a/example.yml b/example.yml
index b2bcc46..891b59d 100644
--- a/example.yml
+++ b/example.yml
@@ -3,6 +3,7 @@ modules:
     prober: http
     timeout: 5s
     http:
+      valid_http_versions: ["HTTP/1.1", "HTTP/2"]
       valid_status_codes: []  # Defaults to 2xx
       method: GET
       headers:
diff --git a/http.go b/http.go
index cb91288..217a2f6 100644
--- a/http.go
+++ b/http.go
@@ -21,6 +21,7 @@ import (
 	"net/http"
 	"net/url"
 	"regexp"
+	"strconv"
 	"strings"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -85,12 +86,18 @@ func probeHTTP(target string, module Module, registry *prometheus.Registry) (suc
 			Name: "probe_ssl_earliest_cert_expiry",
 			Help: "Returns earliest SSL cert expiry in unixtime",
 		})
+
+		probeHTTPVersionGauge = prometheus.NewGauge(prometheus.GaugeOpts{
+			Name: "probe_http_version",
+			Help: "Returns the version of HTTP of the probe response",
+		})
 	)
 
 	registry.MustRegister(contentLengthGauge)
 	registry.MustRegister(redirectsGauge)
 	registry.MustRegister(isSSLGauge)
 	registry.MustRegister(statusCodeGauge)
+	registry.MustRegister(probeHTTPVersionGauge)
 
 	httpConfig := module.HTTP
 
@@ -160,7 +167,6 @@ func probeHTTP(target string, module Module, registry *prometheus.Registry) (suc
 		request.Body = ioutil.NopCloser(strings.NewReader(httpConfig.Body))
 	}
 	resp, err := client.Do(request)
-
 	// Err won't be nil if redirects were turned off. See https://github.com/golang/go/issues/3795
 	if err != nil && resp == nil {
 		log.Warnf("Error for HTTP request to %s: %s", target, err)
@@ -180,6 +186,27 @@ func probeHTTP(target string, module Module, registry *prometheus.Registry) (suc
 		if success && (len(httpConfig.FailIfMatchesRegexp) > 0 || len(httpConfig.FailIfNotMatchesRegexp) > 0) {
 			success = matchRegularExpressions(resp.Body, httpConfig)
 		}
+
+		var httpVersionNumber float64
+		httpVersionNumber, err = strconv.ParseFloat(strings.TrimPrefix(resp.Proto, "HTTP/"), 64)
+		if err != nil {
+			log.Errorf("Error parsing version number from HTTP version: %v", err)
+		}
+		probeHTTPVersionGauge.Set(httpVersionNumber)
+
+		if len(httpConfig.ValidHTTPVersions) != 0 {
+			found := false
+			for _, version := range httpConfig.ValidHTTPVersions {
+				if version == resp.Proto {
+					found = true
+					break
+				}
+			}
+			if !found {
+				success = false
+			}
+		}
+
 	}
 
 	if resp == nil {
diff --git a/http_test.go b/http_test.go
index 976a300..443d689 100644
--- a/http_test.go
+++ b/http_test.go
@@ -58,6 +58,33 @@ func TestHTTPStatusCodes(t *testing.T) {
 	}
 }
 
+func TestValidHTTPVersion(t *testing.T) {
+	tests := []struct {
+		ValidHTTPVersions []string
+		ShouldSucceed     bool
+	}{
+		{[]string{}, true},
+		{[]string{"HTTP/1.1"}, true},
+		{[]string{"HTTP/1.1", "HTTP/2"}, true},
+		{[]string{"HTTP/2"}, false},
+	}
+	for i, test := range tests {
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		}))
+		defer ts.Close()
+		recorder := httptest.NewRecorder()
+		registry := prometheus.NewRegistry()
+		result := probeHTTP(ts.URL,
+			Module{Timeout: time.Second, HTTP: HTTPProbe{
+				ValidHTTPVersions: test.ValidHTTPVersions,
+			}}, registry)
+		body := recorder.Body.String()
+		if result != test.ShouldSucceed {
+			t.Fatalf("Test %v had unexpected result: %s", i, body)
+		}
+	}
+}
+
 func TestRedirectFollowed(t *testing.T) {
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path == "/" {
-- 
2.25.1