Add a configurable timeout to the module configuration (#55)

README.md

@@ -160,6 +160,9 @@ modules: [<module>]
 # The protocol over which the probe will take place (https, tcp)
 prober: <prober_string>
 
+# How long the probe will wait before giving up.
+[ timeout: <duration> ]
+
 # Configuration for TLS
 [ tls_config: <tls_config> ]
 

config/config.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net/url"
 	"os"
+	"time"
 
 	"github.com/prometheus/common/config"
 	yaml "gopkg.in/yaml.v3"
@@ -50,6 +51,7 @@ type Config struct {
 
 type Module struct {
 	Prober    string           `yaml:"prober,omitempty"`
+	Timeout   time.Duration    `yaml:"timeout,omitempty"`
 	TLSConfig config.TLSConfig `yaml:"tls_config,omitempty"`
 	HTTPS     HTTPSProbe       `yaml:"https,omitempty"`
 	TCP       TCPProbe         `yaml:"tcp,omitempty"`

examples/ssl_exporter.yaml

@@ -9,6 +9,9 @@ modules:
     prober: https
     https:
       proxy_url: "socks5://localhost:8123"
+  https_timeout:
+    prober: https
+    timeout: 3s
   tcp:
     prober: tcp
   tcp_servername:

ssl_exporter.go

@@ -227,24 +227,27 @@ func probeHandler(w http.ResponseWriter, r *http.Request, conf *config.Config) {
 		return
 	}
 
-	// The following timeout block was taken wholly from the blackbox exporter
-	//   https://github.com/prometheus/blackbox_exporter/blob/master/main.go
-	var timeoutSeconds float64
-	if v := r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"); v != "" {
-		var err error
-		timeoutSeconds, err = strconv.ParseFloat(v, 64)
-		if err != nil {
-			http.Error(w, fmt.Sprintf("Failed to parse timeout from Prometheus header: %s", err), http.StatusInternalServerError)
-			return
+	timeout := module.Timeout
+	if timeout == 0 {
+		// The following timeout block was taken wholly from the blackbox exporter
+		//   https://github.com/prometheus/blackbox_exporter/blob/master/main.go
+		var timeoutSeconds float64
+		if v := r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"); v != "" {
+			var err error
+			timeoutSeconds, err = strconv.ParseFloat(v, 64)
+			if err != nil {
+				http.Error(w, fmt.Sprintf("Failed to parse timeout from Prometheus header: %s", err), http.StatusInternalServerError)
+				return
+			}
+		} else {
+			timeoutSeconds = 10
+		}
+		if timeoutSeconds == 0 {
+			timeoutSeconds = 10
 		}
-	} else {
-		timeoutSeconds = 10
-	}
-	if timeoutSeconds == 0 {
-		timeoutSeconds = 10
-	}
 
-	timeout := time.Duration((timeoutSeconds) * 1e9)
+		timeout = time.Duration((timeoutSeconds) * 1e9)
+	}
 
 	target := r.URL.Query().Get("target")
 	if target == "" {

ssl_exporter_test.go

@@ -72,6 +72,49 @@ func TestProbeHandlerHTTPS(t *testing.T) {
 	}
 }
 
+// TestProbeHandlerHTTPSTimeout tests that the probe respects the timeout set in
+// the module configuration
+func TestProbeHandlerHTTPSTimeout(t *testing.T) {
+	server, _, _, caFile, teardown, err := test.SetupHTTPSServer()
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
+	defer teardown()
+
+	server.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(3 * time.Second)
+		fmt.Fprintln(w, "Hello world")
+	})
+
+	server.StartTLS()
+	defer server.Close()
+
+	conf := &config.Config{
+		Modules: map[string]config.Module{
+			"https": config.Module{
+				Prober:  "https",
+				Timeout: 1 * time.Second,
+				TLSConfig: pconfig.TLSConfig{
+					CAFile: caFile,
+				},
+			},
+		},
+	}
+
+	rr, err := probe(server.URL, "https", conf)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
+
+	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0"); !ok {
+		t.Errorf("expected `ssl_tls_connect_success 0`")
+	}
+
+	if ok := strings.Contains(rr.Body.String(), "ssl_prober{prober=\"https\"} 1"); !ok {
+		t.Errorf("expected `ssl_prober{prober=\"https\"} 1`")
+	}
+}
+
 // TestProbeHandlerHTTPSVerifiedChains checks that metrics are generated
 // correctly for the verified chains
 func TestProbeHandlerHTTPSVerifiedChains(t *testing.T) {
@@ -319,6 +362,45 @@ func TestProbeHandlerTCP(t *testing.T) {
 	}
 }
 
+// TestProbeHandlerTCPTimeout tests that the probe respects the timeout set in
+// the module configuration
+func TestProbeHandlerTCPTimeout(t *testing.T) {
+	server, _, _, caFile, teardown, err := test.SetupTCPServer()
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
+	defer teardown()
+
+	server.StartTLSWait(3 * time.Second)
+	defer server.Close()
+
+	conf := &config.Config{
+		Modules: map[string]config.Module{
+			"tcp": config.Module{
+				Prober:  "tcp",
+				Timeout: 1 * time.Second,
+				TLSConfig: pconfig.TLSConfig{
+					CAFile: caFile,
+				},
+			},
+		},
+	}
+
+	rr, err := probe(server.Listener.Addr().String(), "tcp", conf)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
+
+	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0"); !ok {
+		t.Errorf("expected `ssl_tls_connect_success 0`")
+	}
+
+	if ok := strings.Contains(rr.Body.String(), "ssl_prober{prober=\"tcp\"} 1"); !ok {
+		t.Errorf("expected `ssl_prober{prober=\"tcp\"} 1`")
+	}
+
+}
+
 // TestProbeHandlerTCPVerifiedChains checks that metrics are generated
 // correctly for the verified chains
 func TestProbeHandlerTCPVerifiedChains(t *testing.T) {

test/tcp.go

@@ -39,6 +39,30 @@ func (t *TCPServer) StartTLS() {
 	}()
 }
 
+// StartTLSWait starts a listener and waits for duration 'd' before performing
+// the TLS handshake
+func (t *TCPServer) StartTLSWait(d time.Duration) {
+	go func() {
+		ln := tls.NewListener(t.Listener, t.TLS)
+		conn, err := ln.Accept()
+		if err != nil {
+			panic(fmt.Sprintf("Error accepting on socket: %s", err))
+		}
+		defer conn.Close()
+
+		time.Sleep(d)
+
+		if err := conn.(*tls.Conn).Handshake(); err != nil {
+			log.Errorln(err)
+		} else {
+			// Send some bytes before terminating the connection.
+			fmt.Fprintf(conn, "Hello World!\n")
+		}
+
+		t.stopCh <- struct{}{}
+	}()
+}
+
 // StartSMTP starts a listener that negotiates a TLS connection with an smtp
 // client using STARTTLS
 func (t *TCPServer) StartSMTP() {