Warp/server/internal/pushserver/gateway.go

package pushserver

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"sync"
	"time"

	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// The push gateway client: the Push Gateway API's one endpoint,
// POST /_matrix/push/v1/notify, called for every pusher of a notified user.
//
// Pushes are sent from a small pool of senders fed by a bounded queue, so
// that a slow or dead gateway delays its own pushes and nothing else: the
// evaluation worker never waits on the network. A push that fails with a
// server error or a timeout is retried a few times with a growing pause;
// one the gateway refuses outright is not. When the queue is full a push
// is dropped and logged. Pushes are a best-effort wake-up — the
// notification itself is already counted, and the device sees it on its
// next sync whether or not the push arrived.

const (
	// gatewaySenders is how many pushes are in flight at once.
	gatewaySenders = 4

	// gatewayQueueSize bounds the pushes waiting for a sender.
	gatewayQueueSize = 1024

	// gatewayAttempts is how many times one push is tried.
	gatewayAttempts = 3

	// maxGatewayResponse bounds what is read of a gateway's answer, which
	// is a list of rejected push keys and should be tiny.
	maxGatewayResponse = 64 << 10
)

// gatewayRetryDelay is the pause before the given retry (1-based).
var gatewayRetryDelay = func(attempt int) time.Duration {
	return time.Duration(1<<(2*(attempt-1))) * time.Second
}

// notificationPush is the notification object of a push, without the
// devices, which are per pusher.
type notificationPush struct {
	EventID           string          `json:"event_id,omitempty"`
	RoomID            string          `json:"room_id,omitempty"`
	Type              string          `json:"type,omitempty"`
	Sender            string          `json:"sender,omitempty"`
	SenderDisplayName string          `json:"sender_display_name,omitempty"`
	RoomName          string          `json:"room_name,omitempty"`
	RoomAlias         string          `json:"room_alias,omitempty"`
	UserIsTarget      *bool           `json:"user_is_target,omitempty"`
	Membership        string          `json:"membership,omitempty"`
	Prio              string          `json:"prio,omitempty"`
	Content           json.RawMessage `json:"content,omitempty"`
	Counts            pushCounts      `json:"counts"`
	Devices           []pushDevice    `json:"devices"`
}

type pushCounts struct {
	Unread int `json:"unread"`
}

type pushDevice struct {
	AppID     string                     `json:"app_id"`
	PushKey   string                     `json:"pushkey"`
	PushKeyTS int64                      `json:"pushkey_ts,omitempty"`
	Data      map[string]json.RawMessage `json:"data,omitempty"`
	Tweaks    map[string]json.RawMessage `json:"tweaks,omitempty"`
}

// pushJob is one push to one pusher.
type pushJob struct {
	localpart string
	pusher    storage.Pusher
	note      notificationPush
	tweaks    map[string]json.RawMessage
}

type gatewayOptions struct {
	client         *http.Client
	includeContent bool

	// rejected is told about a push key the gateway refused.
	rejected func(localpart, appID, pushKey string)

	log *slog.Logger
}

type gateway struct {
	opts  gatewayOptions
	queue chan pushJob
	wg    sync.WaitGroup
	ctx   context.Context
}

func newGateway(opts gatewayOptions) *gateway {
	return &gateway{opts: opts, queue: make(chan pushJob, gatewayQueueSize)}
}

func (g *gateway) start(ctx context.Context) {
	g.ctx = ctx
	for range gatewaySenders {
		g.wg.Add(1)
		go g.send()
	}
}

func (g *gateway) wait() { g.wg.Wait() }

func (g *gateway) enqueue(job pushJob) {
	select {
	case g.queue <- job:
	default:
		g.opts.log.Warn("push queue is full; dropping a push",
			"localpart", job.localpart, "app_id", job.pusher.AppID)
	}
}

func (g *gateway) send() {
	defer g.wg.Done()
	for {
		select {
		case <-g.ctx.Done():
			return
		case job := <-g.queue:
			g.deliver(job)
		}
	}
}

// deliver sends one push, retrying what is worth retrying.
func (g *gateway) deliver(job pushJob) {
	var data pusherData
	_ = json.Unmarshal(job.pusher.Data, &data)
	body, err := g.payload(job, data)
	if err != nil {
		g.opts.log.Error("building a push failed", "localpart", job.localpart, "error", err)
		return
	}

	for attempt := 1; attempt <= gatewayAttempts; attempt++ {
		rejected, retry, err := g.post(data.URL, body)
		if err == nil {
			for _, key := range rejected {
				if key == job.pusher.PushKey {
					g.opts.rejected(job.localpart, job.pusher.AppID, job.pusher.PushKey)
				}
			}
			return
		}
		if !retry || attempt == gatewayAttempts {
			g.opts.log.Warn("push gateway request failed",
				"localpart", job.localpart, "app_id", job.pusher.AppID, "attempts", attempt, "error", err)
			return
		}
		select {
		case <-g.ctx.Done():
			return
		case <-time.After(gatewayRetryDelay(attempt)):
		}
	}
}

// payload builds the request body for one pusher.
//
// The pusher's data goes to the gateway as the device's data, less the
// URL, which is the homeserver's business. A pusher whose format is
// event_id_only gets only the IDs and the counts: the app fetches the
// event itself, so nothing of it passes through the gateway.
func (g *gateway) payload(job pushJob, data pusherData) ([]byte, error) {
	deviceData := map[string]json.RawMessage{}
	if err := json.Unmarshal(job.pusher.Data, &deviceData); err != nil {
		return nil, fmt.Errorf("pusher data: %w", err)
	}
	delete(deviceData, "url")

	device := pushDevice{
		AppID:     job.pusher.AppID,
		PushKey:   job.pusher.PushKey,
		PushKeyTS: job.pusher.PushKeyTSMS / 1000,
		Data:      deviceData,
		Tweaks:    job.tweaks,
	}

	note := job.note
	if data.Format == "event_id_only" {
		note = notificationPush{
			EventID: note.EventID,
			RoomID:  note.RoomID,
			Prio:    note.Prio,
			Counts:  note.Counts,
		}
	} else if !g.opts.includeContent {
		note.Content = nil
	}
	note.Devices = []pushDevice{device}
	return json.Marshal(struct {
		Notification notificationPush `json:"notification"`
	}{note})
}

// post sends a body to a gateway. It returns the push keys the gateway
// rejected, or whether a failure is worth another try.
func (g *gateway) post(url string, body []byte) (rejected []string, retry bool, err error) {
	req, err := http.NewRequestWithContext(g.ctx, http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return nil, false, err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := g.opts.client.Do(req)
	if err != nil {
		return nil, g.ctx.Err() == nil, err
	}
	defer func() { _ = resp.Body.Close() }()

	raw, err := io.ReadAll(io.LimitReader(resp.Body, maxGatewayResponse))
	if err != nil {
		return nil, true, err
	}
	switch {
	case resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500:
		return nil, true, fmt.Errorf("gateway answered %d", resp.StatusCode)
	case resp.StatusCode < 200 || resp.StatusCode >= 300:
		return nil, false, fmt.Errorf("gateway answered %d", resp.StatusCode)
	}

	var answer struct {
		Rejected []string `json:"rejected"`
	}
	// A gateway that says 200 and sends something unreadable delivered
	// the push all the same.
	_ = json.Unmarshal(raw, &answer)
	return answer.Rejected, false, nil
}