Repository: raviqqe/muffet
Pull request #369 (2024-02-28 10:32:28)

Checked HTTP client


Patch:
diff --git a/checked_http_client.go b/checked_http_client.go
new file mode 100644
index 0000000..ef4f86b
--- /dev/null
+++ b/checked_http_client.go
@@ -0,0 +1,27 @@
+package main
+
+import (
+	"fmt"
+	"net/http"
+	"net/url"
+)
+
+type checkedHttpClient struct {
+	client              httpClient
+	acceptedStatusCodes statusCodeSet
+}
+
+func newCheckedHttpClient(c httpClient, acceptedStatusCodes statusCodeSet) httpClient {
+	return &checkedHttpClient{c, acceptedStatusCodes}
+}
+
+func (c *checkedHttpClient) Get(u *url.URL, header http.Header) (httpResponse, error) {
+	r, err := c.client.Get(u, header)
+	if err != nil {
+		return nil, err
+	} else if code := r.StatusCode(); !c.acceptedStatusCodes.Contains(code) {
+		return nil, fmt.Errorf("%v", code)
+	}
+
+	return r, nil
+}
diff --git a/command.go b/command.go
index 0512b6e..bb5f8e6 100644
--- a/command.go
+++ b/command.go
@@ -44,24 +44,26 @@ func (c *command) runWithError(ss []string) (bool, error) {
 		return true, nil
 	}
 
-	client := newRedirectHttpClient(
-		newThrottledHttpClient(
-			c.httpClientFactory.Create(
-				httpClientOptions{
-					MaxConnectionsPerHost: args.MaxConnectionsPerHost,
-					MaxResponseBodySize:   args.MaxResponseBodySize,
-					BufferSize:            args.BufferSize,
-					Proxy:                 args.Proxy,
-					SkipTLSVerification:   args.SkipTLSVerification,
-					Timeout:               time.Duration(args.Timeout) * time.Second,
-					Header:                args.Header,
-				},
+	client := newCheckedHttpClient(
+		newRedirectHttpClient(
+			newThrottledHttpClient(
+				c.httpClientFactory.Create(
+					httpClientOptions{
+						MaxConnectionsPerHost: args.MaxConnectionsPerHost,
+						MaxResponseBodySize:   args.MaxResponseBodySize,
+						BufferSize:            args.BufferSize,
+						Proxy:                 args.Proxy,
+						SkipTLSVerification:   args.SkipTLSVerification,
+						Timeout:               time.Duration(args.Timeout) * time.Second,
+						Header:                args.Header,
+					},
+				),
+				args.RateLimit,
+				args.MaxConnections,
+				args.MaxConnectionsPerHost,
 			),
-			args.RateLimit,
-			args.MaxConnections,
-			args.MaxConnectionsPerHost,
+			args.MaxRedirections,
 		),
-		args.MaxRedirections,
 		args.AcceptedStatusCodes,
 	)
 
diff --git a/go.mod b/go.mod
index 933693a..839b9ac 100644
--- a/go.mod
+++ b/go.mod
@@ -1,6 +1,6 @@
 module github.com/raviqqe/muffet/v2
 
-go 1.19
+go 1.22.0
 
 require (
 	github.com/andybalholm/brotli v1.1.0
diff --git a/go.sum b/go.sum
index 3c921f6..3709a88 100644
--- a/go.sum
+++ b/go.sum
@@ -42,6 +42,7 @@ github.com/valyala/fasthttp v1.52.0/go.mod h1:hf5C4QnVMkNXMspnsUlfM3WitlgYflyhHY
 github.com/yhat/scrape v0.0.0-20161128144610-24b7890b0945 h1:6Ju8pZBYFTN9FaV/JvNBiIHcsgEmP4z4laciqjfjY8E=
 github.com/yhat/scrape v0.0.0-20161128144610-24b7890b0945/go.mod h1:4vRFPPNYllgCacoj+0FoKOjTW68rUhEfqPLiEJaK2w8=
 go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
+go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/ratelimit v0.3.0 h1:IdZd9wqvFXnvLvSEBo0KPcGfkoBGNkpTHlrE3Rcjkjw=
 go.uber.org/ratelimit v0.3.0/go.mod h1:So5LG7CV1zWpY1sHe+DXTJqQvOx+FFPFaAs2SnoyBaI=
 golang.org/x/net v0.21.0 h1:AQyQV4dYCvJ7vGmJyKki9+PBdyvhkSd8EIx/qb0AYv4=
diff --git a/redirect_http_client.go b/redirect_http_client.go
index 89151f0..5dbf8a2 100644
--- a/redirect_http_client.go
+++ b/redirect_http_client.go
@@ -9,13 +9,12 @@ import (
 )
 
 type redirectHttpClient struct {
-	client              httpClient
-	maxRedirections     int
-	acceptedStatusCodes statusCodeSet
+	client          httpClient
+	maxRedirections int
 }
 
-func newRedirectHttpClient(c httpClient, maxRedirections int, acceptedStatusCodes statusCodeSet) httpClient {
-	return &redirectHttpClient{c, maxRedirections, acceptedStatusCodes}
+func newRedirectHttpClient(c httpClient, maxRedirections int) httpClient {
+	return &redirectHttpClient{c, maxRedirections}
 }
 
 func (c *redirectHttpClient) Get(u *url.URL, header http.Header) (httpResponse, error) {
@@ -24,59 +23,39 @@ func (c *redirectHttpClient) Get(u *url.URL, header http.Header) (httpResponse,
 	}
 
 	cj, err := cookiejar.New(nil)
-
 	if err != nil {
 		return nil, err
 	}
 
-	i := 0
-
-	for {
+	for i := range c.maxRedirections + 1 {
 		for _, c := range cj.Cookies(u) {
 			header.Add("cookie", c.String())
 		}
 
 		r, err := c.client.Get(u, header)
-		if err != nil {
-			return nil, c.formatError(err, i, u)
-		}
-
-		code := r.StatusCode()
-
-		if c.acceptedStatusCodes.Contains(code) {
+		if err != nil && i == 0 {
+			return nil, err
+		} else if err != nil {
+			return nil, fmt.Errorf("%w (following redirect %v)", err, u.String())
+		} else if c := r.StatusCode(); c < 300 || c >= 400 {
 			return r, nil
-		} else if code >= 300 && code <= 399 {
-			i++
-
-			if i > c.maxRedirections {
-				return nil, errors.New("too many redirections")
-			}
-
-			s := r.Header("Location")
-
-			if len(s) == 0 {
-				return nil, errors.New("location header not set")
-			}
+		}
 
-			u, err = u.Parse(s)
+		s := r.Header("Location")
 
-			if err != nil {
-				return nil, err
-			}
+		if len(s) == 0 {
+			return nil, errors.New("location header not set")
+		}
 
-			cj.SetCookies(u, parseCookies(r.Header("set-cookie")))
-		} else {
-			return nil, c.formatError(fmt.Errorf("%v", code), i, u)
+		u, err = u.Parse(s)
+		if err != nil {
+			return nil, err
 		}
-	}
-}
 
-func (*redirectHttpClient) formatError(err error, redirections int, u *url.URL) error {
-	if redirections == 0 {
-		return err
+		cj.SetCookies(u, parseCookies(r.Header("set-cookie")))
 	}
 
-	return fmt.Errorf("%w (following redirect %v)", err, u.String())
+	return nil, errors.New("too many redirections")
 }
 
 func parseCookies(s string) []*http.Cookie {


Test patch:
diff --git a/checked_http_client_test.go b/checked_http_client_test.go
new file mode 100644
index 0000000..e3ba8d1
--- /dev/null
+++ b/checked_http_client_test.go
@@ -0,0 +1,44 @@
+package main
+
+import (
+	"net/url"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCheckedHttpClientFailWithValidStatusCode(t *testing.T) {
+	u, err := url.Parse(testUrl)
+
+	assert.Nil(t, err)
+
+	r, err := newCheckedHttpClient(
+		newFakeHttpClient(
+			func(u *url.URL) (*fakeHttpResponse, error) {
+				return newFakeHttpResponse(200, testUrl, nil, nil), nil
+			},
+		),
+		statusCodeSet{{200, 201}: {}},
+	).Get(u, nil)
+
+	assert.Nil(t, err)
+	assert.NotNil(t, r)
+}
+
+func TestCheckedHttpClientFailWithInvalidStatusCode(t *testing.T) {
+	u, err := url.Parse(testUrl)
+
+	assert.Nil(t, err)
+
+	r, err := newCheckedHttpClient(
+		newFakeHttpClient(
+			func(u *url.URL) (*fakeHttpResponse, error) {
+				return newFakeHttpResponse(404, testUrl, nil, nil), nil
+			},
+		),
+		statusCodeSet{{200, 201}: {}},
+	).Get(u, nil)
+
+	assert.Nil(t, r)
+	assert.Equal(t, err.Error(), "404")
+}
diff --git a/redirect_http_client_test.go b/redirect_http_client_test.go
index 2676759..487cf66 100644
--- a/redirect_http_client_test.go
+++ b/redirect_http_client_test.go
@@ -3,6 +3,7 @@ package main
 import (
 	"errors"
 	"net/url"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -10,10 +11,8 @@ import (
 
 const testUrl = "http://foo.com"
 
-var acceptedStatusCodes = statusCodeSet{{200, 300}: struct{}{}}
-
 func TestNewRedirectHttpClient(t *testing.T) {
-	newRedirectHttpClient(newFakeHttpClient(nil), 42, acceptedStatusCodes)
+	newRedirectHttpClient(newFakeHttpClient(nil), 42)
 }
 
 func TestRedirectHttpClientGet(t *testing.T) {
@@ -32,7 +31,28 @@ func TestRedirectHttpClientGet(t *testing.T) {
 			},
 		),
 		42,
-		acceptedStatusCodes,
+	).Get(u, nil)
+
+	assert.Nil(t, err)
+	assert.Equal(t, 200, r.StatusCode())
+}
+
+func TestRedirectHttpClientGetWithoutRedirect(t *testing.T) {
+	u, err := url.Parse(testUrl)
+
+	assert.Nil(t, err)
+
+	r, err := newRedirectHttpClient(
+		newFakeHttpClient(
+			func(u *url.URL) (*fakeHttpResponse, error) {
+				if u.String() != testUrl {
+					return nil, errors.New("")
+				}
+
+				return newFakeHtmlResponse(testUrl, ""), nil
+			},
+		),
+		0,
 	).Get(u, nil)
 
 	assert.Nil(t, err)
@@ -65,7 +85,6 @@ func TestRedirectHttpClientGetWithRedirect(t *testing.T) {
 			},
 		),
 		42,
-		acceptedStatusCodes,
 	).Get(u, nil)
 
 	assert.Nil(t, err)
@@ -100,7 +119,6 @@ func TestRedirectHttpClientGetWithRedirects(t *testing.T) {
 			},
 		),
 		maxRedirections,
-		acceptedStatusCodes,
 	).Get(u, nil)
 
 	assert.Nil(t, err)
@@ -139,7 +157,6 @@ func TestRedirectHttpClientGetWithRelativeRedirect(t *testing.T) {
 			},
 		),
 		maxRedirections,
-		acceptedStatusCodes,
 	).Get(u, nil)
 
 	assert.Nil(t, err)
@@ -148,33 +165,34 @@ func TestRedirectHttpClientGetWithRelativeRedirect(t *testing.T) {
 }
 
 func TestRedirectHttpClientFailWithTooManyRedirects(t *testing.T) {
-	const maxRedirections = 42
+	for _, n := range []int{0, 1, 42} {
+		t.Run(strconv.Itoa(n), func(t *testing.T) {
+			u, err := url.Parse(testUrl)
 
-	u, err := url.Parse(testUrl)
-
-	assert.Nil(t, err)
-
-	i := 0
-	r, err := newRedirectHttpClient(
-		newFakeHttpClient(
-			func(u *url.URL) (*fakeHttpResponse, error) {
-				i++
+			assert.Nil(t, err)
 
-				return newFakeHttpResponse(
-					300,
-					testUrl,
-					nil,
-					map[string]string{"location": testUrl},
-				), nil
-			},
-		),
-		maxRedirections,
-		acceptedStatusCodes,
-	).Get(u, nil)
+			i := 0
+			r, err := newRedirectHttpClient(
+				newFakeHttpClient(
+					func(u *url.URL) (*fakeHttpResponse, error) {
+						i++
 
-	assert.Nil(t, r)
-	assert.Equal(t, err.Error(), "too many redirections")
-	assert.Equal(t, maxRedirections+1, i)
+						return newFakeHttpResponse(
+							300,
+							testUrl,
+							nil,
+							map[string]string{"location": testUrl},
+						), nil
+					},
+				),
+				n,
+			).Get(u, nil)
+
+			assert.Nil(t, r)
+			assert.Equal(t, err.Error(), "too many redirections")
+			assert.Equal(t, n+1, i)
+		})
+	}
 }
 
 func TestRedirectHttpClientFailWithUnsetLocationHeader(t *testing.T) {
@@ -189,7 +207,6 @@ func TestRedirectHttpClientFailWithUnsetLocationHeader(t *testing.T) {
 			},
 		),
 		42,
-		acceptedStatusCodes,
 	).Get(u, nil)
 
 	assert.Nil(t, r)
@@ -213,32 +230,12 @@ func TestRedirectHttpClientFailWithInvalidLocationURL(t *testing.T) {
 			},
 		),
 		42,
-		acceptedStatusCodes,
 	).Get(u, nil)
 
 	assert.Nil(t, r)
 	assert.Contains(t, err.Error(), "parse")
 }
 
-func TestRedirectHttpClientFailWithInvalidStatusCode(t *testing.T) {
-	u, err := url.Parse(testUrl)
-
-	assert.Nil(t, err)
-
-	r, err := newRedirectHttpClient(
-		newFakeHttpClient(
-			func(u *url.URL) (*fakeHttpResponse, error) {
-				return newFakeHttpResponse(404, testUrl, nil, nil), nil
-			},
-		),
-		42,
-		acceptedStatusCodes,
-	).Get(u, nil)
-
-	assert.Nil(t, r)
-	assert.Equal(t, err.Error(), "404")
-}
-
 func TestRedirectHttpClientFailAfterRedirect(t *testing.T) {
 	u, err := url.Parse(testUrl)
 
@@ -259,11 +256,10 @@ func TestRedirectHttpClientFailAfterRedirect(t *testing.T) {
 					), nil
 				}
 
-				return newFakeHttpResponse(404, "", nil, nil), nil
+				return nil, errors.New("foo")
 			},
 		),
 		42,
-		acceptedStatusCodes,
 	).Get(u, nil)
 
 	assert.Nil(t, r)