Support limit image size

cmd/crproxy/main.go

@@ -31,6 +31,7 @@ var (
 	userpass             []string
 	disableKeepAlives    []string
 	limitDelay           bool
+	blockImageSize       string
 	blobsSpeedLimit      string
 	ipsSpeedLimit        string
 	totalBlobsSpeedLimit string
@@ -50,6 +51,7 @@ func init() {
 	pflag.StringVarP(&address, "address", "a", ":8080", "listen on the address")
 	pflag.StringSliceVar(&disableKeepAlives, "disable-keep-alives", nil, "disable keep alives for the host")
 	pflag.BoolVar(&limitDelay, "limit-delay", false, "limit with delay")
+	pflag.StringVar(&blockImageSize, "block-image-size", "", "block image size")
 	pflag.StringVar(&blobsSpeedLimit, "blobs-speed-limit", "", "blobs speed limit per second (default unlimited)")
 	pflag.StringVar(&ipsSpeedLimit, "ips-speed-limit", "", "ips speed limit per second (default unlimited)")
 	pflag.StringVar(&totalBlobsSpeedLimit, "total-blobs-speed-limit", "", "total blobs speed limit per second (default unlimited)")
@@ -190,6 +192,15 @@ func main() {
 		opts = append(opts, crproxy.WithBlobsSpeedLimit(b, d))
 	}
 
+	if blockImageSize != "" {
+		b, err := geario.FromBytesSize(blockImageSize)
+		if err != nil {
+			logger.Println("failed to FromBytesSize:", err)
+			os.Exit(1)
+		}
+		opts = append(opts, crproxy.WithBlockImageSize(b))
+	}
+
 	if totalBlobsSpeedLimit != "" {
 		b, err := geario.FromBytesSize(totalBlobsSpeedLimit)
 		if err != nil {

crproxy.go

@@ -1,6 +1,7 @@
 package crproxy
 
 import (
+	"bytes"
 	"context"
 	"crypto/sha256"
 	"encoding/hex"
@@ -16,6 +17,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/distribution/distribution/v3/manifest/manifestlist"
 	"github.com/distribution/distribution/v3/registry/api/errcode"
 	"github.com/distribution/distribution/v3/registry/client/auth"
 	"github.com/distribution/distribution/v3/registry/client/auth/challenge"
@@ -57,6 +59,7 @@ type CRProxy struct {
 	ipsSpeedLimit           *geario.B
 	ipsSpeedLimitDuration   time.Duration
 	blockFunc               func(*PathInfo) bool
+	blockImageSize          *geario.B
 	retry                   int
 	retryInterval           time.Duration
 	storageDriver           storagedriver.StorageDriver
@@ -86,6 +89,12 @@ func WithLimitDelay(b bool) Option {
 	}
 }
 
+func WithBlockImageSize(limit geario.B) Option {
+	return func(c *CRProxy) {
+		c.blockImageSize = &limit
+	}
+}
+
 func WithLinkExpires(d time.Duration) Option {
 	return func(c *CRProxy) {
 		c.linkExpires = d
@@ -465,6 +474,29 @@ func (c *CRProxy) directResponse(rw http.ResponseWriter, r *http.Request, info *
 		}
 	}
 
+	if r.Method != http.MethodHead && !c.isPrivileged(r.RemoteAddr) {
+		if c.blockImageSize != nil && info.Manifests != "" {
+			ct := resp.Header.Get("Content-Type")
+			content, err := io.ReadAll(resp.Body)
+			resp.Body.Close()
+			resp.Body = io.NopCloser(bytes.NewBuffer(content))
+			if err == nil {
+				manifest, _, err := UnmarshalManifest(ct, content)
+				if err == nil {
+					if _, ok := manifest.(*manifestlist.DeserializedManifestList); !ok {
+						var sum int64
+						for _, refer := range manifest.References() {
+							sum += refer.Size
+						}
+						if int64(*c.blockImageSize) < sum {
+							errcode.ServeJSON(rw, errcode.ErrorCodeDenied)
+							return
+						}
+					}
+				}
+			}
+		}
+	}
 	header := rw.Header()
 	for k, v := range resp.Header {
 		key := textproto.CanonicalMIMEHeaderKey(k)
@@ -523,13 +555,14 @@ func (c *CRProxy) cacheBlobResponse(rw http.ResponseWriter, r *http.Request, inf
 
 	stat, err := c.storageDriver.Stat(ctx, blobPath)
 	if err == nil {
+		size := stat.Size()
 		if r.Method == http.MethodHead {
-			rw.Header().Set("Content-Length", strconv.FormatInt(stat.Size(), 10))
+			rw.Header().Set("Content-Length", strconv.FormatInt(size, 10))
 			rw.Header().Set("Content-Type", "application/octet-stream")
 			doneCache()
 			return
 		}
-		c.accumulativeLimit(rw, r, info, stat.Size())
+		c.accumulativeLimit(rw, r, info, size)
 		err = c.redirect(rw, r, blobPath)
 		if err == nil {
 			doneCache()

go.mod

@@ -26,6 +26,7 @@ require (
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/docker/go-metrics v0.0.1 // indirect
+	github.com/docker/libtrust v0.0.0-20150114040149-fa567046d9b1 // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect
 	github.com/gofrs/uuid v4.0.0+incompatible // indirect
 	github.com/golang-jwt/jwt/v4 v4.5.0 // indirect

manifest.go

@@ -0,0 +1,14 @@
+package crproxy
+
+import (
+	"github.com/distribution/distribution/v3"
+
+	_ "github.com/distribution/distribution/v3/manifest/manifestlist"
+	_ "github.com/distribution/distribution/v3/manifest/ocischema"
+	_ "github.com/distribution/distribution/v3/manifest/schema1"
+	_ "github.com/distribution/distribution/v3/manifest/schema2"
+)
+
+func UnmarshalManifest(ctHeader string, p []byte) (distribution.Manifest, distribution.Descriptor, error) {
+	return distribution.UnmarshalManifest(ctHeader, p)
+}