From 837542102c30448e86915abdc7f443f5351cbf3b Mon Sep 17 00:00:00 2001
From: Emmanuel Gomez <emmanuel@gomez.io>
Date: Fri, 14 Dec 2018 11:01:28 -0800
Subject: [PATCH] Add support for `body_path` in HTTP probe config.

Fixes prometheus/blackbox_exporter#391.

Signed-off-by: Emmanuel Gomez <emmanuel@gomez.io>
---
 config/config.go    | 35 +++++++++++++++++++++++------------
 prober/http.go      |  6 +++++-
 prober/http_test.go | 27 +++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/config/config.go b/config/config.go
index 2bd5d946..d65864c3 100644
--- a/config/config.go
+++ b/config/config.go
@@ -3,7 +3,9 @@ package config
 import (
 	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
+	"os"
 	"runtime"
 	"sync"
 	"time"
@@ -52,18 +54,20 @@ type Module struct {
 
 type HTTPProbe struct {
 	// Defaults to 2xx.
-	ValidStatusCodes       []int                   `yaml:"valid_status_codes,omitempty"`
-	ValidHTTPVersions      []string                `yaml:"valid_http_versions,omitempty"`
-	IPProtocol             string                  `yaml:"preferred_ip_protocol,omitempty"`
-	IPProtocolFallback     bool                    `yaml:"ip_protocol_fallback,omitempty"`
-	NoFollowRedirects      bool                    `yaml:"no_follow_redirects,omitempty"`
-	FailIfSSL              bool                    `yaml:"fail_if_ssl,omitempty"`
-	FailIfNotSSL           bool                    `yaml:"fail_if_not_ssl,omitempty"`
-	Method                 string                  `yaml:"method,omitempty"`
-	Headers                map[string]string       `yaml:"headers,omitempty"`
-	FailIfMatchesRegexp    []string                `yaml:"fail_if_matches_regexp,omitempty"`
-	FailIfNotMatchesRegexp []string                `yaml:"fail_if_not_matches_regexp,omitempty"`
-	Body                   string                  `yaml:"body,omitempty"`
+	ValidStatusCodes       []int             `yaml:"valid_status_codes,omitempty"`
+	ValidHTTPVersions      []string          `yaml:"valid_http_versions,omitempty"`
+	IPProtocol             string            `yaml:"preferred_ip_protocol,omitempty"`
+	IPProtocolFallback     bool              `yaml:"ip_protocol_fallback,omitempty"`
+	NoFollowRedirects      bool              `yaml:"no_follow_redirects,omitempty"`
+	FailIfSSL              bool              `yaml:"fail_if_ssl,omitempty"`
+	FailIfNotSSL           bool              `yaml:"fail_if_not_ssl,omitempty"`
+	Method                 string            `yaml:"method,omitempty"`
+	Headers                map[string]string `yaml:"headers,omitempty"`
+	FailIfMatchesRegexp    []string          `yaml:"fail_if_matches_regexp,omitempty"`
+	FailIfNotMatchesRegexp []string          `yaml:"fail_if_not_matches_regexp,omitempty"`
+	Body                   string            `yaml:"body,omitempty"`
+	BodyPath               string            `yaml:"body_path,omitempty"`
+	BodyFile               io.Reader
 	HTTPClientConfig       config.HTTPClientConfig `yaml:"http_client_config,inline"`
 }
 
@@ -135,6 +139,13 @@ func (s *HTTPProbe) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	if err := s.HTTPClientConfig.Validate(); err != nil {
 		return err
 	}
+	if s.BodyPath != "" {
+		file, err := os.Open(s.BodyPath)
+		if err != nil {
+			return err
+		}
+		s.BodyFile = file
+	}
 	return nil
 }
 
diff --git a/prober/http.go b/prober/http.go
index 542f3129..3054ad86 100644
--- a/prober/http.go
+++ b/prober/http.go
@@ -257,10 +257,14 @@ func ProbeHTTP(ctx context.Context, target string, module config.Module, registr
 
 	var body io.Reader
 
-	// If a body is configured, add it to the request.
+	// If a body string is configured, add it to the request.
 	if httpConfig.Body != "" {
 		body = strings.NewReader(httpConfig.Body)
 	}
+	// If a body path is configured, add it to the request.
+	if httpConfig.BodyFile != nil {
+		body = httpConfig.BodyFile
+	}
 
 	request, err := http.NewRequest(httpConfig.Method, targetURL.String(), body)
 	request.Host = origHost
diff --git a/prober/http_test.go b/prober/http_test.go
index a7deaa89..080ea0fa 100644
--- a/prober/http_test.go
+++ b/prober/http_test.go
@@ -164,6 +164,33 @@ func TestPost(t *testing.T) {
 	}
 }
 
+func TestPostBody(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != "POST" {
+			w.WriteHeader(http.StatusBadRequest)
+		}
+	}))
+	defer ts.Close()
+
+	recorder := httptest.NewRecorder()
+	registry := prometheus.NewRegistry()
+	testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	result := ProbeHTTP(testCTX, ts.URL,
+		config.Module{
+			Timeout: time.Second,
+			HTTP: config.HTTPProbe{
+				IPProtocolFallback: true,
+				Method:             "POST",
+				BodyPath:           "./http.go",
+			},
+		}, registry, log.NewNopLogger())
+	body := recorder.Body.String()
+	if !result {
+		t.Fatalf("Post test failed unexpectedly, got %s", body)
+	}
+}
+
 func TestBasicAuth(t *testing.T) {
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 	}))