ITWorx-Pulse-Public/internal/probe/executor.go

package probe

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"net"
	"net/http"
	"net/netip"
	"net/url"
	"reflect"
	"strconv"
	"strings"
	"time"
)

type HTTPDoer interface {
	Do(*http.Request) (*http.Response, error)
}

type ProbeExecutor struct {
	Policy      NetworkPolicy
	Resolver    Resolver
	HTTP        HTTPDoer
	DialContext func(context.Context, string, string) (net.Conn, error)
	Now         func() time.Time
}

func (e ProbeExecutor) Execute(ctx context.Context, definition Definition) (Result, error) {
	if ctx == nil {
		return Result{ProbeID: definition.ID, State: "unknown", ErrorClass: "invalid_context"}, errors.New("probe context is nil")
	}
	if err := ctx.Err(); err != nil {
		return Result{ProbeID: definition.ID, State: "unknown", ErrorClass: "canceled"}, err
	}
	if err := definition.Validate(); err != nil {
		return Result{ProbeID: definition.ID, State: "unknown", ErrorClass: "invalid_probe"}, err
	}
	if e.Resolver == nil {
		e.Resolver = NetResolver{}
	}
	if e.Now == nil {
		e.Now = func() time.Time { return time.Now().UTC() }
	}
	probeContext, cancel := context.WithTimeout(ctx, definition.Timeout)
	defer cancel()
	switch definition.Type {
	case TypeHTTP:
		return e.executeHTTP(probeContext, definition)
	case TypeTCP:
		return e.executeTCP(probeContext, definition)
	case TypeDNS:
		return e.executeDNS(probeContext, definition)
	case TypeTLS:
		return e.executeTLS(probeContext, definition)
	case TypeICMP:
		return Result{ProbeID: definition.ID, ObservedAt: e.Now().UTC(), CompletedAt: e.Now().UTC(), State: "unknown", ErrorClass: "unsupported"}, nil
	default:
		return Result{ProbeID: definition.ID, ObservedAt: e.Now().UTC(), CompletedAt: e.Now().UTC(), State: "unknown", ErrorClass: "unsupported"}, nil
	}
}

func (e ProbeExecutor) executeHTTP(ctx context.Context, definition Definition) (Result, error) {
	started := e.Now().UTC()
	scheme := strings.ToLower(strings.TrimSpace(definition.Target.Scheme))
	if scheme == "" {
		scheme = "http"
	}
	port := defaultPort(scheme, 80)
	if definition.Target.Port > 0 {
		port = definition.Target.Port
	}
	if _, err := e.Policy.ValidateTarget(ctx, e.Resolver, scheme, definition.Target.Host, port); err != nil {
		return transportFailure(definition.ID, err), err
	}
	target := &url.URL{Scheme: scheme, Host: net.JoinHostPort(definition.Target.Host, strconv.Itoa(port)), Path: definition.Target.Path}
	if target.Path == "" {
		target.Path = "/"
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, target.String(), nil)
	if err != nil {
		return transportFailure(definition.ID, err), err
	}
	if e.HTTP == nil {
		client, clientErr := NewSafeClient(e.Policy, e.Resolver)
		if clientErr != nil {
			return transportFailure(definition.ID, clientErr), clientErr
		}
		if !definition.FollowRedirects {
			client.HTTP.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }
		}
		e.HTTP = client
	}
	response, err := e.HTTP.Do(request)
	if err != nil {
		return transportFailure(definition.ID, err), err
	}
	if response == nil {
		err = errors.New("probe response is invalid")
		return transportFailure(definition.ID, err), err
	}
	body, bodyErr := ReadLimitedBody(response, e.Policy.withDefaults().MaxResponseBytes)
	result := Result{ProbeID: definition.ID, ObservedAt: started, CompletedAt: e.Now().UTC(), ResponseTimeMS: intPointer(int(time.Since(started).Milliseconds())), StatusCode: intPointer(response.StatusCode)}
	if bodyErr != nil {
		result.State = "unknown"
		result.ErrorClass = "response_too_large"
		return result, nil
	}
	if !expectedStatus(definition.ExpectedStatusCodes, response.StatusCode) {
		result.State = "down"
		result.ErrorClass = "status_not_expected"
		return result, nil
	}
	if assertionErr := assertBody(body, definition.ContentAssertion); assertionErr != nil {
		result.State = "degraded"
		result.ErrorClass = "content_assertion_failed"
		return result, nil
	}
	result.State = "up"
	return result, nil
}

func (e ProbeExecutor) executeTCP(ctx context.Context, definition Definition) (Result, error) {
	started := e.Now().UTC()
	port := definition.Target.Port
	if port == 0 {
		return transportFailure(definition.ID, errors.New("tcp port is required")), errors.New("tcp port is required")
	}
	addresses, err := e.Policy.ValidateTarget(ctx, e.Resolver, TypeTCP, definition.Target.Host, port)
	if err != nil {
		return transportFailure(definition.ID, err), err
	}
	connection, err := e.dial(ctx, "tcp", addresses[0], port)
	if err != nil {
		return transportFailure(definition.ID, err), err
	}
	if connection == nil {
		err = errors.New("probe dial returned nil connection")
		return transportFailure(definition.ID, err), err
	}
	_ = connection.Close()
	return Result{ProbeID: definition.ID, ObservedAt: started, CompletedAt: e.Now().UTC(), State: "up", ResponseTimeMS: intPointer(int(time.Since(started).Milliseconds()))}, nil
}

func (e ProbeExecutor) executeDNS(ctx context.Context, definition Definition) (Result, error) {
	started := e.Now().UTC()
	addresses, err := e.Policy.ValidateTarget(ctx, e.Resolver, TypeDNS, definition.Target.Host, defaultPort("dns", 53))
	if err != nil {
		return transportFailure(definition.ID, err), err
	}
	return Result{ProbeID: definition.ID, ObservedAt: started, CompletedAt: e.Now().UTC(), State: "up", ResponseTimeMS: intPointer(int(time.Since(started).Milliseconds())), Attributes: map[string]any{"addressCount": len(addresses)}}, nil
}

func (e ProbeExecutor) executeTLS(ctx context.Context, definition Definition) (Result, error) {
	started := e.Now().UTC()
	port := defaultPort(definition.Target.Scheme, 443)
	if definition.Target.Port > 0 {
		port = definition.Target.Port
	}
	addresses, err := e.Policy.ValidateTarget(ctx, e.Resolver, TypeTLS, definition.Target.Host, port)
	if err != nil {
		return transportFailure(definition.ID, err), err
	}
	connection, err := e.dial(ctx, "tcp", addresses[0], port)
	if err != nil {
		return transportFailure(definition.ID, err), err
	}
	if connection == nil {
		err = errors.New("probe dial returned nil connection")
		return transportFailure(definition.ID, err), err
	}
	defer connection.Close()
	tlsConnection := tls.Client(connection, &tls.Config{ServerName: definition.Target.Host, MinVersion: tls.VersionTLS12, InsecureSkipVerify: !definition.VerifyTLS})
	if err := tlsConnection.HandshakeContext(ctx); err != nil {
		return transportFailure(definition.ID, err), err
	}
	state := tlsConnection.ConnectionState()
	result := Result{ProbeID: definition.ID, ObservedAt: started, CompletedAt: e.Now().UTC(), State: "up", ResponseTimeMS: intPointer(int(time.Since(started).Milliseconds()))}
	if len(state.PeerCertificates) > 0 {
		result.Certificate = certificateFromX509(state.PeerCertificates[0], definition, e.Now)
		if result.Certificate.VerificationState == "attention" || result.Certificate.VerificationState == "invalid" {
			result.State = "degraded"
		}
	}
	return result, nil
}

func (e ProbeExecutor) dial(ctx context.Context, network string, address netip.Addr, port int) (net.Conn, error) {
	if e.DialContext != nil {
		return e.DialContext(ctx, network, net.JoinHostPort(address.String(), strconv.Itoa(port)))
	}
	dialer := &net.Dialer{Timeout: e.Policy.withDefaults().Timeout}
	return dialer.DialContext(ctx, network, net.JoinHostPort(address.String(), strconv.Itoa(port)))
}

func certificateFromX509(certificate *x509.Certificate, definition Definition, now func() time.Time) *Certificate {
	if certificate == nil {
		return nil
	}
	observed := now().UTC()
	hostnameValid := certificate.VerifyHostname(definition.Target.Host) == nil
	verification := "valid"
	if !hostnameValid {
		verification = "invalid"
	} else if !certificate.NotAfter.After(observed) {
		verification = "invalid"
	} else if certificate.NotAfter.Before(observed.Add(30 * 24 * time.Hour)) {
		verification = "attention"
	}
	return &Certificate{ID: definition.ID + "-certificate", ServiceID: definition.ServiceID, EndpointID: definition.EndpointID, ObservedAt: observed, ExpiresAt: &certificate.NotAfter, Issuer: certificate.Issuer.String(), Subject: certificate.Subject.String(), HostnameValid: &hostnameValid, VerificationState: verification}
}

func assertBody(body []byte, assertion map[string]any) error {
	if len(assertion) == 0 {
		return nil
	}
	if keyword, ok := assertion["keyword"].(string); ok && !strings.Contains(string(body), keyword) {
		return errors.New("keyword assertion failed")
	}
	if expected, ok := assertion["json"].(map[string]any); ok {
		var actual map[string]any
		if err := json.Unmarshal(body, &actual); err != nil {
			return errors.New("json assertion body is invalid")
		}
		for key, value := range expected {
			if !reflect.DeepEqual(actual[key], value) {
				return errors.New("json assertion failed")
			}
		}
	}
	return nil
}

func expectedStatus(expected []int, status int) bool {
	if len(expected) == 0 {
		return status >= 200 && status < 400
	}
	for _, candidate := range expected {
		if candidate == status {
			return true
		}
	}
	return false
}

func defaultPort(scheme string, fallback int) int {
	if scheme == "http" {
		return 80
	}
	if scheme == "https" || scheme == "tls" {
		return 443
	}
	return fallback
}

func transportFailure(id string, err error) Result {
	result := Result{ProbeID: id, State: "unknown", ErrorClass: "transport_error"}
	if err != nil {
		result.ErrorMessage = boundText(err.Error(), 256)
	}
	return result
}

func intPointer(value int) *int { return &value }