authz_client/client.go

package client

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"math"
	"net/http"
	"reflect"
	"strconv"
	"sync"
	"time"

	goamqp "codeberg.org/messaging/go-messaging-amqp"
	spec "codeberg.org/messaging/messaging"
)

// SequenceHeader carries the global sequence number of the authz-service read view the
// /authz snapshot was read at.
const SequenceHeader = "X-Authz-Sequence"

// pending marks a fact cleared by an event without a sequence number. No event can
// override it; only a snapshot can.
const pending = math.MaxInt

// CompanyPrivileges contains the privileges for a combination of email address and company id
type CompanyPrivileges struct {
	Admin      bool `json:"admin"`
	Company    bool `json:"company"`
	Consumer   bool `json:"consumer"`
	Time       bool `json:"time"`
	Invoicing  bool `json:"invoicing"`
	Accounting bool `json:"accounting"`
	Supplier   bool `json:"supplier"`
	Salary     bool `json:"salary"`
}

// field returns the flag for privilege, or nil for an unknown privilege.
func (c *CompanyPrivileges) field(privilege Privilege) *bool {
	switch privilege {
	case PrivilegeAdmin:
		return &c.Admin
	case PrivilegeCompany:
		return &c.Company
	case PrivilegeConsumer:
		return &c.Consumer
	case PrivilegeTime:
		return &c.Time
	case PrivilegeInvoicing:
		return &c.Invoicing
	case PrivilegeAccounting:
		return &c.Accounting
	case PrivilegeSupplier:
		return &c.Supplier
	case PrivilegeSalary:
		return &c.Salary
	}
	return nil
}

// companyState is what the handler knows about one email in one company. Every
// privilege event for a company comes from authz-service's Company aggregate, so
// its global sequence number orders the events for the pair. The four routing
// keys arrive on separate queues in any order; a fact is only overwritten by a
// fact with a higher sequence number.
type companyState struct {
	privileges CompanyPrivileges
	// privilegeSeq is the sequence number of the last fact about each privilege.
	privilegeSeq map[Privilege]int
	// member is whether the user belongs to the company, as of memberSeq. A
	// User.Removed stamps every privilege with its sequence number, so a grant
	// older than the removal can't come back.
	member    bool
	memberSeq int
}

func newCompanyState() *companyState {
	return &companyState{privilegeSeq: map[Privilege]int{}}
}

// clearPending turns facts cleared without a sequence number into unordered facts, so
// the snapshot being merged replaces them.
func (s *companyState) clearPending() {
	if s.memberSeq == pending {
		s.memberSeq = 0
	}
	for p, seq := range s.privilegeSeq {
		if seq == pending {
			s.privilegeSeq[p] = 0
		}
	}
}

// maxSeq is the highest sequence number of any fact in the state.
func (s *companyState) maxSeq() int {
	m := s.memberSeq
	for _, seq := range s.privilegeSeq {
		m = max(m, seq)
	}
	return m
}

// PrivilegeHandler processes PrivilegeAdded-events and fetches the initial set of privileges from an authz-service
type PrivilegeHandler struct {
	*sync.RWMutex
	client  *http.Client
	baseURL string
	apiKey  string
	state   map[string]map[string]*companyState
	// floor is the sequence number of the newest snapshot. Its effects are all in
	// the state, so an event at or below it is stale.
	floor int
	// retries is how many times Fetch retries a snapshot authz-service isn't ready to serve.
	retries    int
	retryDelay time.Duration
}

// OptsFunc is used to configure the PrivilegeHandler
type OptsFunc func(handler *PrivilegeHandler)

// WithBaseURL sets the base URL to the authz-service
func WithBaseURL(url string) OptsFunc {
	return func(handler *PrivilegeHandler) {
		handler.baseURL = url
	}
}

// WithAPIKey sets an API key used as a Bearer token when fetching privileges
func WithAPIKey(key string) OptsFunc {
	return func(handler *PrivilegeHandler) {
		handler.apiKey = key
	}
}

// New creates a new PrivilegeHandler. Pass OptsFuncs to configure.
func New(opts ...OptsFunc) *PrivilegeHandler {
	handler := &PrivilegeHandler{
		RWMutex:    &sync.RWMutex{},
		client:     &http.Client{Timeout: 30 * time.Second},
		baseURL:    "http://authz-service",
		state:      map[string]map[string]*companyState{},
		retries:    60,
		retryDelay: time.Second,
	}
	for _, opt := range opts {
		opt(handler)
	}
	return handler
}

// Fetch a snapshot of all privileges from an authz-service and merge it into the state.
//
// Call it after the AMQP connection has started: authz-service only serves a snapshot
// that includes every stored event, so together with the bound queues no event is
// missed. While its read view is behind it answers 503, and Fetch retries. The snapshot
// is applied as facts at its sequence number: a fact from an event newer than the
// snapshot is kept, everything else is replaced, and a pair missing from the snapshot is
// removed. A snapshot older than one already merged is ignored.
func (h *PrivilegeHandler) Fetch() error {
	for attempt := 0; ; attempt++ {
		err := h.fetch()
		if !errors.Is(err, errNotReady) || attempt >= h.retries {
			return err
		}
		time.Sleep(h.retryDelay)
	}
}

var errNotReady = errors.New("fetch privileges: authz-service read view not ready")

func (h *PrivilegeHandler) fetch() error {
	req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("%s/authz", h.baseURL), nil)
	if err != nil {
		return err
	}

	if h.apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+h.apiKey)
	}

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

	if resp.StatusCode == http.StatusServiceUnavailable {
		return errNotReady
	}
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("fetch privileges: unexpected status %s", resp.Status)
	}

	buff, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	var snapshot map[string]map[string]CompanyPrivileges
	if err := json.Unmarshal(buff, &snapshot); err != nil {
		return err
	}

	seq := 0
	if v := resp.Header.Get(SequenceHeader); v != "" {
		seq, err = strconv.Atoi(v)
		if err != nil || seq < 0 {
			return fmt.Errorf("fetch privileges: invalid %s header %q", SequenceHeader, v)
		}
	} else {
		// ponytail: only during the rollout window, before authz-service sends the header.
		slog.Warn("authz snapshot has no sequence number; events older than it can revert it", "header", SequenceHeader)
	}

	h.Lock()
	defer h.Unlock()
	h.merge(snapshot, seq)
	return nil
}

func (h *PrivilegeHandler) merge(snapshot map[string]map[string]CompanyPrivileges, seq int) {
	if seq < h.floor {
		// A newer snapshot is already merged; this one could only bring back removed state.
		slog.Warn("ignoring authz snapshot older than the one already merged", "seq", seq, "floor", h.floor)
		return
	}
	for email, companies := range h.state {
		for companyID, s := range companies {
			s.clearPending()
			if _, exists := snapshot[email][companyID]; exists {
				continue
			}
			if s.maxSeq() <= seq {
				// Nothing newer than the snapshot; the floor keeps stale events out.
				delete(companies, companyID)
				continue
			}
			removeUser(s, seq, false)
		}
		if len(companies) == 0 {
			delete(h.state, email)
		}
	}
	for email, companies := range snapshot {
		for companyID, privileges := range companies {
			s := h.company(email, companyID)
			if s.memberSeq <= seq {
				s.member = true
				s.memberSeq = seq
			}
			for _, p := range AllPrivilege {
				if s.privilegeSeq[p] <= seq {
					*s.privileges.field(p) = *privileges.field(p)
					s.privilegeSeq[p] = seq
				}
			}
		}
	}
	h.floor = max(h.floor, seq)
}

func (h *PrivilegeHandler) Setup() []goamqp.Setup {
	return []goamqp.Setup{
		goamqp.TransientEventStreamConsumer("User.Added", process[UserAdded](h)),
		goamqp.TransientEventStreamConsumer("User.Removed", process[UserRemoved](h)),
		goamqp.TransientEventStreamConsumer("Privilege.Added", process[PrivilegeAdded](h)),
		goamqp.TransientEventStreamConsumer("Privilege.Removed", process[PrivilegeRemoved](h)),
	}
}

// privilegeEvent is the set of events Process handles.
type privilegeEvent interface {
	UserAdded | UserRemoved | PrivilegeAdded | PrivilegeRemoved
}

// process adapts Process to a typed go-messaging-amqp handler.
func process[T privilegeEvent](h *PrivilegeHandler) spec.EventHandler[T] {
	return func(_ context.Context, event spec.ConsumableEvent[T]) error {
		return h.Process(&event.Payload)
	}
}

// Process privilege-related events and update the internal state.
//
// Events are applied by sequence number, not by arrival order. An event without a
// sequence number can't be ordered, so it fails closed: a removal is applied and an
// addition is dropped.
func (h *PrivilegeHandler) Process(msg any) error {
	h.Lock()
	defer h.Unlock()

	switch ev := msg.(type) {
	case *UserAdded:
		if h.stale(ev.SequenceNo, true, ev) {
			return nil
		}
		s := h.company(ev.Email, ev.CompanyID)
		if ev.SequenceNo > s.memberSeq {
			s.member = true
			s.memberSeq = ev.SequenceNo
		}
		return nil
	case *UserRemoved:
		if h.stale(ev.SequenceNo, false, ev) {
			return nil
		}
		removeUser(h.company(ev.Email, ev.CompanyID), ev.SequenceNo, ev.SequenceNo == 0)
		return nil
	case *PrivilegeAdded:
		if h.stale(ev.SequenceNo, true, ev) {
			return nil
		}
		h.setPrivilege(ev.Email, ev.CompanyID, ev.Privilege, ev.SequenceNo, true)
		return nil
	case *PrivilegeRemoved:
		if h.stale(ev.SequenceNo, false, ev) {
			return nil
		}
		h.setPrivilege(ev.Email, ev.CompanyID, ev.Privilege, ev.SequenceNo, false)
		return nil
	default:
		slog.Error("unexpected privilege message type", "type", reflect.TypeOf(msg).String())
		return fmt.Errorf("unexpected event type: '%s'", reflect.TypeOf(msg))
	}
}

// stale reports whether an event must be skipped: an addition without a sequence
// number, or any event already covered by the latest snapshot.
func (h *PrivilegeHandler) stale(seq int, addition bool, ev any) bool {
	if seq < 0 || seq >= pending {
		slog.Error("dropping privilege event with invalid sequence number", "type", reflect.TypeOf(ev).String(), "seq", seq)
		return true
	}
	if seq == 0 {
		if addition {
			slog.Warn("dropping privilege addition without sequence number", "type", reflect.TypeOf(ev).String())
		}
		return addition
	}
	return seq <= h.floor
}

func (h *PrivilegeHandler) company(email, companyID string) *companyState {
	companies, exists := h.state[email]
	if !exists {
		companies = map[string]*companyState{}
		h.state[email] = companies
	}
	s, exists := companies[companyID]
	if !exists {
		s = newCompanyState()
		companies[companyID] = s
	}
	return s
}

// removeUser removes the user at seq, overriding every fact older than seq. An
// unordered removal (all) overrides every fact and marks them pending, so no event can
// bring them back before the next snapshot.
func removeUser(s *companyState, seq int, all bool) {
	stamp := seq
	if all {
		stamp = pending
	}
	if all || s.memberSeq < seq {
		s.member = false
		s.memberSeq = max(s.memberSeq, stamp)
	}
	for _, p := range AllPrivilege {
		if all || s.privilegeSeq[p] < seq {
			*s.privileges.field(p) = false
			s.privilegeSeq[p] = max(s.privilegeSeq[p], stamp)
		}
	}
}

func (h *PrivilegeHandler) setPrivilege(email, companyID string, privilege Privilege, seq int, set bool) {
	if !privilege.IsValid() {
		return
	}
	s := h.company(email, companyID)
	if seq != 0 && seq <= s.privilegeSeq[privilege] {
		return
	}
	*s.privileges.field(privilege) = set
	if seq == 0 {
		// Only an unordered removal gets here; keep it until the next snapshot.
		seq = pending
	}
	s.privilegeSeq[privilege] = max(s.privilegeSeq[privilege], seq)
	// authz-service's aggregate adds the user when a privilege is granted.
	if set && seq > s.memberSeq {
		s.member = true
		s.memberSeq = seq
	}
}

// CompaniesByUser return a slice of company ids matching the provided email and predicate func
func (h *PrivilegeHandler) CompaniesByUser(email string, predicate func(privileges CompanyPrivileges) bool) []string {
	h.RLock()
	defer h.RUnlock()
	result := []string{}
	for k, s := range h.state[email] {
		if s.member && predicate(s.privileges) {
			result = append(result, k)
		}
	}
	return result
}

// IsAllowed return true if the provided predicate return true for the privileges matching the provided email and companyID, return false otherwise
func (h *PrivilegeHandler) IsAllowed(email, companyID string, predicate func(privileges CompanyPrivileges) bool) bool {
	h.RLock()
	defer h.RUnlock()
	if s, exists := h.state[email][companyID]; exists && s.member {
		return predicate(s.privileges)
	}
	return false
}