From a0554138716afdd2469d0f0db553a8a3a64f5f61 Mon Sep 17 00:00:00 2001
From: conorbroderick <cjayjayb@gmail.com>
Date: Fri, 4 Aug 2017 12:30:00 +0100
Subject: [PATCH] Fixed timeout bug

---
 tcp.go      |  4 ++--
 tcp_test.go | 20 ++++++--------------
 2 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/tcp.go b/tcp.go
index 61dc9fe..5abbb51 100644
--- a/tcp.go
+++ b/tcp.go
@@ -20,7 +20,6 @@ import (
 	"fmt"
 	"net"
 	"regexp"
-	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/config"
@@ -73,7 +72,7 @@ func probeTCP(ctx context.Context, target string, module Module, registry *prome
 		Help: "Indicates if probe failed due to regex",
 	})
 	registry.MustRegister(probeFailedDueToRegex)
-	deadline := time.Now().Add(module.Timeout)
+	deadline, _ := ctx.Deadline()
 	conn, err := dialTCP(ctx, target, module, registry)
 	if err != nil {
 		log.Errorf("Error dialing TCP: %v", err)
@@ -114,6 +113,7 @@ func probeTCP(ctx context.Context, target string, module Module, registry *prome
 				}
 			}
 			if scanner.Err() != nil {
+				log.Errorf("Error reading from connection: %v", scanner.Err().Error())
 				return false
 			}
 			if match == nil {
diff --git a/tcp_test.go b/tcp_test.go
index 0551792..37a3241 100644
--- a/tcp_test.go
+++ b/tcp_test.go
@@ -43,7 +43,7 @@ func TestTCPConnection(t *testing.T) {
 	testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 	registry := prometheus.NewRegistry()
-	if !probeTCP(testCTX, ln.Addr().String(), Module{Timeout: time.Second}, registry) {
+	if !probeTCP(testCTX, ln.Addr().String(), Module{}, registry) {
 		t.Fatalf("TCP module failed, expected success.")
 	}
 	<-ch
@@ -54,7 +54,7 @@ func TestTCPConnectionFails(t *testing.T) {
 	registry := prometheus.NewRegistry()
 	testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
-	if probeTCP(testCTX, ":0", Module{Timeout: time.Second}, registry) {
+	if probeTCP(testCTX, ":0", Module{}, registry) {
 		t.Fatalf("TCP module suceeded, expected failure.")
 	}
 }
@@ -70,7 +70,6 @@ func TestTCPConnectionQueryResponseIRC(t *testing.T) {
 	defer cancel()
 
 	module := Module{
-		Timeout: time.Second,
 		TCP: TCPProbe{
 			QueryResponse: []QueryResponse{
 				{Send: "NICK prober"},
@@ -137,9 +136,8 @@ func TestTCPConnectionQueryResponseMatching(t *testing.T) {
 
 	testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
-
+	time.Sleep(time.Millisecond * 100)
 	module := Module{
-		Timeout: time.Second,
 		TCP: TCPProbe{
 			QueryResponse: []QueryResponse{
 				{
@@ -206,7 +204,6 @@ func TestTCPConnectionProtocol(t *testing.T) {
 
 	// Force IPv4
 	module := Module{
-		Timeout: time.Second,
 		TCP: TCPProbe{
 			PreferredIPProtocol: "ip4",
 		},
@@ -228,8 +225,7 @@ func TestTCPConnectionProtocol(t *testing.T) {
 
 	// Force IPv6
 	module = Module{
-		Timeout: time.Second,
-		TCP:     TCPProbe{},
+		TCP: TCPProbe{},
 	}
 
 	registry = prometheus.NewRegistry()
@@ -248,7 +244,6 @@ func TestTCPConnectionProtocol(t *testing.T) {
 
 	// Prefer IPv4
 	module = Module{
-		Timeout: time.Second,
 		TCP: TCPProbe{
 			PreferredIPProtocol: "ip4",
 		},
@@ -270,7 +265,6 @@ func TestTCPConnectionProtocol(t *testing.T) {
 
 	// Prefer IPv6
 	module = Module{
-		Timeout: time.Second,
 		TCP: TCPProbe{
 			PreferredIPProtocol: "ip6",
 		},
@@ -292,8 +286,7 @@ func TestTCPConnectionProtocol(t *testing.T) {
 
 	// Prefer nothing
 	module = Module{
-		Timeout: time.Second,
-		TCP:     TCPProbe{},
+		TCP: TCPProbe{},
 	}
 
 	registry = prometheus.NewRegistry()
@@ -312,8 +305,7 @@ func TestTCPConnectionProtocol(t *testing.T) {
 
 	// No protocol
 	module = Module{
-		Timeout: time.Second,
-		TCP:     TCPProbe{},
+		TCP: TCPProbe{},
 	}
 
 	registry = prometheus.NewRegistry()
-- 
2.25.1