aghio: create package

internal/aghio/limitedbody.go

@@ -0,0 +1,14 @@
+package aghio
+
+import (
+	"net/http"
+)
+
+// LimitRequestBody substitutes body of the request with LimitedReadCloser.
+func LimitRequestBody(h http.Handler) (limited http.Handler) {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		r.Body = LimitReadCloser(r.Body, 1024)
+
+		h.ServeHTTP(w, r)
+	})
+}

internal/aghio/limitedreader.go

@@ -0,0 +1,48 @@
+// Package aghio contains extensions for io package's types and methods
+package aghio
+
+import (
+	"fmt"
+	"io"
+
+	"github.com/AdguardTeam/AdGuardHome/internal/agherr"
+)
+
+// ErrLimitReached is returned if the limit of LimitedReader is reached.
+const ErrLimitReached agherr.Error = "limit reached"
+
+// limitedReadCloser is a wrapper for io.ReadCloser with limited reader and
+// dealing with agherr package.
+type limitedReadCloser struct {
+	limit int64
+	N     int64
+	io.ReadCloser
+}
+
+// Read implements Reader interface.
+func (lrc *limitedReadCloser) Read(p []byte) (n int, err error) {
+	if lrc.N <= 0 {
+		return 0, fmt.Errorf("read %d bytes: %w", lrc.limit, ErrLimitReached)
+	}
+	if int64(len(p)) > lrc.N {
+		p = p[0:lrc.N]
+	}
+	n, err = lrc.ReadCloser.Read(p)
+	lrc.N -= int64(n)
+	return n, err
+}
+
+// Close implements Closer interface.
+func (lrc *limitedReadCloser) Close() error {
+	return lrc.ReadCloser.Close()
+}
+
+// LimitReadCloser returns a ReadCloser with original Closer and Reader that
+// stops with ErrLimitReached after n bytes read.
+func LimitReadCloser(rc io.ReadCloser, n int64) io.ReadCloser {
+	return &limitedReadCloser{
+		limit:      n,
+		N:          n,
+		ReadCloser: rc,
+	}
+}

internal/aghio/limitedreader_test.go

@@ -0,0 +1,58 @@
+package aghio
+
+import (
+	"errors"
+	"io"
+	"io/ioutil"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLimitedReadCloser_Read(t *testing.T) {
+	testCases := []struct {
+		name  string
+		limit int64
+		rStr  string
+		want  int
+		err   error
+	}{{
+		name:  "perfectly_match",
+		limit: 3,
+		rStr:  "abc",
+		want:  3,
+		err:   nil,
+	}, {
+		name:  "eof",
+		limit: 3,
+		rStr:  "",
+		want:  0,
+		err:   io.EOF,
+	}, {
+		name:  "limit_reached",
+		limit: 0,
+		rStr:  "abc",
+		want:  0,
+		err:   ErrLimitReached,
+	}, {
+		name:  "truncated",
+		limit: 2,
+		rStr:  "abc",
+		want:  2,
+		err:   nil,
+	}}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			readCloser := ioutil.NopCloser(strings.NewReader(tc.rStr))
+			buf := make([]byte, tc.limit+1)
+
+			lreader := LimitReadCloser(readCloser, tc.limit)
+			n, err := lreader.Read(buf)
+
+			assert.Equal(t, n, tc.want)
+			assert.True(t, errors.Is(err, tc.err), buf)
+		})
+	}
+}