server-scripts/haven-notify/discord.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"
)

const (
	deliveryTimeout   = 10 * time.Second
	attemptTimeout    = 5 * time.Second
	maximumAttempts   = 3
	maximumRetryDelay = 2 * time.Second
	truncationMarker  = "… [truncated]"
)

var (
	errWebhookDelivery = errors.New("webhook delivery failed")
	errWebhookTimeout  = errors.New("webhook delivery timed out")
)

type sleepFunc func(context.Context, time.Duration) error

type discordClient struct {
	webhookURL string
	httpClient *http.Client
	sleep      sleepFunc
}

type discordPayload struct {
	Content         string                 `json:"content,omitempty"`
	Embeds          []discordEmbed         `json:"embeds,omitempty"`
	AllowedMentions discordAllowedMentions `json:"allowed_mentions"`
}

type discordAllowedMentions struct {
	Parse []string `json:"parse"`
}

type discordEmbed struct {
	Title       string              `json:"title,omitempty"`
	Description string              `json:"description,omitempty"`
	Color       int                 `json:"color,omitempty"`
	Timestamp   string              `json:"timestamp,omitempty"`
	Fields      []discordEmbedField `json:"fields,omitempty"`
	Footer      *discordEmbedFooter `json:"footer,omitempty"`
}

type discordEmbedField struct {
	Name   string `json:"name"`
	Value  string `json:"value"`
	Inline bool   `json:"inline"`
}

type discordEmbedFooter struct {
	Text string `json:"text"`
}

func newDiscordClient(webhookURL string, client *http.Client, sleep sleepFunc) (*discordClient, error) {
	parsed, err := url.Parse(webhookURL)
	if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") {
		return nil, errors.New("WEBHOOK_URL must be an absolute HTTP(S) URL")
	}
	if client == nil {
		client = &http.Client{Timeout: attemptTimeout}
	}
	if sleep == nil {
		sleep = sleepContext
	}
	return &discordClient{webhookURL: webhookURL, httpClient: client, sleep: sleep}, nil
}

func (c *discordClient) send(ctx context.Context, payload discordPayload) error {
	normalizeDiscordPayload(&payload)
	body, err := json.Marshal(payload)
	if err != nil {
		return fmt.Errorf("%w: marshal payload", errWebhookDelivery)
	}

	deliveryCtx, cancel := context.WithTimeout(ctx, deliveryTimeout)
	defer cancel()

	var lastErr error
	for attempt := 0; attempt < maximumAttempts; attempt++ {
		attemptCtx, attemptCancel := context.WithTimeout(deliveryCtx, attemptTimeout)
		request, requestErr := http.NewRequestWithContext(attemptCtx, http.MethodPost, c.webhookURL, bytes.NewReader(body))
		if requestErr != nil {
			attemptCancel()
			return fmt.Errorf("%w: create request", errWebhookDelivery)
		}
		request.Header.Set("Content-Type", "application/json")

		response, requestErr := c.httpClient.Do(request)
		if requestErr == nil {
			delay := retryDelay(response, attempt)
			_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4<<10))
			_ = response.Body.Close()
			attemptCancel()

			if response.StatusCode >= 200 && response.StatusCode < 300 {
				return nil
			}
			lastErr = fmt.Errorf("%w: HTTP %d", errWebhookDelivery, response.StatusCode)
			if !retryableStatus(response.StatusCode) || attempt == maximumAttempts-1 {
				return lastErr
			}
			if err := c.sleep(deliveryCtx, delay); err != nil {
				return errWebhookTimeout
			}
			continue
		}

		attemptCancel()
		lastErr = requestErr
		if deliveryCtx.Err() != nil {
			return errWebhookTimeout
		}
		if attempt == maximumAttempts-1 {
			if isTimeout(requestErr) {
				return fmt.Errorf("%w: %v", errWebhookTimeout, requestErr)
			}
			return fmt.Errorf("%w: %v", errWebhookDelivery, requestErr)
		}
		if err := c.sleep(deliveryCtx, fallbackRetryDelay(attempt)); err != nil {
			return errWebhookTimeout
		}
	}

	return fmt.Errorf("%w: %v", errWebhookDelivery, lastErr)
}

func retryableStatus(status int) bool {
	return status == http.StatusTooManyRequests || status >= 500
}

func retryDelay(response *http.Response, attempt int) time.Duration {
	if response.StatusCode == http.StatusTooManyRequests {
		for _, header := range []string{"Retry-After", "X-RateLimit-Reset-After"} {
			if delay, ok := parseRetryDelay(response.Header.Get(header)); ok {
				return min(delay, maximumRetryDelay)
			}
		}
	}
	return fallbackRetryDelay(attempt)
}

func parseRetryDelay(value string) (time.Duration, bool) {
	value = strings.TrimSpace(value)
	if value == "" {
		return 0, false
	}
	if seconds, err := strconv.ParseFloat(value, 64); err == nil && seconds >= 0 {
		return time.Duration(seconds * float64(time.Second)), true
	}
	if timestamp, err := http.ParseTime(value); err == nil {
		return max(time.Until(timestamp), 0), true
	}
	return 0, false
}

func fallbackRetryDelay(attempt int) time.Duration {
	return 250 * time.Millisecond * time.Duration(1<<attempt)
}

func sleepContext(ctx context.Context, delay time.Duration) error {
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}

func isTimeout(err error) bool {
	if errors.Is(err, context.DeadlineExceeded) {
		return true
	}
	var netErr net.Error
	return errors.As(err, &netErr) && netErr.Timeout()
}

func normalizeDiscordPayload(payload *discordPayload) {
	payload.AllowedMentions.Parse = []string{}
	payload.Content = truncateRunes(payload.Content, 2000)
	if len(payload.Embeds) > 10 {
		payload.Embeds = payload.Embeds[:10]
	}

	for embedIndex := range payload.Embeds {
		embed := &payload.Embeds[embedIndex]
		embed.Title = truncateRunes(embed.Title, 256)
		embed.Description = truncateRunes(embed.Description, 4096)
		if len(embed.Fields) > 25 {
			embed.Fields = embed.Fields[:25]
		}
		for fieldIndex := range embed.Fields {
			embed.Fields[fieldIndex].Name = truncateRunes(embed.Fields[fieldIndex].Name, 256)
			embed.Fields[fieldIndex].Value = truncateRunes(embed.Fields[fieldIndex].Value, 1024)
		}
		if embed.Footer != nil {
			embed.Footer.Text = truncateRunes(embed.Footer.Text, 2048)
		}
	}

	for discordEmbedCharacters(payload.Embeds) > 6000 {
		removed := false
		for embedIndex := len(payload.Embeds) - 1; embedIndex >= 0; embedIndex-- {
			embed := &payload.Embeds[embedIndex]
			if len(embed.Fields) > 0 {
				embed.Fields = embed.Fields[:len(embed.Fields)-1]
				removed = true
				break
			}
		}
		if !removed {
			break
		}
	}

	remainingOverflow := discordEmbedCharacters(payload.Embeds) - 6000
	for embedIndex := len(payload.Embeds) - 1; remainingOverflow > 0 && embedIndex >= 0; embedIndex-- {
		embed := &payload.Embeds[embedIndex]
		if embed.Footer != nil && embed.Footer.Text != "" {
			newLimit := max(utf8.RuneCountInString(embed.Footer.Text)-remainingOverflow, 0)
			embed.Footer.Text = truncateRunes(embed.Footer.Text, newLimit)
		}
		remainingOverflow = discordEmbedCharacters(payload.Embeds) - 6000
		if remainingOverflow > 0 && embed.Description != "" {
			newLimit := max(utf8.RuneCountInString(embed.Description)-remainingOverflow, 0)
			embed.Description = truncateRunes(embed.Description, newLimit)
		}
		remainingOverflow = discordEmbedCharacters(payload.Embeds) - 6000
		if remainingOverflow > 0 && embed.Title != "" {
			newLimit := max(utf8.RuneCountInString(embed.Title)-remainingOverflow, 0)
			embed.Title = truncateRunes(embed.Title, newLimit)
		}
		remainingOverflow = discordEmbedCharacters(payload.Embeds) - 6000
	}
}

func discordEmbedCharacters(embeds []discordEmbed) int {
	total := 0
	for _, embed := range embeds {
		total += utf8.RuneCountInString(embed.Title)
		total += utf8.RuneCountInString(embed.Description)
		if embed.Footer != nil {
			total += utf8.RuneCountInString(embed.Footer.Text)
		}
		for _, field := range embed.Fields {
			total += utf8.RuneCountInString(field.Name)
			total += utf8.RuneCountInString(field.Value)
		}
	}
	return total
}

func truncateRunes(value string, limit int) string {
	runes := []rune(value)
	if len(runes) <= limit {
		return value
	}
	if limit <= 0 {
		return ""
	}
	marker := []rune(truncationMarker)
	if limit <= len(marker) {
		return string(marker[:limit])
	}
	return string(runes[:limit-len(marker)]) + truncationMarker
}