Count attachments and use the accept time when recording receive usage

README.md

@@ -42,6 +42,10 @@ Requires Go 1.27 or newer.
 
 
 
+## Quotas and usage
+
+For each local recipient of a message from another host, fmsgd asks fmsgid for the address's receive limits and usage, and rejects the recipient with `101` (user full) if one more message of this size would exceed its messages per day, size per day or total size (`-1` is unlimited). The size is the stored size: the message body plus every attachment, decompressed. Once the message is stored, fmsgd records that size for each recipient that accepted it with `POST /fmsgid/recv`, timestamped when this host accepted the message (not the sender's timestamp), so a backdated message still counts towards today. fmsgd doesn't record send usage or mail between local addresses: the client that sends and delivers locally (such as fmsg-webapi) does.
+
 ## Running
 
 An up and running [fmsg Id API](https://github.com/markmnl/fmsgid) needs to be reachable by fmsgd to know users and their quotas for this fmsgd service. See also [fmsg-docker](https://github.com/markmnl/fmsg-docker) - a docker compose stack for a fmsg host including fmsgid, fmsg-webpi and fmsgd.

cmd/fmsgd/host.go

@@ -1253,21 +1253,48 @@ func validateMsgRecvForAddr(h *FMsgHeader, addr *FMsgAddress, msgHash []byte) (c
 		return RejectCodeUserNotAccepting, nil
 	}
 
-	// check user limits
-	if detail.LimitRecvCountPer1d > -1 && detail.RecvCountPer1d+1 > detail.LimitRecvCountPer1d {
-		log.Printf("WARN: Message rejected: RecvCountPer1d would exceed LimitRecvCountPer1d %d", detail.LimitRecvCountPer1d)
+	// check user limits against the size that will be recorded once accepted
+	if exceeded := recvLimitExceeded(detail, storedSize(h)); exceeded != "" {
+		log.Printf("WARN: Message rejected: %s", exceeded)
 		return RejectCodeUserFull, nil
 	}
-	if detail.LimitRecvSizePer1d > -1 && detail.RecvSizePer1d+int64(h.Size) > detail.LimitRecvSizePer1d {
-		log.Printf("WARN: Message rejected: RecvSizePer1d would exceed LimitRecvSizePer1d %d", detail.LimitRecvSizePer1d)
-		return RejectCodeUserFull, nil
+
+	return RejectCodeAccept, nil
+}
+
+// storedSize is the size a received message takes in storage, and so the size
+// recorded as the recipient's usage: the message body plus every attachment,
+// each expanded (stored decompressed) when sent compressed. Limits are checked
+// against this same size before the message is accepted.
+func storedSize(h *FMsgHeader) int64 {
+	size := int64(h.Size)
+	if h.Flags&FlagDeflate != 0 {
+		size = int64(h.ExpandedSize)
 	}
-	if detail.LimitRecvSizeTotal > -1 && detail.RecvSizeTotal+int64(h.Size) > detail.LimitRecvSizeTotal {
-		log.Printf("WARN: Message rejected: RecvSizeTotal would exceed LimitRecvSizeTotal %d", detail.LimitRecvSizeTotal)
-		return RejectCodeUserFull, nil
+	for _, a := range h.Attachments {
+		if a.Flags&(1<<1) != 0 {
+			size += int64(a.ExpandedSize)
+		} else {
+			size += int64(a.Size)
+		}
 	}
+	return size
+}
 
-	return RejectCodeAccept, nil
+// recvLimitExceeded checks one more message of size bytes against an
+// address's receive limits (a limit of -1 is unlimited). It returns which
+// limit would be exceeded, or "" when the message fits.
+func recvLimitExceeded(detail *AddressDetail, size int64) string {
+	if detail.LimitRecvCountPer1d > -1 && detail.RecvCountPer1d+1 > detail.LimitRecvCountPer1d {
+		return fmt.Sprintf("RecvCountPer1d would exceed LimitRecvCountPer1d %d", detail.LimitRecvCountPer1d)
+	}
+	if detail.LimitRecvSizePer1d > -1 && detail.RecvSizePer1d+size > detail.LimitRecvSizePer1d {
+		return fmt.Sprintf("RecvSizePer1d would exceed LimitRecvSizePer1d %d", detail.LimitRecvSizePer1d)
+	}
+	if detail.LimitRecvSizeTotal > -1 && detail.RecvSizeTotal+size > detail.LimitRecvSizeTotal {
+		return fmt.Sprintf("RecvSizeTotal would exceed LimitRecvSizeTotal %d", detail.LimitRecvSizeTotal)
+	}
+	return ""
 }
 
 // uniqueFilepath generates a unique file path in the given directory,
@@ -1515,13 +1542,22 @@ func storeAcceptedMessage(h *FMsgHeader, codes []byte, acceptedTo []FMsgAddress,
 		return false
 	}
 
-	allAccepted := append(acceptedTo, acceptedAddTo...)
-	for i := range allAccepted {
-		if err := postMsgStatRecv(&allAccepted[i], h.Timestamp, int(h.Size)); err != nil {
-			log.Printf("WARN: Failed to post msg recv stat: %s", err)
+	recordRecvUsage(h, append(acceptedTo, acceptedAddTo...), timeutil.TimestampNow().Float64())
+	return true
+}
+
+// recordRecvUsage records a stored message against each recipient that
+// accepted it: its stored size (body plus attachments), at now — the time this
+// host accepted it, not the sender's timestamp, so a backdated message still
+// counts towards today's limits. A failure is logged: the message is already
+// stored.
+func recordRecvUsage(h *FMsgHeader, accepted []FMsgAddress, now float64) {
+	size := storedSize(h)
+	for i := range accepted {
+		if err := postMsgStatRecv(&accepted[i], now, size); err != nil {
+			log.Printf("WARN: Failed to post msg recv stat for %s: %s", accepted[i].ToString(), err)
 		}
 	}
-	return true
 }
 
 func downloadMessage(c net.Conn, r io.Reader, h *FMsgHeader, skipData bool) error {

cmd/fmsgd/id.go

@@ -52,15 +52,15 @@ func getAddressDetail(addr *FMsgAddress) (*AddressDetail, error) {
 	return &detail, nil
 }
 
-func postMsgStatSend(addr *FMsgAddress, timestamp float64, size int) error {
+func postMsgStatSend(addr *FMsgAddress, timestamp float64, size int64) error {
 	return postMsgStat(addr, timestamp, size, true)
 }
 
-func postMsgStatRecv(addr *FMsgAddress, timestamp float64, size int) error {
+func postMsgStatRecv(addr *FMsgAddress, timestamp float64, size int64) error {
 	return postMsgStat(addr, timestamp, size, false)
 }
 
-func postMsgStat(addr *FMsgAddress, timestamp float64, size int, isSending bool) error {
+func postMsgStat(addr *FMsgAddress, timestamp float64, size int64, isSending bool) error {
 	var part string
 	if isSending {
 		part = "send"

cmd/fmsgd/usage_test.go

@@ -0,0 +1,130 @@
+package main
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+)
+
+func TestStoredSize(t *testing.T) {
+	tests := []struct {
+		name string
+		h    FMsgHeader
+		want int64
+	}{
+		{name: "body only", h: FMsgHeader{Size: 100}, want: 100},
+		{name: "compressed body counts expanded", h: FMsgHeader{Flags: FlagDeflate, Size: 40, ExpandedSize: 400}, want: 400},
+		{
+			name: "attachments included",
+			h: FMsgHeader{Size: 100, Attachments: []FMsgAttachmentHeader{
+				{Size: 1000},
+				{Flags: 1 << 1, Size: 50, ExpandedSize: 5000},
+			}},
+			want: 100 + 1000 + 5000,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := storedSize(&tt.h); got != tt.want {
+				t.Fatalf("storedSize = %d, want %d", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestRecvLimitExceeded(t *testing.T) {
+	unlimited := AddressDetail{
+		LimitRecvSizeTotal: -1, LimitRecvSizePer1d: -1, LimitRecvCountPer1d: -1,
+		RecvSizeTotal: 1 << 40, RecvSizePer1d: 1 << 40, RecvCountPer1d: 1 << 40,
+	}
+	tests := []struct {
+		name   string
+		detail AddressDetail
+		size   int64
+		want   string // substring of the reason, "" when within limits
+	}{
+		{name: "unlimited", detail: unlimited, size: 1 << 30},
+		{name: "count fits", detail: AddressDetail{LimitRecvCountPer1d: 10, RecvCountPer1d: 9, LimitRecvSizePer1d: -1, LimitRecvSizeTotal: -1}, size: 1},
+		{name: "count full", detail: AddressDetail{LimitRecvCountPer1d: 10, RecvCountPer1d: 10, LimitRecvSizePer1d: -1, LimitRecvSizeTotal: -1}, size: 1, want: "RecvCountPer1d"},
+		{name: "day size fits exactly", detail: AddressDetail{LimitRecvCountPer1d: -1, LimitRecvSizePer1d: 100, RecvSizePer1d: 60, LimitRecvSizeTotal: -1}, size: 40},
+		{name: "day size over", detail: AddressDetail{LimitRecvCountPer1d: -1, LimitRecvSizePer1d: 100, RecvSizePer1d: 60, LimitRecvSizeTotal: -1}, size: 41, want: "RecvSizePer1d"},
+		{name: "total over", detail: AddressDetail{LimitRecvCountPer1d: -1, LimitRecvSizePer1d: -1, LimitRecvSizeTotal: 1000, RecvSizeTotal: 999}, size: 2, want: "RecvSizeTotal"},
+		{name: "zero limit refuses", detail: AddressDetail{LimitRecvCountPer1d: 0, LimitRecvSizePer1d: -1, LimitRecvSizeTotal: -1}, size: 0, want: "RecvCountPer1d"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := recvLimitExceeded(&tt.detail, tt.size)
+			if tt.want == "" && got != "" {
+				t.Fatalf("got %q, want within limits", got)
+			}
+			if tt.want != "" && !strings.Contains(got, tt.want) {
+				t.Fatalf("got %q, want %q", got, tt.want)
+			}
+		})
+	}
+}
+
+// TestRecvLimitIncludesAttachments shows a message whose body fits but whose
+// attachments take it over the limit is refused: limits are checked against
+// the stored size, the same size that is recorded.
+func TestRecvLimitIncludesAttachments(t *testing.T) {
+	h := FMsgHeader{Size: 10, Attachments: []FMsgAttachmentHeader{{Size: 500}}}
+	detail := AddressDetail{LimitRecvCountPer1d: -1, LimitRecvSizePer1d: 100, LimitRecvSizeTotal: -1}
+	if recvLimitExceeded(&detail, int64(h.Size)) != "" {
+		t.Fatal("body alone should fit")
+	}
+	if recvLimitExceeded(&detail, storedSize(&h)) == "" {
+		t.Fatal("body plus attachments should exceed the daily size limit")
+	}
+}
+
+func TestRecordRecvUsage(t *testing.T) {
+	type tx struct {
+		Address string  `json:"address"`
+		Ts      float64 `json:"ts"`
+		Size    int64   `json:"size"`
+	}
+	var mu sync.Mutex
+	var got []tx
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost || r.URL.Path != "/fmsgid/recv" {
+			t.Errorf("unexpected %s %s", r.Method, r.URL.Path)
+			w.WriteHeader(http.StatusNotFound)
+			return
+		}
+		var v tx
+		if err := json.NewDecoder(r.Body).Decode(&v); err != nil {
+			t.Errorf("decode: %v", err)
+		}
+		mu.Lock()
+		got = append(got, v)
+		mu.Unlock()
+	}))
+	defer srv.Close()
+	origURI := IDURI
+	IDURI = srv.URL
+	defer func() { IDURI = origURI }()
+
+	// The sender's timestamp is a day old; usage must use the accept time.
+	h := FMsgHeader{
+		Timestamp:    1_000,
+		Flags:        FlagDeflate,
+		Size:         20,
+		ExpandedSize: 200,
+		Attachments:  []FMsgAttachmentHeader{{Size: 30}},
+	}
+	accepted := []FMsgAddress{{User: "alice", Domain: "example.com"}, {User: "bob", Domain: "example.com"}}
+	recordRecvUsage(&h, accepted, 90_000)
+
+	if len(got) != 2 {
+		t.Fatalf("posted %d recv stats, want 2", len(got))
+	}
+	for i, v := range got {
+		if v.Address != accepted[i].ToString() || v.Ts != 90_000 || v.Size != 230 {
+			t.Fatalf("stat %d = %+v, want %s at 90000 size 230", i, v, accepted[i].ToString())
+		}
+	}
+}