Disable head method

cmd/crproxy/main.go

@@ -72,6 +72,8 @@ var (
 	enableInternalAPI bool
 
 	readmeURL string
+
+	allowHeadMethod bool
 )
 
 func init() {
@@ -114,6 +116,7 @@ func init() {
 	pflag.BoolVar(&enableInternalAPI, "enable-internal-api", false, "enable internal api")
 
 	pflag.StringVar(&readmeURL, "readme-url", "", "redirect readme url when not found")
+	pflag.BoolVar(&allowHeadMethod, "allow-head-method", false, "allow head method")
 	pflag.Parse()
 }
 
@@ -440,6 +443,10 @@ func main() {
 		}))
 	}
 
+	if allowHeadMethod {
+		opts = append(opts, crproxy.WithAllowHeadMethod(allowHeadMethod))
+	}
+
 	crp, err := crproxy.NewCRProxy(opts...)
 	if err != nil {
 		logger.Println("failed to NewCRProxy:", err)

crproxy.go

@@ -12,6 +12,7 @@ import (
 	"strings"
 	"sync"
 	"time"
+	"bytes"
 
 	"github.com/daocloud/crproxy/internal/maps"
 	"github.com/distribution/distribution/v3/registry/api/errcode"
@@ -81,6 +82,7 @@ type CRProxy struct {
 
 	privilegedFunc           func(r *http.Request, info *ImageInfo) bool
 	redirectToOriginBlobFunc func(r *http.Request, info *ImageInfo) bool
+	allowHeadMethod          bool
 }
 
 type Option func(c *CRProxy)
@@ -243,6 +245,12 @@ func WithRetry(retry int, retryInterval time.Duration) Option {
 	}
 }
 
+func WithAllowHeadMethod(allowHeadMethod bool) Option {
+	return func(c *CRProxy) {
+		c.allowHeadMethod = allowHeadMethod
+	}
+}
+
 func NewCRProxy(opts ...Option) (*CRProxy, error) {
 	c := &CRProxy{
 		challengeManager: challenge.NewSimpleManager(),
@@ -419,25 +427,18 @@ func emptyTagsList(w http.ResponseWriter, r *http.Request) {
 	fmt.Fprint(w, emptyTagsList)
 }
 
-func (c *CRProxy) do(cli *http.Client, r *http.Request) (*http.Response, error) {
-	resp, err := cli.Do(r)
-	if err == nil {
-		return resp, nil
-	}
-
-	if r.Method != http.MethodHead {
-		return nil, err
-	}
+var emptyBody = io.NopCloser(bytes.NewBuffer([]byte{}))
 
-	r.Method = http.MethodGet
-	defer func() {
-		r.Method = http.MethodHead
-	}()
-	resp0, err0 := cli.Do(r)
-	if err0 != nil {
-		return nil, err
-	}
-	return resp0, nil
+func (c *CRProxy) do(cli *http.Client, r *http.Request) (*http.Response, error) {
+	if !c.allowHeadMethod && r.Method == http.MethodHead {
+		r.Method = http.MethodGet
+		defer func() {
+			r.Method = http.MethodHead
+			_ = r.Body.Close()
+			r.Body = emptyBody
+		}()
+	}
+	return cli.Do(r)
 }
 
 func (c *CRProxy) doWithAuth(cli *http.Client, r *http.Request, host string) (*http.Response, error) {