Quota pools; accepting_new and quota_pool on POST /fmsgid

README.md

@@ -87,14 +87,31 @@ Available columns:
 
 See `addresses.csv.example` for a complete example with all columns.
 
+## Quota pools
+
+Addresses can share one quota. The nullable `address.quota_pool` column holds the `address_lower` (the case folded address) of the pool owner; `NULL` means the address is its own pool. For example linked addresses such as `@alice_bot@example.com` can share the quota of `@alice@example.com`:
+
+```sql
+update address set quota_pool = '@alice@example.com'
+where address_lower = '@alice_bot@example.com';
+```
+
+- An address in a pool is checked against the pool owner's limits (`limit_*`) and the combined usage of the owner and every address whose `quota_pool` is the owner. Its own limits are kept but not used while it is pooled.
+- The pool owner reports the same combined usage, so a pool member's usage counts against the owner too.
+- `accepting_new` and `display_name` stay per address.
+- Pools are one level deep: a pool owner's own `quota_pool` should be `NULL`. `POST /fmsgid` enforces that; anything writing the column directly should too.
+- `quota_pool` references `address (address_lower)`, so the owner must exist first; deleting the owner sets its members' `quota_pool` back to `NULL`.
+
+Set `quota_pool` with `POST /fmsgid`, or have the system syncing your identity provider update the column directly. A role that does so needs `UPDATE (quota_pool)` on `address` besides `SELECT`.
+
 ## API Routes
 
 All routes are served over HTTPS under the `/fmsgid` path.
 
 | Method | Route | Description |
 |--------|-------|-------------|
-| `GET` | `/fmsgid/:address` | Lookup an fmsg address and return its details including display name, quotas, and usage. The address must be in fmsg format (`@user@example.com`). Returns `AddressDetail` JSON on success, `400` if the address is invalid, `404` if not found. |
-| `POST` | `/fmsgid` | Register a new address with default quotas (the `address` table's column defaults in `dd.sql`), idempotently. Accepts a JSON body with `address` and optional `display_name`. Returns `201` if the address was created, `200` if it already existed (never modifies an existing row — use CSV sync to change quotas or `accepting_new` on an existing address), `400` if the address is invalid. |
+| `GET` | `/fmsgid/:address` | Lookup an fmsg address and return its details including display name, quotas, and usage. The address must be in fmsg format (`@user@example.com`). Returns `AddressDetail` JSON on success, `400` if the address is invalid, `404` if not found. For an address in a [quota pool](#quota-pools) the limits are the pool owner's, the usage is the whole pool's and `quotaPool` is the pool owner's address; `quotaPool` is omitted otherwise. `acceptingNew` is always the address's own. |
+| `POST` | `/fmsgid` | Register an address, idempotently. Accepts a JSON body with `address` and optional `display_name`, `accepting_new` (bool) and `quota_pool` (an fmsg address, see [Quota pools](#quota-pools)). A new address takes the `address` table's column defaults in `dd.sql` except for `accepting_new` and `quota_pool` when given. An existing address is left as it is except for `accepting_new` and `quota_pool` when given (`display_name` and quotas are never changed — use CSV sync for those). Setting `accepting_new` `false` stops the address accepting new messages; setting it `true` again reinstates it. `quota_pool` `""` or the address itself takes the address out of any pool. Returns `201` if the address was created, `200` if it already existed, both with `{"created":bool,"reinstated":bool}` where `reinstated` is `true` when the address existed, wasn't accepting new messages and now is. Returns `400` if the body or an address is invalid, `422` with `{"error":"..."}` if `quota_pool` names an address that doesn't exist or is in another pool, or the address owns a pool itself. |
 | `POST` | `/fmsgid/send` | Record a send transaction. Accepts an `AddressTx` JSON body with `address`, `ts` (timestamp), and `size`. The address is case folded, and every call is recorded, including several with the same timestamp. Returns `200`, or `400` if the body or address is invalid. |
 | `POST` | `/fmsgid/recv` | Record a receive transaction. Accepts an `AddressTx` JSON body with `address`, `ts` (timestamp), and `size`. The address is case folded, and every call is recorded, including several with the same timestamp. Returns `200`, or `400` if the body or address is invalid. |
 

dd.sql

@@ -32,6 +32,26 @@ create table if not exists address (
     limit_send_count_per_1d 	bigint	not null default 1000
 );
 
+-- quota pools: an address whose quota_pool is set (to the address_lower of
+-- another address, the pool owner) shares the pool owner's quota: it is
+-- checked against the owner's limits and the combined usage of the owner and
+-- every address whose quota_pool is the owner. NULL means the address is its
+-- own pool; the owner's usage includes its pool members' too. Pools are one
+-- level deep, so a pool owner's own quota_pool should be NULL. A nullable
+-- column without a default is added without rewriting the table.
+alter table address add column if not exists quota_pool text;
+
+do $$
+begin
+	if not exists (select 1 from pg_constraint
+			where conrelid = 'address'::regclass and conname = 'address_quota_pool_fkey') then
+		alter table address add constraint address_quota_pool_fkey
+			foreign key (quota_pool) references address (address_lower) on delete set null;
+	end if;
+end $$;
+
+create index if not exists address_quota_pool on address (quota_pool);
+
 -- one row per message sent or received; several may share a timestamp
 create table if not exists address_tx (
 	id		bigserial			primary key,

src/fmsgid.go

@@ -2,12 +2,15 @@ package main
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"log"
 	"net/http"
 	"os"
 	"strings"
 
 	"github.com/gin-gonic/gin"
+	"github.com/jackc/pgx/v4"
 	"github.com/jackc/pgx/v4/pgxpool"
 	"github.com/joho/godotenv"
 	"golang.org/x/text/cases"
@@ -39,8 +42,32 @@ type AddressTx struct {
 type CreateAddressReq struct {
 	Address     string `json:"address"`
 	DisplayName string `json:"display_name"`
+	// AcceptingNew, when given, is set on the address whether it is new or
+	// already exists. When omitted a new address takes the column default and
+	// an existing one is left as it is.
+	AcceptingNew *bool `json:"accepting_new"`
+	// QuotaPool, when given, puts the address in the quota pool of the
+	// address it names (which must exist and be in no other pool) whether the
+	// address is new or already exists. "" or the address itself takes the
+	// address out of any pool. When omitted quota_pool is left as it is.
+	QuotaPool *string `json:"quota_pool"`
 }
 
+// CreateAddressResp is the result of POST /fmsgid.
+type CreateAddressResp struct {
+	// Created is true when the address didn't exist before.
+	Created bool `json:"created"`
+	// Reinstated is true when the address already existed, wasn't accepting
+	// new messages and now is, because the request set accepting_new true.
+	Reinstated bool `json:"reinstated"`
+}
+
+// quotaPoolError is a quota_pool the address can't join; its text is returned
+// to the caller.
+type quotaPoolError string
+
+func (e quotaPoolError) Error() string { return string(e) }
+
 type AddressDetail struct {
 	Address             string   `json:"address"`
 	DisplayName         string   `json:"displayName"`
@@ -60,6 +87,9 @@ type AddressDetail struct {
 	SendSizePer1d       int64    `json:"sendSizePer1d"`
 	SendCountPer1d      int64    `json:"sendCountPer1d"`
 	Tags                []string `json:"tags"`
+	// QuotaPool is the address of the pool owner whose limits and combined
+	// usage are reported, omitted when the address is in no pool.
+	QuotaPool *string `json:"quotaPool,omitempty"`
 }
 
 func initPool() error {
@@ -111,19 +141,23 @@ func getAddressDetail(c *gin.Context) {
 	}
 
 	var ad AddressDetail
+	var poolOwner string
 
+	// the limits are the quota pool's: the pool owner's when the address is in
+	// a pool, else its own
 	err = rows.Scan(&ad.Address, &ad.DisplayName, &ad.AcceptingNew, &ad.LimitRecvSizeTotal,
 		&ad.LimitRecvSizePerMsg, &ad.LimitRecvSizePer1d, &ad.LimitRecvCountPer1d,
 		&ad.LimitSendSizeTotal, &ad.LimitSendSizePerMsg, &ad.LimitSendSizePer1d,
-		&ad.LimitSendCountPer1d)
+		&ad.LimitSendCountPer1d, &poolOwner, &ad.QuotaPool)
 	if err != nil {
 		c.AbortWithError(500, err)
 		return
 	}
+	rows.Close()
 
-	// select actuals; an aggregate always returns one row, zeros when the
-	// address hasn't sent or received anything
-	err = pool.QueryRow(ctx, sqlActuals, addr).Scan(
+	// select actuals summed over the quota pool; an aggregate always returns
+	// one row, zeros when the pool hasn't sent or received anything
+	err = pool.QueryRow(ctx, sqlActuals, poolOwner).Scan(
 		&ad.RecvSizeTotal, &ad.RecvSizePer1d, &ad.RecvCountPer1d,
 		&ad.SendSizeTotal, &ad.SendSizePer1d, &ad.SendCountPer1d)
 	if err != nil {
@@ -134,9 +168,11 @@ func getAddressDetail(c *gin.Context) {
 	c.JSON(http.StatusOK, ad)
 }
 
-// postCreateAddress registers a new address with default quotas, idempotently.
-// It never modifies an address that already exists — use CSV sync to change
-// quotas or accepting_new on an existing address.
+// postCreateAddress registers an address, idempotently. A new address takes
+// the address table's column defaults except for any accepting_new and
+// quota_pool given. An existing address is left as it is except for any
+// accepting_new and quota_pool given; display_name and the limits are never
+// changed — use CSV sync for those.
 func postCreateAddress(c *gin.Context) {
 	ctx := c.Request.Context()
 
@@ -151,20 +187,122 @@ func postCreateAddress(c *gin.Context) {
 		c.AbortWithStatus(http.StatusBadRequest)
 		return
 	}
+	if req.QuotaPool != nil && *req.QuotaPool != "" && !isValidFmsgAddr(*req.QuotaPool) {
+		c.AbortWithStatus(http.StatusBadRequest)
+		return
+	}
 
-	addrLower := cases.Fold().String(req.Address)
-
-	tag, err := pool.Exec(ctx, sqlInsertAddressIfNotExists, addrLower, req.Address, req.DisplayName)
+	resp, err := registerAddress(ctx, req)
+	var qpErr quotaPoolError
+	if errors.As(err, &qpErr) {
+		c.AbortWithStatusJSON(http.StatusUnprocessableEntity, gin.H{"error": qpErr.Error()})
+		return
+	}
 	if err != nil {
 		c.AbortWithError(500, err)
 		return
 	}
 
-	if tag.RowsAffected() == 0 {
-		c.Status(http.StatusOK)
+	if resp.Created {
+		c.JSON(http.StatusCreated, resp)
 		return
 	}
-	c.Status(http.StatusCreated)
+	c.JSON(http.StatusOK, resp)
+}
+
+// registerAddress inserts or updates the requested address in one transaction.
+func registerAddress(ctx context.Context, req CreateAddressReq) (CreateAddressResp, error) {
+	var resp CreateAddressResp
+	addrLower := cases.Fold().String(req.Address)
+
+	// quotaPool is the case folded pool owner to set when req.QuotaPool is
+	// given, nil to take the address out of any pool
+	var quotaPool *string
+	if req.QuotaPool != nil && *req.QuotaPool != "" {
+		if p := cases.Fold().String(*req.QuotaPool); p != addrLower {
+			quotaPool = &p
+		}
+	}
+
+	tx, err := pool.Begin(ctx)
+	if err != nil {
+		return resp, err
+	}
+	defer tx.Rollback(ctx) //nolint:errcheck // a no-op after commit
+
+	// pools are one level deep: the pool owner must be in no other pool, and
+	// an address that owns a pool can't join another
+	if quotaPool != nil {
+		var ownerPool *string
+		err := tx.QueryRow(ctx, sqlSelectQuotaPoolForShare, *quotaPool).Scan(&ownerPool)
+		if errors.Is(err, pgx.ErrNoRows) {
+			return resp, quotaPoolError("quota_pool address not found")
+		}
+		if err != nil {
+			return resp, err
+		}
+		if ownerPool != nil && *ownerPool != *quotaPool {
+			return resp, quotaPoolError("quota_pool address is in another quota pool")
+		}
+		var isOwner bool
+		if err := tx.QueryRow(ctx, sqlIsPoolOwner, addrLower).Scan(&isOwner); err != nil {
+			return resp, err
+		}
+		if isOwner {
+			return resp, quotaPoolError("address owns a quota pool")
+		}
+	}
+
+	// only the columns given are inserted so the rest take their defaults
+	cols := []string{"address_lower", "address", "display_name"}
+	args := []any{addrLower, req.Address, req.DisplayName}
+	if req.AcceptingNew != nil {
+		cols = append(cols, "accepting_new")
+		args = append(args, *req.AcceptingNew)
+	}
+	if req.QuotaPool != nil {
+		cols = append(cols, "quota_pool")
+		args = append(args, quotaPool)
+	}
+	tag, err := tx.Exec(ctx, insertAddressSQL(cols), args...)
+	if err != nil {
+		return resp, err
+	}
+	resp.Created = tag.RowsAffected() == 1
+
+	if !resp.Created && (req.AcceptingNew != nil || req.QuotaPool != nil) {
+		var wasAccepting bool
+		if err := tx.QueryRow(ctx, sqlSelectAcceptingNewForUpdate, addrLower).Scan(&wasAccepting); err != nil {
+			return resp, err
+		}
+		var set []string
+		args := []any{addrLower}
+		if req.AcceptingNew != nil {
+			args = append(args, *req.AcceptingNew)
+			set = append(set, fmt.Sprintf("accepting_new = $%d", len(args)))
+		}
+		if req.QuotaPool != nil {
+			args = append(args, quotaPool)
+			set = append(set, fmt.Sprintf("quota_pool = $%d", len(args)))
+		}
+		if _, err := tx.Exec(ctx, "update address set "+strings.Join(set, ", ")+" where address_lower = $1;", args...); err != nil {
+			return resp, err
+		}
+		resp.Reinstated = !wasAccepting && req.AcceptingNew != nil && *req.AcceptingNew
+	}
+
+	return resp, tx.Commit(ctx)
+}
+
+// insertAddressSQL returns an insert of the given address columns, in order as
+// $1..$n, that does nothing when the address already exists.
+func insertAddressSQL(cols []string) string {
+	params := make([]string, len(cols))
+	for i := range cols {
+		params[i] = fmt.Sprintf("$%d", i+1)
+	}
+	return "insert into address (" + strings.Join(cols, ", ") + ") values (" +
+		strings.Join(params, ", ") + ") on conflict (address_lower) do nothing;"
 }
 
 func postAddressTxSend(c *gin.Context) {