server-scripts/haven-notify/server.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"mime"
	"net"
	"net/http"
	"strings"
	"time"
)

var netListen = net.Listen

const (
	maxRequestBody  = 64 << 10
	shutdownTimeout = 10 * time.Second
)

var (
	errInvalidPayload       = errors.New("invalid payload")
	errUnsupportedMediaType = errors.New("unsupported media type")
)

type applicationConfig struct {
	WebhookURL string
	HTTPClient *http.Client
	Sleep      sleepFunc
	Logger     *log.Logger
}

type application struct {
	discord  *discordClient
	renderer *templateRenderer
	logger   *log.Logger
}

type notificationRequest struct {
	Title   *string `json:"title"`
	Message *string `json:"message"`
}

type extraRequest struct {
	Name  *string `json:"name"`
	Value *string `json:"value"`
}

type extraRequests []extraRequest

type backupRequest struct {
	Title          *string       `json:"title"`
	Asset          *string       `json:"asset"`
	BackupSizeInMB *float64      `json:"backupSizeInMB"`
	Extra          extraRequests `json:"extra"`
}

type updateRequest struct {
	Host  *string  `json:"host"`
	Asset *string  `json:"asset"`
	Time  *float64 `json:"time"`
}

type errorRequest struct {
	Caller   *string       `json:"caller"`
	Message  *string       `json:"message"`
	Critical *bool         `json:"critical"`
	Extra    extraRequests `json:"extra"`
}

func newApplication(cfg applicationConfig) (*application, error) {
	logger := cfg.Logger
	if logger == nil {
		logger = log.New(io.Discard, "", 0)
	}

	renderer, err := newTemplateRenderer()
	if err != nil {
		return nil, fmt.Errorf("load templates: %w", err)
	}

	discord, err := newDiscordClient(cfg.WebhookURL, cfg.HTTPClient, cfg.Sleep)
	if err != nil {
		return nil, err
	}

	return &application{discord: discord, renderer: renderer, logger: logger}, nil
}

func (a *application) handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/notify", a.requireMethod(http.MethodPost, a.notifyHandler))
	mux.HandleFunc("/template/notify/backup", a.requireMethod(http.MethodPost, a.backupHandler))
	mux.HandleFunc("/template/notify/update", a.requireMethod(http.MethodPost, a.updateHandler))
	mux.HandleFunc("/template/notify/error", a.requireMethod(http.MethodPost, a.errorHandler))
	mux.HandleFunc("/ready", a.requireProbeMethod(a.readinessHandler))
	mux.HandleFunc("/live", a.requireProbeMethod(a.livenessHandler))
	mux.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
		writeText(w, http.StatusNotFound, "Not found")
	})
	return mux
}

func newHTTPServer(handler http.Handler) *http.Server {
	return &http.Server{
		Handler:           handler,
		ReadHeaderTimeout: 5 * time.Second,
		ReadTimeout:       10 * time.Second,
		WriteTimeout:      15 * time.Second,
		IdleTimeout:       60 * time.Second,
	}
}

func serve(ctx context.Context, server *http.Server, listener net.Listener) error {
	errCh := make(chan error, 1)
	go func() {
		errCh <- server.Serve(listener)
	}()

	select {
	case err := <-errCh:
		return err
	case <-ctx.Done():
		shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
		defer cancel()
		if err := server.Shutdown(shutdownCtx); err != nil {
			return fmt.Errorf("graceful shutdown: %w", err)
		}
		return <-errCh
	}
}

func (a *application) requireMethod(method string, next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != method {
			w.Header().Set("Allow", method)
			writeText(w, http.StatusMethodNotAllowed, "Method not allowed")
			return
		}
		next(w, r)
	}
}

func (a *application) requireProbeMethod(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet && r.Method != http.MethodHead {
			w.Header().Set("Allow", "GET, HEAD")
			writeText(w, http.StatusMethodNotAllowed, "Method not allowed")
			return
		}
		next(w, r)
	}
}

func (a *application) notifyHandler(w http.ResponseWriter, r *http.Request) {
	var request notificationRequest
	if err := decodeRequest(w, r, &request, "title", "message"); err != nil {
		a.writeDecodeError(w, r, err)
		return
	}
	if !requiredString(request.Title) || !requiredString(request.Message) {
		a.writeDecodeError(w, r, errInvalidPayload)
		return
	}

	payload := discordPayload{
		Content: "**" + *request.Title + "**\n" + *request.Message,
	}
	if !a.send(w, r, payload) {
		return
	}
	writeText(w, http.StatusOK, "Notification sent")
}

func (a *application) backupHandler(w http.ResponseWriter, r *http.Request) {
	var request backupRequest
	if err := decodeRequest(w, r, &request, "title", "asset", "backupsizeinmb", "extra"); err != nil {
		a.writeDecodeError(w, r, err)
		return
	}
	if !requiredString(request.Title) || !requiredString(request.Asset) || request.BackupSizeInMB == nil || *request.BackupSizeInMB < 0 || !validExtra(request.Extra) {
		a.writeDecodeError(w, r, errInvalidPayload)
		return
	}

	payload, err := a.renderer.renderBackup(request)
	if err != nil {
		a.logger.Printf("Template rendering failed for %s", r.URL.Path)
		writeText(w, http.StatusInternalServerError, "Failed to render notification")
		return
	}
	if !a.send(w, r, payload) {
		return
	}
	writeText(w, http.StatusOK, "Notification sent")
}

func (a *application) updateHandler(w http.ResponseWriter, r *http.Request) {
	var request updateRequest
	if err := decodeRequest(w, r, &request, "host", "asset", "time"); err != nil {
		a.writeDecodeError(w, r, err)
		return
	}
	if !requiredString(request.Host) || !requiredString(request.Asset) || request.Time == nil || *request.Time < 0 {
		a.writeDecodeError(w, r, errInvalidPayload)
		return
	}

	payload, err := a.renderer.renderUpdate(request)
	if err != nil {
		a.logger.Printf("Template rendering failed for %s", r.URL.Path)
		writeText(w, http.StatusInternalServerError, "Failed to render notification")
		return
	}
	if !a.send(w, r, payload) {
		return
	}
	writeText(w, http.StatusOK, "Notification sent")
}

func (a *application) errorHandler(w http.ResponseWriter, r *http.Request) {
	var request errorRequest
	if err := decodeRequest(w, r, &request, "caller", "message", "critical", "extra"); err != nil {
		a.writeDecodeError(w, r, err)
		return
	}
	if !requiredString(request.Caller) || !requiredString(request.Message) || request.Critical == nil || !validExtra(request.Extra) {
		a.writeDecodeError(w, r, errInvalidPayload)
		return
	}

	payload, err := a.renderer.renderError(request)
	if err != nil {
		a.logger.Printf("Template rendering failed for %s", r.URL.Path)
		writeText(w, http.StatusInternalServerError, "Failed to render notification")
		return
	}
	if !a.send(w, r, payload) {
		return
	}
	writeText(w, http.StatusOK, "Notification sent")
}

func (a *application) readinessHandler(w http.ResponseWriter, _ *http.Request) {
	writeText(w, http.StatusOK, "Ready")
}

func (a *application) livenessHandler(w http.ResponseWriter, _ *http.Request) {
	writeText(w, http.StatusOK, "Alive")
}

func (a *application) send(w http.ResponseWriter, r *http.Request, payload discordPayload) bool {
	if err := a.discord.send(r.Context(), payload); err != nil {
		if errors.Is(err, errWebhookTimeout) {
			a.logger.Printf("Discord delivery timed out for %s", r.URL.Path)
			writeText(w, http.StatusGatewayTimeout, "Webhook timed out")
			return false
		}
		a.logger.Printf("Discord delivery failed for %s", r.URL.Path)
		writeText(w, http.StatusBadGateway, "Failed to send notification")
		return false
	}
	return true
}

func (a *application) writeDecodeError(w http.ResponseWriter, r *http.Request, err error) {
	if errors.Is(err, errUnsupportedMediaType) {
		a.logger.Printf("Rejected unsupported media type for %s", r.URL.Path)
		writeText(w, http.StatusUnsupportedMediaType, "Content-Type must be application/json")
		return
	}
	a.logger.Printf("Rejected invalid payload for %s", r.URL.Path)
	writeText(w, http.StatusBadRequest, "Invalid payload")
}

func decodeRequest(w http.ResponseWriter, r *http.Request, destination any, knownKeys ...string) error {
	mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err != nil || !strings.EqualFold(mediaType, "application/json") {
		return errUnsupportedMediaType
	}

	r.Body = http.MaxBytesReader(w, r.Body, maxRequestBody)
	data, err := io.ReadAll(r.Body)
	if err != nil {
		return fmt.Errorf("read request: %w", errInvalidPayload)
	}
	if len(bytes.TrimSpace(data)) == 0 {
		return errInvalidPayload
	}
	if err := rejectDuplicateKnownKeys(data, knownKeys); err != nil {
		return errInvalidPayload
	}
	if err := json.Unmarshal(data, destination); err != nil {
		return errInvalidPayload
	}
	return nil
}

func rejectDuplicateKnownKeys(data []byte, knownKeys []string) error {
	known := make(map[string]struct{}, len(knownKeys))
	for _, key := range knownKeys {
		known[strings.ToLower(key)] = struct{}{}
	}

	decoder := json.NewDecoder(bytes.NewReader(data))
	token, err := decoder.Token()
	if err != nil {
		return err
	}
	delim, ok := token.(json.Delim)
	if !ok || delim != '{' {
		return errInvalidPayload
	}

	seen := make(map[string]struct{}, len(knownKeys))
	for decoder.More() {
		token, err = decoder.Token()
		if err != nil {
			return err
		}
		key, ok := token.(string)
		if !ok {
			return errInvalidPayload
		}
		canonical := strings.ToLower(key)
		if _, isKnown := known[canonical]; isKnown {
			if _, duplicate := seen[canonical]; duplicate {
				return errInvalidPayload
			}
			seen[canonical] = struct{}{}
		}
		var value json.RawMessage
		if err := decoder.Decode(&value); err != nil {
			return err
		}
	}
	if _, err := decoder.Token(); err != nil {
		return err
	}
	if decoder.Decode(&struct{}{}) != io.EOF {
		return errInvalidPayload
	}
	return nil
}

func (e *extraRequest) UnmarshalJSON(data []byte) error {
	if err := rejectDuplicateKnownKeys(data, []string{"name", "value"}); err != nil {
		return err
	}
	type rawExtra extraRequest
	var decoded rawExtra
	if err := json.Unmarshal(data, &decoded); err != nil {
		return err
	}
	*e = extraRequest(decoded)
	return nil
}

func (e *extraRequests) UnmarshalJSON(data []byte) error {
	if bytes.Equal(bytes.TrimSpace(data), []byte("null")) {
		return errInvalidPayload
	}
	var decoded []extraRequest
	if err := json.Unmarshal(data, &decoded); err != nil {
		return err
	}
	*e = decoded
	return nil
}

func requiredString(value *string) bool {
	return value != nil && strings.TrimSpace(*value) != ""
}

func validExtra(extra extraRequests) bool {
	for _, field := range extra {
		if !requiredString(field.Name) || !requiredString(field.Value) {
			return false
		}
	}
	return true
}

func writeText(w http.ResponseWriter, status int, body string) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.WriteHeader(status)
	_, _ = io.WriteString(w, body)
}