websocket support

Signed-off-by: ad <[email protected]>

Author: adnull

CONFIGURATION.md

@@ -41,6 +41,7 @@ modules:
   [ dns: <dns_probe> ]
   [ icmp: <icmp_probe> ]
   [ grpc: <grpc_probe> ]
+  [ websocket: <websocket_probe> ]
 
 ```
 
@@ -318,9 +319,105 @@ tls_config:
   [ <tls_config> ]
 ```
 
+### `<websocket_probe>`
+
+```yml
+# Optional HTTP request configuration
+http_config: 
+  
+  # The HTTP basic authentification credentials
+  basic_auth:
+    [ username: <string> ]
+    [ password: <string >]
+  
+  # Sets the `Authorization: Bearer <token>` header on every request with
+  # the configured token.
+  [ bearer_token: <string>
+
+  # Sets HTTP headers for the request
+  headers:
+    [ - [ header_name: <string> ], ... ]
+  
+
+  # Whether to skip certificate verification on connect
+  [insecure_skip_verify: <boolean> | default = true ]
+
+# The query sent after connection upgrade and the expected associated response.
+query_response:
+  [ - [ [ expect: <string> ],
+        [ send: <string> ],
+        [ starttls: <boolean | default = false> ]
+      ], ...
+  ]
+
+```
+
+### `<tls_config>`
+
+```yml
+
+# Disable target certificate validation.
+[ insecure_skip_verify: <boolean> | default = false ]
+
+# The CA cert to use for the targets.
+[ ca_file: <filename> ]
+
+# The client cert file for the targets.
+[ cert_file: <filename> ]
+
+# The client key file for the targets.
+[ key_file: <filename> ]
+
+# Used to verify the hostname for the targets.
+[ server_name: <string> ]
+
+# Minimum acceptable TLS version. Accepted values: TLS10 (TLS 1.0), TLS11 (TLS
+# 1.1), TLS12 (TLS 1.2), TLS13 (TLS 1.3).
+# If unset, Prometheus will use Go default minimum version, which is TLS 1.2.
+# See MinVersion in https://pkg.go.dev/crypto/tls#Config.
+[ min_version: <string> ]
+
+# Maximum acceptable TLS version. Accepted values: TLS10 (TLS 1.0), TLS11 (TLS
+# 1.1), TLS12 (TLS 1.2), TLS13 (TLS 1.3).
+# Can be used to test for the presence of insecure TLS versions.
+# If unset, Prometheus will use Go default maximum version, which is TLS 1.3.
+# See MaxVersion in https://pkg.go.dev/crypto/tls#Config.
+[ max_version: <string> ]
+```
+
+#### `<oauth2>`
+
+OAuth 2.0 authentication using the client credentials grant type. Blackbox
+exporter fetches an access token from the specified endpoint with the given
+client access and secret keys.
+
+NOTE: This is *experimental* in the blackbox exporter and might not be
+reflected properly in the probe metrics at the moment.
+
+```yml
+client_id: <string>
+[ client_secret: <secret> ]
+
+# Read the client secret from a file.
+# It is mutually exclusive with `client_secret`.
+[ client_secret_file: <filename> ]
+
+# Scopes for the token request.
+scopes:
+  [ - <string> ... ]
+
+# The URL to fetch the token from.
+token_url: <string>
+
+# Optional parameters to append to the token URL.
+endpoint_params:
+  [ <string>: <string> ... ]
+```
+
 ### `<tls_config>`
 
 ```yml
+[ http_config: <websocket_http_config>]
 
 # Disable target certificate validation.
 [ insecure_skip_verify: <boolean> | default = false ]

blackbox.yml

@@ -49,3 +49,5 @@ modules:
     timeout: 5s
     icmp:
       ttl: 5
+  websocket:
+    prober: websocket

config/config.go

@@ -14,6 +14,7 @@
 package config
 
 import (
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"math"
@@ -193,13 +194,14 @@ func MustNewRegexp(s string) Regexp {
 }
 
 type Module struct {
-	Prober  string        `yaml:"prober,omitempty"`
-	Timeout time.Duration `yaml:"timeout,omitempty"`
-	HTTP    HTTPProbe     `yaml:"http,omitempty"`
-	TCP     TCPProbe      `yaml:"tcp,omitempty"`
-	ICMP    ICMPProbe     `yaml:"icmp,omitempty"`
-	DNS     DNSProbe      `yaml:"dns,omitempty"`
-	GRPC    GRPCProbe     `yaml:"grpc,omitempty"`
+	Prober    string         `yaml:"prober,omitempty"`
+	Timeout   time.Duration  `yaml:"timeout,omitempty"`
+	HTTP      HTTPProbe      `yaml:"http,omitempty"`
+	TCP       TCPProbe       `yaml:"tcp,omitempty"`
+	ICMP      ICMPProbe      `yaml:"icmp,omitempty"`
+	DNS       DNSProbe       `yaml:"dns,omitempty"`
+	GRPC      GRPCProbe      `yaml:"grpc,omitempty"`
+	Websocket WebsocketProbe `yaml:"websocket,omitempty"`
 }
 
 type HTTPProbe struct {
@@ -287,6 +289,27 @@ type DNSRRValidator struct {
 	FailIfNoneMatchesRegexp []string `yaml:"fail_if_none_matches_regexp,omitempty"`
 }
 
+type WebsocketProbe struct {
+	HTTPClientConfig HTTPClientConfig `yaml:"http_config,omitempty"`
+	QueryResponse    []QueryResponse  `yaml:"query_response,omitempty"`
+}
+
+type HTTPClientConfig struct {
+	HTTPHeaders        map[string]interface{} `yaml:"headers,omitempty"`
+	BasicAuth          HTTPBasicAuth          `yaml:"basic_auth,omitempty"`
+	BearerToken        string                 `yaml:"bearer_token,omitempty"`
+	InsecureSkipVerify bool                   `yaml:"insecure_skip_verify,omitempty"`
+}
+
+type HTTPBasicAuth struct {
+	Username string `yaml:"username"`
+	Password string `yaml:"password"`
+}
+
+func (c *HTTPBasicAuth) BasicAuthHeader() string {
+	return "Basic " + base64.StdEncoding.EncodeToString([]byte(c.Username+":"+c.Password))
+}
+
 // UnmarshalYAML implements the yaml.Unmarshaler interface.
 func (s *Config) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	type plain Config

example.yml

@@ -181,3 +181,18 @@ modules:
       transport_protocol: "tcp" # defaults to "udp"
       preferred_ip_protocol: "ip4" # defaults to "ip6"
       query_name: "www.prometheus.io"
+  websocket_example:
+    prober: websocket
+    websocket:
+      http_config:
+        basic_auth:
+          username: "user"
+          password: "password"
+        bearer_token: "secret_token" 
+        headers:
+          X-Some-Header: "my_header"
+        insecure_skip_verify: true
+      query_response:
+        - expect: ^Hello,\s(.+)"
+        - send: "Hello server, i'am ${1}"
+        - expect: ^Welcome

go.mod

@@ -7,12 +7,14 @@ require (
 	github.com/alecthomas/units v0.0.0-20231202071711-9a357b53e9c9
 	github.com/andybalholm/brotli v1.1.0
 	github.com/go-kit/log v0.2.1
+	github.com/gorilla/websocket v1.5.3
 	github.com/miekg/dns v1.1.61
 	github.com/prometheus/client_golang v1.19.1
 	github.com/prometheus/client_model v0.6.1
 	github.com/prometheus/common v0.55.0
 	github.com/prometheus/exporter-toolkit v0.11.0
 	golang.org/x/net v0.27.0
+	golang.org/x/text v0.16.0
 	google.golang.org/grpc v1.65.0
 	gopkg.in/yaml.v2 v2.4.0
 	gopkg.in/yaml.v3 v3.0.1
@@ -33,7 +35,6 @@ require (
 	golang.org/x/oauth2 v0.21.0 // indirect
 	golang.org/x/sync v0.7.0 // indirect
 	golang.org/x/sys v0.22.0 // indirect
-	golang.org/x/text v0.16.0 // indirect
 	golang.org/x/tools v0.22.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240528184218-531527333157 // indirect
 	google.golang.org/protobuf v1.34.2 // indirect

go.sum

@@ -20,6 +20,8 @@ github.com/go-logfmt/logfmt v0.5.1/go.mod h1:WYhtIu8zTZfxdn5+rREduYbwxfcBr/Vr6KE
 github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/jpillora/backoff v1.0.0 h1:uvFg412JmmHBHw7iwprIxkPMI+sGQ4kzOWsMeHnm2EA=
 github.com/jpillora/backoff v1.0.0/go.mod h1:J/6gKK9jxlEcS3zixgDgUAsiuZ7yrSoa/FX5e0EB2j4=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=

prober/handler.go

@@ -34,11 +34,12 @@ import (
 
 var (
 	Probers = map[string]ProbeFn{
-		"http": ProbeHTTP,
-		"tcp":  ProbeTCP,
-		"icmp": ProbeICMP,
-		"dns":  ProbeDNS,
-		"grpc": ProbeGRPC,
+		"http":      ProbeHTTP,
+		"tcp":       ProbeTCP,
+		"icmp":      ProbeICMP,
+		"dns":       ProbeDNS,
+		"grpc":      ProbeGRPC,
+		"websocket": ProbeWebsocket,
 	}
 )
 

prober/websocket.go

@@ -0,0 +1,138 @@
+// Copyright 2016 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package prober
+
+import (
+	"context"
+	"crypto/tls"
+	"net/http"
+	"net/url"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+	"github.com/gorilla/websocket"
+	"github.com/prometheus/blackbox_exporter/config"
+	"github.com/prometheus/client_golang/prometheus"
+	"golang.org/x/text/cases"
+	"golang.org/x/text/language"
+)
+
+func ProbeWebsocket(ctx context.Context, target string, module config.Module, registry *prometheus.Registry, logger log.Logger) (success bool) {
+
+	targetURL, err := url.Parse(target)
+	if err != nil {
+		logger.Log("msg", "Could not parse target URL", "err", err)
+		return false
+	}
+
+	level.Debug(logger).Log("msg", "probing websocket", "target", targetURL.String())
+
+	httpStatusCode := prometheus.NewGauge(prometheus.GaugeOpts{
+		Name: "probe_http_status_code",
+		Help: "Response HTTP status code",
+	})
+	isConnected := prometheus.NewGauge(prometheus.GaugeOpts{
+		Name: "probe_is_upgraded",
+		Help: "Indicates if the websocket connection was successfully upgraded",
+	})
+
+	registry.MustRegister(isConnected)
+	registry.MustRegister(httpStatusCode)
+
+	dialer := websocket.Dialer{
+		TLSClientConfig: &tls.Config{
+			InsecureSkipVerify: module.Websocket.HTTPClientConfig.InsecureSkipVerify,
+		},
+	}
+
+	connection, resp, err := dialer.DialContext(ctx, targetURL.String(), constructHeadersFromConfig(module.Websocket.HTTPClientConfig, logger))
+	if resp != nil {
+		httpStatusCode.Set(float64(resp.StatusCode))
+	}
+	if err != nil {
+		logger.Log("msg", "Error dialing websocket", "err", err)
+		return false
+	}
+	defer connection.Close()
+
+	isConnected.Set(1)
+
+	if len(module.Websocket.QueryResponse) > 0 {
+		probeFailedDueToRegex := prometheus.NewGauge(prometheus.GaugeOpts{
+			Name: "probe_failed_due_to_regex",
+			Help: "Indicates if probe failed due to regex",
+		})
+		registry.MustRegister(probeFailedDueToRegex)
+
+		queryMatched := true
+		for _, qr := range module.Websocket.QueryResponse {
+			send := qr.Send
+
+			if qr.Expect.Regexp != nil {
+				var match []int
+				_, message, err := connection.ReadMessage()
+				if err != nil {
+					logger.Log("msg", "Error reading message", "err", err)
+					queryMatched = false
+					break
+				}
+				match = qr.Expect.Regexp.FindSubmatchIndex(message)
+				if match != nil {
+					level.Debug(logger).Log("msg", "regexp matched", "regexp", qr.Expect.Regexp, "line", message)
+				} else {
+					level.Error(logger).Log("msg", "Regexp did not match", "regexp", qr.Expect.Regexp, "line", message)
+					queryMatched = false
+					break
+				}
+				send = string(qr.Expect.Regexp.Expand(nil, []byte(send), message, match))
+			}
+
+			if send != "" {
+				err = connection.WriteMessage(websocket.TextMessage, []byte(send))
+				if err != nil {
+					queryMatched = false
+					logger.Log("msg", "Error sending message", "err", err)
+					break
+				}
+				level.Debug(logger).Log("msg", "message sent", "message", send)
+			}
+		}
+		if queryMatched {
+			probeFailedDueToRegex.Set(0)
+		} else {
+			probeFailedDueToRegex.Set(1)
+		}
+	}
+
+	return true
+}
+
+func constructHeadersFromConfig(config config.HTTPClientConfig, logger log.Logger) map[string][]string {
+	headers := http.Header{}
+	if config.BasicAuth.Username != "" || config.BasicAuth.Password != "" {
+		headers.Add("Authorization", config.BasicAuth.BasicAuthHeader())
+	} else if config.BearerToken != "" {
+		headers.Add("Authorization", "Bearer "+config.BearerToken)
+	}
+	for key, value := range config.HTTPHeaders {
+		if _, ok := value.(string); ok {
+			headers.Add(key, value.(string))
+		} else if _, ok := value.([]string); ok {
+			headers[cases.Title(language.English).String(key)] = append(headers[key], value.([]string)...)
+		}
+	}
+
+	level.Debug(logger).Log("msg", "Constructed headers", "headers", headers)
+	return headers
+}