Merge pull request #28 from servak/add-http-header

Enhanced HTTP Prober with Additional Capabilities
servak · Jun 18, 2023 · 061a453 · 061a453
2 parents 60be025 + 2cce176
commit 061a453
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 5 deletions.
diff --git a/internal/command/mping.go b/internal/command/mping.go
@@ -204,10 +204,10 @@ func newProber(cfg *prober.ProberConfig, manager *stats.MetricsManager, targets
 			manager.Register(ip.String(), name)
 		}
 		probe, err = prober.NewICMPProber(cfg.Probe, uniqueStringer(addrs), cfg.ICMP)
-	case prober.HTTP:
+	case prober.HTTP, prober.HTTPS:
 		var ts []string
 		for _, h := range targets {
-			t := "http:" + h
+			t := fmt.Sprintf("%s:%s", cfg.Probe, h)
 			ts = append(ts, t)
 			manager.Register(t, t)
 		}

diff --git a/internal/config/config.go b/internal/config/config.go
@@ -39,6 +39,13 @@ func DefaultConfig() *Config {
 					ExpectBody: "",
 				},
 			},
+			string(prober.HTTPS): {
+				Probe: prober.HTTPS,
+				HTTP: &prober.HTTPConfig{
+					ExpectCode: 200,
+					ExpectBody: "",
+				},
+			},
 		},
 		UI: &ui.UIConfig{
 			CUI: &ui.CUIConfig{

diff --git a/internal/prober/http.go b/internal/prober/http.go
@@ -1,6 +1,7 @@
 package prober
 
 import (
+	"crypto/tls"
 	"errors"
 	"fmt"
 	"io"
@@ -26,14 +27,37 @@ type (
 	}
 
 	HTTPConfig struct {
-		ExpectCode int    `yaml:"expect_code"`
-		ExpectBody string `yaml:"expect_body"`
+		Header              http.Header `yaml:"headers"`
+		ExpectCode          int         `yaml:"expect_code"`
+		ExpectBody          string      `yaml:"expect_body"`
+		SkipSSLVerification bool        `yaml:"skip_ssl_verification"`
+		RedirectOFF         bool        `yaml:"redirect_off"`
+	}
+
+	customTransport struct {
+		transport http.RoundTripper
+		headers   http.Header
 	}
 )
 
 func NewHTTPProber(targets []string, cfg *HTTPConfig) *HTTPProber {
+	var rd func(req *http.Request, via []*http.Request) error
+	if cfg.RedirectOFF {
+		rd = func(req *http.Request, via []*http.Request) error {
+			return http.ErrUseLastResponse
+		}
+	}
+	client := &http.Client{
+		Transport: &customTransport{
+			transport: &http.Transport{
+				TLSClientConfig: &tls.Config{InsecureSkipVerify: cfg.SkipSSLVerification},
+			},
+			headers: cfg.Header,
+		},
+		CheckRedirect: rd,
+	}
 	return &HTTPProber{
-		client:   &http.Client{},
+		client:   client,
 		targets:  targets,
 		config:   cfg,
 		exitChan: make(chan bool),
@@ -126,3 +150,10 @@ func (p *HTTPProber) Start(r chan *Event, interval, timeout time.Duration) error
 func (p *HTTPProber) Stop() {
 	p.exitChan <- true
 }
+
+func (c *customTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	for k, v := range c.headers {
+		req.Header[k] = v
+	}
+	return c.transport.RoundTrip(req)
+}