From 4f7d87a3e3bc56ef3a5c0549441b240857b3e6f9 Mon Sep 17 00:00:00 2001
From: Shiming Zhang <wzshiming@hotmail.com>
Date: Thu, 5 Sep 2024 19:01:32 +0800
Subject: [PATCH] Disable head method

---
 cmd/crproxy/main.go |  7 +++++++
 crproxy.go          | 37 +++++++++++++++++++------------------
 2 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/cmd/crproxy/main.go b/cmd/crproxy/main.go
index a567518..54401f8 100644
--- a/cmd/crproxy/main.go
+++ b/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)
diff --git a/crproxy.go b/crproxy.go
index fdfe337..844af7d 100644
--- a/crproxy.go
+++ b/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) {