From ac5bfa7ad7ffa47795c7fa775244a84b2290c852 Mon Sep 17 00:00:00 2001
From: magical <magical@tilde.town>
Date: Fri, 31 Dec 2021 22:24:58 +0000
Subject: [PATCH] limit response bodies

---
 mastodon.go | 62 ++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 61 insertions(+), 1 deletion(-)

diff --git a/mastodon.go b/mastodon.go
index 6d5befe..b3e9f78 100644
--- a/mastodon.go
+++ b/mastodon.go
@@ -3,6 +3,7 @@ package main
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -17,9 +18,12 @@ import (
 	"golang.org/x/net/html/atom"
 )
 
+const maxResponseSize = 1e6 // 1MB
+
 // https://docs.joinmastodon.org/spec/microformats/
 
 // Root elements (h-*)
+//
 // h-feed
 //     Represents a stream of entries. Attached to a profile's toots. Also
 //     attached to the parent thread within detailed status views.
@@ -127,8 +131,9 @@ func (src *MastoSource) update(ctx context.Context) {
 	}
 
 	// TODO: cache body
+	body := MaxBytesReader(resp.Body, maxResponseSize)
 
-	feed, err := parseMicroformats(resp.Body)
+	feed, err := parseMicroformats(body)
 	if err != nil {
 		err := fmt.Errorf("error parsing %q: %w", src.URL, err)
 		log.Println(err)
@@ -202,3 +207,58 @@ func text(s *goquery.Selection) string {
 
 	return buf.String()
 }
+
+// MaxBytesReader is similar to io.LimitReader but is intended for
+// limiting the size of incoming request bodies. In contrast to
+// io.LimitReader, MaxBytesReader's result is a ReadCloser, returns a
+// non-EOF error for a Read beyond the limit, and closes the
+// underlying reader when its Close method is called.
+//
+// MaxBytesReader prevents clients from accidentally or maliciously
+// sending a large request and wasting server resources.
+//
+// Based on http.MaxBytesReader
+func MaxBytesReader(r io.ReadCloser, n int64) io.ReadCloser {
+	if n < 0 { // Treat negative limits as equivalent to 0.
+		n = 0
+	}
+	return &maxBytesReader{r: r, n: n}
+}
+
+type maxBytesReader struct {
+	r   io.ReadCloser // underlying reader
+	n   int64         // max bytes remaining
+	err error         // sticky error
+}
+
+func (l *maxBytesReader) Read(p []byte) (n int, err error) {
+	if l.err != nil {
+		return 0, l.err
+	}
+	if len(p) == 0 {
+		return 0, nil
+	}
+	// If they asked for a 32KB byte read but only 5 bytes are
+	// remaining, no need to read 32KB. 6 bytes will answer the
+	// question of the whether we hit the limit or go past it.
+	if int64(len(p)) > l.n+1 {
+		p = p[:l.n+1]
+	}
+	n, err = l.r.Read(p)
+
+	if int64(n) <= l.n {
+		l.n -= int64(n)
+		l.err = err
+		return n, err
+	}
+
+	n = int(l.n)
+	l.n = 0
+
+	l.err = errors.New("http: response body too large")
+	return n, l.err
+}
+
+func (l *maxBytesReader) Close() error {
+	return l.r.Close()
+}