Add block-image-list

cmd/crproxy/main.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"slices"
 	"strings"
 
 	"github.com/gorilla/handlers"
@@ -22,6 +23,7 @@ var (
 	disableKeepAlives    []string
 	blobsSpeedLimit      string
 	totalBlobsSpeedLimit string
+	blockImageList       []string
 )
 
 func init() {
@@ -30,6 +32,7 @@ func init() {
 	pflag.StringSliceVar(&disableKeepAlives, "disable-keep-alives", nil, "disable keep alives for the host")
 	pflag.StringVar(&blobsSpeedLimit, "blobs-speed-limit", "", "blobs speed limit per second (default unlimited)")
 	pflag.StringVar(&totalBlobsSpeedLimit, "total-blobs-speed-limit", "", "total blobs speed limit per second (default unlimited)")
+	pflag.StringSliceVar(&blockImageList, "block-image-list", nil, "block image list")
 	pflag.Parse()
 }
 
@@ -80,6 +83,10 @@ func main() {
 		crproxy.WithBaseClient(cli),
 		crproxy.WithLogger(logger),
 		crproxy.WithMaxClientSizeForEachRegistry(16),
+		crproxy.WithBlockFunc(func(info *crproxy.PathInfo) bool {
+			image := info.Host + "/" + info.Image
+			return slices.Contains(blockImageList, image)
+		}),
 		crproxy.WithDomainAlias(map[string]string{
 			"docker.io": "registry-1.docker.io",
 			"ollama.ai": "registry.ollama.ai",

crproxy.go

@@ -43,6 +43,7 @@ type CRProxy struct {
 	logger                  Logger
 	totalBlobsSpeedLimit    *geario.Gear
 	blobsSpeedLimit         *geario.B
+	blockFunc               func(*PathInfo) bool
 }
 
 type Option func(c *CRProxy)
@@ -104,6 +105,12 @@ func WithDisableKeepAlives(disableKeepAlives []string) Option {
 	}
 }
 
+func WithBlockFunc(blockFunc func(info *PathInfo) bool) Option {
+	return func(c *CRProxy) {
+		c.blockFunc = blockFunc
+	}
+}
+
 func NewCRProxy(opts ...Option) (*CRProxy, error) {
 	c := &CRProxy{
 		challengeManager: challenge.NewSimpleManager(),
@@ -316,6 +323,11 @@ func (c *CRProxy) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 		info = c.modify(info)
 	}
 
+	if c.blockFunc != nil && c.blockFunc(info) {
+		errcode.ServeJSON(rw, errcode.ErrorCodeDenied)
+		return
+	}
+
 	path, err := info.Path()
 	if err != nil {
 		if c.logger != nil {