smallwebwaf/internal/sww/healthcheck.go

package sww

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"time"

	"github.com/spf13/viper"

	"sneak.berlin/go/smallwebwaf/internal/config"
	"sneak.berlin/go/smallwebwaf/internal/discriminator"
)

// healthCheckTimeout bounds the whole health check.
const healthCheckTimeout = 5 * time.Second

var errHealthEndpoint = errors.New("the health endpoint answered")

// HealthCheck is the container's health check. In discriminator's container, where
// SWW_COMMAND is discriminator, it returns 0 while discriminator answers its health
// endpoint on 127.0.0.1, at the port in SWW_WAF_LISTEN_ADDR. In rproxy's container,
// where SWW_COMMAND is rproxy, unset or empty, it returns 0 while rproxy
// answers its health endpoint on 127.0.0.1, at the port in SWW_RPROXY_LISTEN_ADDR, and
// the app accepts connections at the address in SWW_RPROXY_UPSTREAM_URL. Otherwise,
// any other value of SWW_COMMAND included, it writes why to stderr and
// returns 1. It reads no other setting, nor a file that another names, so neither
// can fail it.
// args are the arguments after `healthcheck`; it takes none, and given
// one it names it on stderr and returns 1 without checking anything.
func HealthCheck(
	ctx context.Context, args []string, v *viper.Viper,
	stderr io.Writer,
) int {
	if len(args) > 0 {
		_, _ = fmt.Fprintf(stderr,
			"sww healthcheck: unexpected argument %q\n", args[0])

		return 1
	}

	ctx, cancel := context.WithTimeout(ctx, healthCheckTimeout)
	defer cancel()

	command, err := config.Command(v)

	switch {
	case err != nil:
		err = fmt.Errorf("invalid setting: %w", err)
	case command == "discriminator":
		err = checkDiscriminator(ctx, v)
	default:
		err = checkRProxyAndApp(ctx, v)
	}

	if err != nil {
		_, _ = fmt.Fprintln(stderr, "unhealthy:", err)

		return 1
	}

	return 0
}

// checkDiscriminator asks discriminator's health endpoint.
func checkDiscriminator(ctx context.Context, v *viper.Viper) error {
	listenAddr, err := config.DiscriminatorListenAddr(v)
	if err != nil {
		return fmt.Errorf("invalid setting: %w", err)
	}

	err = askHealthEndpoint(ctx, listenAddr)
	if err != nil {
		return fmt.Errorf("ask discriminator: %w", err)
	}

	return nil
}

// checkRProxyAndApp asks rproxy's health endpoint, and connects to the app.
func checkRProxyAndApp(ctx context.Context, v *viper.Viper) error {
	listenAddr, upstreamURL, err := config.RProxyListenAddrAndUpstreamURL(v)
	if err != nil {
		return fmt.Errorf("invalid setting: %w", err)
	}

	err = askHealthEndpoint(ctx, listenAddr)
	if err != nil {
		return fmt.Errorf("ask rproxy: %w", err)
	}

	conn, err := (&net.Dialer{}).DialContext(ctx, "tcp", appAddress(upstreamURL))
	if err != nil {
		return fmt.Errorf("connect to the app: %w", err)
	}

	_ = conn.Close()

	return nil
}

// askHealthEndpoint asks the health endpoint on 127.0.0.1 at the port of
// listenAddr, which the settings have checked has one.
func askHealthEndpoint(ctx context.Context, listenAddr string) error {
	_, port, _ := net.SplitHostPort(listenAddr)
	health := "http://" + net.JoinHostPort("127.0.0.1", port) + discriminator.HealthPath

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, health, http.NoBody)
	if err != nil {
		return fmt.Errorf("make the request: %w", err)
	}

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("send the request: %w", err)
	}

	_ = res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("%w %s", errHealthEndpoint, res.Status)
	}

	return nil
}

// appAddress is the host and port of the app's URL, the port being the
// scheme's own when the URL names none.
func appAddress(app *url.URL) string {
	port := app.Port()
	if port == "" {
		port = "80"
		if app.Scheme == "https" {
			port = "443"
		}
	}

	return net.JoinHostPort(app.Hostname(), port)
}