Skip to content

Commit

Permalink
websocket support
Browse files Browse the repository at this point in the history
Signed-off-by: ad <[email protected]>
  • Loading branch information
adnull committed Aug 11, 2024
1 parent 3dd5dfe commit 9d8f9ec
Show file tree
Hide file tree
Showing 9 changed files with 490 additions and 13 deletions.
97 changes: 97 additions & 0 deletions CONFIGURATION.md
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ modules:
[ dns: <dns_probe> ]
[ icmp: <icmp_probe> ]
[ grpc: <grpc_probe> ]
[ websocket: <websocket_probe> ]

```

Expand Down Expand Up @@ -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 ]
Expand Down
2 changes: 2 additions & 0 deletions blackbox.yml
Original file line number Diff line number Diff line change
Expand Up @@ -49,3 +49,5 @@ modules:
timeout: 5s
icmp:
ttl: 5
websocket:
prober: websocket
37 changes: 30 additions & 7 deletions config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
package config

import (
"encoding/base64"
"errors"
"fmt"
"math"
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down
15 changes: 15 additions & 0 deletions example.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
3 changes: 2 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
11 changes: 6 additions & 5 deletions prober/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
)

Expand Down
138 changes: 138 additions & 0 deletions prober/websocket.go
Original file line number Diff line number Diff line change
@@ -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
}
Loading

0 comments on commit 9d8f9ec

Please sign in to comment.