smallwebwaf/internal/sww/sww.go

// Package sww runs the sww process, as the command it is given.
// `sww discriminator` reads the settings, the rule files, the lookup database
// and the state files, answers rproxy's questions until it is told to stop, and
// then stops in an orderly way, writing the state files. `sww rproxy`
// reads rproxy's settings alone, and runs the reverse proxy in front of one app
// that asks discriminator about each request. `sww healthcheck` is the
// container's health check.
package sww

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/spf13/viper"

	"sneak.berlin/go/smallwebwaf/internal/alerts"
	"sneak.berlin/go/smallwebwaf/internal/config"
	"sneak.berlin/go/smallwebwaf/internal/discriminator"
	"sneak.berlin/go/smallwebwaf/internal/lookup"
	"sneak.berlin/go/smallwebwaf/internal/remotelog"
	"sneak.berlin/go/smallwebwaf/internal/reputation"
	"sneak.berlin/go/smallwebwaf/internal/requestlog"
	"sneak.berlin/go/smallwebwaf/internal/rules"
	"sneak.berlin/go/smallwebwaf/internal/state"
	"sneak.berlin/go/smallwebwaf/internal/waf"
)

// shutdownTimeout is how long requests in progress may take to finish
// once sww is told to stop, before their connections are closed.
// runit and docker wait a little longer before they kill the process.
const shutdownTimeout = 5 * time.Second

// remoteLogStopTimeout is how long, as discriminator stops, the log lines still
// waiting are sent to SWW_WAF_LOG_REMOTE_URL before they are given up. stdout has
// carried them.
const remoteLogStopTimeout = 2 * time.Second

// usageStatus is the exit status with no command, or one sww does not
// know.
const usageStatus = 2

// usage is what sww prints then.
const usage = `usage: sww <command>

commands:
  rproxy          the reverse proxy in front of one app
  discriminator   decides for rproxy, and keeps bans, counters and history
  healthcheck     the container's health check
`

// Params are what RunDiscriminator and RunRProxy need from the process.
type Params struct {
	// Version is the version of the binary, set when it is built.
	Version string
	// Viper reads the settings, normally the viper config.NewViper returns
	// for the process's environment.
	Viper *viper.Viper
	// Stdout receives the request log and the process's own messages.
	Stdout io.Writer
}

// Main runs the command os.Args names, and returns the process's exit status.
// `sww rproxy` and `sww discriminator` run until SIGTERM or SIGINT.
// With no command, or one it does not know, it prints the commands and returns
// usageStatus.
func Main(version string) int {
	args := os.Args[1:]
	if len(args) == 0 {
		return printUsage(os.Stderr)
	}

	ctx, stop := signal.NotifyContext(context.Background(),
		syscall.SIGTERM, os.Interrupt)
	defer stop()

	params := Params{Version: version, Viper: config.NewViper(), Stdout: os.Stdout}

	switch args[0] {
	case "rproxy":
		return RunRProxy(ctx, params)
	case "discriminator":
		return RunDiscriminator(ctx, params)
	case "healthcheck":
		return HealthCheck(ctx, args[1:], params.Viper, os.Stderr)
	default:
		return printUsage(os.Stderr)
	}
}

// printUsage writes the commands to w, and returns usageStatus.
func printUsage(w io.Writer) int {
	_, _ = fmt.Fprint(w, usage)

	return usageStatus
}

// RunDiscriminator is `sww discriminator`: it reads the settings, the rule files,
// the lookup database and the state files, then serves requests until ctx is
// done. It returns the process's exit status, 1 when discriminator cannot start.
func RunDiscriminator(ctx context.Context, params Params) int {
	// Until the settings are read, the one message is an invalid setting's
	// error, which every SWW_WAF_LOG_LEVEL lets through.
	processLog := requestlog.NewProcessLogger(params.Stdout,
		config.InstanceName(params.Viper), slog.LevelError)

	cfg, err := config.FromEnvironment(params.Viper)
	if err != nil {
		processLog.Error("invalid setting", "error", err.Error())

		return 1
	}

	// While SWW_WAF_LOG_REMOTE_URL is set, every line on stdout from here on
	// is sent there too.
	stdout := params.Stdout

	var remote *remotelog.Sender

	if cfg.LogRemoteURL != nil {
		remote = newRemoteLogSender(cfg)
		stdout = io.MultiWriter(params.Stdout, remote)
	}

	processLog = requestlog.NewProcessLogger(stdout, cfg.InstanceName, cfg.LogLevel)

	if remote != nil {
		stopSending := startSending(ctx, remote, processLog)
		defer stopSending()
	}

	// The state files and the alerts give times in UTC.
	now := func() time.Time { return time.Now().UTC() }

	alertQueue := newAlertQueue(cfg, now, processLog)

	ruleFiles, err := rules.Load(rules.Params{
		Dir:        cfg.RulesDir,
		Enabled:    cfg.RulesEnabled,
		ProcessLog: processLog,
		Alerts:     alertQueue,
	})
	if err != nil {
		processLog.Error("cannot use the rule files", "error", err.Error())

		return 1
	}

	crs, err := discriminator.NewCoreRuleSet(cfg)
	if err != nil {
		processLog.Error("cannot use the Core Rule Set", "error", err.Error())

		return 1
	}

	server, err := newServer(cfg, stdout, processLog, now, ruleFiles, alertQueue, crs)
	if err != nil {
		processLog.Error("cannot use the lookup database", "error", err.Error())

		return 1
	}

	if remote != nil {
		server.Metrics.AddRemoteLog(remote)
	}

	files, err := loadStateFiles(cfg, server, alertQueue, now, processLog)
	if err != nil {
		processLog.Error("cannot use the state files", "error", err.Error())

		return 1
	}

	listener, err := (&net.ListenConfig{}).Listen(ctx, "tcp", cfg.ListenAddr)
	if err != nil {
		processLog.Error("cannot listen on SWW_WAF_LISTEN_ADDR", "error", err.Error())

		return 1
	}

	processLog.Info("starting",
		"version", params.Version,
		"address", listener.Addr().String(),
		"settings", cfg)

	return serve(ctx, server, listener, files, ruleFiles, alertQueue, processLog)
}

// newServer returns the server discriminator runs, with the metrics of the alerts,
// after reading the lookup database while SWW_WAF_LOOKUP_SOURCE is file. A lookup
// database that cannot be read is an error.
func newServer(
	cfg *config.Config, stdout io.Writer, processLog *slog.Logger,
	now func() time.Time, ruleFiles *rules.Files, alertQueue *alerts.Queue,
	crs *waf.CoreRuleSet,
) (*discriminator.Server, error) {
	var lookupFile *lookup.File

	if cfg.LookupSource == "file" {
		var err error

		lookupFile, err = lookup.OpenFile(lookup.FileParams{
			Path:       cfg.LookupDBPath,
			Now:        now,
			ProcessLog: processLog,
			Alerts:     alertQueue,
		})
		if err != nil {
			return nil, err
		}
	}

	server := discriminator.New(discriminator.Params{
		Config:       cfg,
		RequestLog:   stdout,
		ProcessLog:   processLog,
		GeoJSURL:     lookup.URL,
		AbuseIPDBURL: reputation.AbuseIPDBURL,
		LookupFile:   lookupFile,
		CoreRuleSet:  crs,
		Now:          now,
		Rules:        ruleFiles,
		Alerts:       alertQueue,
	})
	server.Metrics.AddAlerts(alertQueue)

	if lookupFile != nil {
		server.Metrics.AddLookupFile(lookupFile.LastRead, lookupFile.ReadFailures)
	}

	return server, nil
}

// loadStateFiles reads the state files into the parts of server and into
// alertQueue, as state.Load does.
func loadStateFiles(
	cfg *config.Config, server *discriminator.Server, alertQueue *alerts.Queue,
	now func() time.Time, processLog *slog.Logger,
) (*state.Files, error) {
	return state.Load(state.Params{
		Dir:             cfg.StateDir,
		WriteDelay:      cfg.StateWriteDelay,
		CounterInterval: cfg.StateCounterInterval,
		Ledger:          server.Ledger,
		Limiter:         server.Limiter,
		GeoJS:           server.GeoJS,
		Lists:           server.Lists,
		DNSBL:           server.DNSBL,
		AbuseIPDB:       server.AbuseIPDB,
		Alerts:          alertQueue,
		Anomalies:       server.Anomalies,
		Now:             now,
		ProcessLog:      processLog,
		Metrics:         server.Metrics,
	})
}

// newAlertQueue returns the queue of the alerts to the webhook, Slack and
// ntfy, with the settings for them.
func newAlertQueue(
	cfg *config.Config, now func() time.Time, processLog *slog.Logger,
) *alerts.Queue {
	return alerts.New(alerts.Params{
		WebhookURL:     cfg.AlertWebhookURL,
		WebhookHeaders: cfg.AlertWebhookHeaders,
		SlackURL:       cfg.AlertSlackWebhookURL,
		NtfyURL:        cfg.AlertNtfyURL,
		NtfyToken:      cfg.AlertNtfyToken,
		Events:         cfg.AlertEvents,
		Cooldown:       cfg.AlertCooldown,
		MaxPerHour:     cfg.AlertMaxPerHour,
		Instance:       cfg.InstanceName,
		Now:            now,
		ProcessLog:     processLog,
	})
}

// newRemoteLogSender returns a sender of the log lines to
// SWW_WAF_LOG_REMOTE_URL, with the settings for it.
func newRemoteLogSender(cfg *config.Config) *remotelog.Sender {
	return remotelog.New(remotelog.Params{
		URL:      cfg.LogRemoteURL,
		RootCAs:  cfg.LogRemoteTLSCAs,
		Buffer:   cfg.LogRemoteBuffer,
		Facility: cfg.LogRemoteFacility,
		AppName:  cfg.LogRemoteAppName,
	})
}

// startSending runs remote until the function it returns is called, which
// then waits at most remoteLogStopTimeout for the lines still waiting to
// be sent. Sending goes on after ctx is done, so that the lines written
// while sww stops are sent too.
func startSending(
	ctx context.Context, remote *remotelog.Sender, processLog *slog.Logger,
) func() {
	sending, stop := context.WithCancel(context.WithoutCancel(ctx))
	sent := make(chan struct{})

	go func() {
		remote.Run(sending, processLog)
		close(sent)
	}()

	return func() {
		stop()

		select {
		case <-sent:
		case <-time.After(remoteLogStopTimeout):
		}
	}
}

// serve serves requests on listener, writes the state files as they are
// due, takes in an admin's edits of them, reads the rule files again as
// they change, and the lookup database when it is replaced, fetches the
// lists the settings name by URL as they are due, and sends the alerts,
// until ctx is done. Then it gives the requests in progress
// shutdownTimeout to finish, and writes every state file, alerts.json with
// the alerts still waiting.
func serve(
	ctx context.Context, server *discriminator.Server, listener net.Listener,
	files *state.Files, ruleFiles *rules.Files, alertQueue *alerts.Queue,
	processLog *slog.Logger,
) int {
	served := make(chan error, 1)

	go func() {
		served <- server.Serve(listener)
	}()

	writing, stopWriting := context.WithCancel(ctx)
	defer stopWriting()

	written := inBackground(func() { files.Run(writing) })
	watched := inBackground(func() { files.Watch(writing) })
	rulesWatched := inBackground(func() { ruleFiles.Watch(writing) })
	lookupFileWatched := inBackground(func() {
		if server.LookupFile != nil {
			server.LookupFile.Watch(writing)
		}
	})
	listsFetched := inBackground(func() { server.Lists.Run(writing) })
	alertsSent := inBackground(func() { alertQueue.Run(writing) })

	select {
	case err := <-served:
		processLog.Error("serving failed", "error", err.Error())

		return 1
	case <-ctx.Done():
	}

	processLog.Info("stopping")

	shutdownCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx),
		shutdownTimeout)
	defer cancel()

	err := server.Shutdown(shutdownCtx)
	if err != nil {
		processLog.Warn("requests still in progress were cut off",
			"error", err.Error())

		_ = server.Close()
	}

	err = <-served
	if !errors.Is(err, http.ErrServerClosed) {
		processLog.Error("serving failed", "error", err.Error())

		return 1
	}

	// Run and Watch have ended, so nothing else reads or writes the
	// files, and no alert is being sent, so that alerts.json keeps every
	// alert not yet sent. Every request has ended too, but for one cut off
	// because Shutdown timed out, which adds to its client's history only as
	// it ends. That can be after this write, and then that request is missing
	// from clients.json.
	<-written
	<-watched
	<-rulesWatched
	<-lookupFileWatched
	<-listsFetched
	<-alertsSent

	err = files.WriteAll()
	if err != nil {
		processLog.Error("writing the state files failed", "error", err.Error())

		return 1
	}

	processLog.Info("stopped")

	return 0
}

// inBackground runs task on a goroutine of its own, and returns a channel
// that is closed once task has returned.
func inBackground(task func()) <-chan struct{} {
	done := make(chan struct{})

	go func() {
		task()
		close(done)
	}()

	return done
}