Add allow list and block message

DaoCloud · Jun 20, 2024 · 2e3b0ed · 2e3b0ed
1 parent 72bca9f
commit 2e3b0ed
Show file tree

Hide file tree

Showing 5 changed files with 126 additions and 63 deletions.
diff --git a/cmd/crproxy/main.go b/cmd/crproxy/main.go
@@ -1,8 +1,11 @@
 package main
 
 import (
+	"bufio"
+	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
@@ -16,6 +19,7 @@ import (
 	"github.com/gorilla/handlers"
 	"github.com/spf13/pflag"
 	"github.com/wzshiming/geario"
+	"github.com/wzshiming/hostmatcher"
 
 	_ "github.com/daocloud/crproxy/storage/driver/oss"
 	_ "github.com/distribution/distribution/v3/registry/storage/driver/azure"
@@ -26,27 +30,29 @@ import (
 )
 
 var (
-	behind               bool
-	address              string
-	userpass             []string
-	disableKeepAlives    []string
-	limitDelay           bool
-	blobsSpeedLimit      string
-	ipsSpeedLimit        string
-	totalBlobsSpeedLimit string
-	allowHostList        []string
-	blockImageList       []string
-	privilegedIPList     []string
-	retry                int
-	retryInterval        time.Duration
-	storageDriver        string
-	storageParameters    map[string]string
-	linkExpires          time.Duration
-	redirectLinks        string
-	disableTagsList      bool
-	enablePprof          bool
-	defaultRegistry      string
-	simpleAuth           bool
+	behind                 bool
+	address                string
+	userpass               []string
+	disableKeepAlives      []string
+	limitDelay             bool
+	blobsSpeedLimit        string
+	ipsSpeedLimit          string
+	totalBlobsSpeedLimit   string
+	allowHostList          []string
+	allowImageListFromFile string
+	blockImageList         []string
+	blockMessage           string
+	privilegedIPList       []string
+	retry                  int
+	retryInterval          time.Duration
+	storageDriver          string
+	storageParameters      map[string]string
+	linkExpires            time.Duration
+	redirectLinks          string
+	disableTagsList        bool
+	enablePprof            bool
+	defaultRegistry        string
+	simpleAuth             bool
 )
 
 func init() {
@@ -59,7 +65,9 @@ func init() {
 	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)")
 	pflag.StringSliceVar(&allowHostList, "allow-host-list", nil, "allow host list")
+	pflag.StringVar(&allowImageListFromFile, "allow-image-list-from-file", "", "allow image list from file")
 	pflag.StringSliceVar(&blockImageList, "block-image-list", nil, "block image list")
+	pflag.StringVar(&blockMessage, "block-message", "", "block message")
 	pflag.StringSliceVar(&privilegedIPList, "privileged-ip-list", nil, "privileged IP list")
 	pflag.IntVar(&retry, "retry", 0, "retry times")
 	pflag.DurationVar(&retryInterval, "retry-interval", 0, "retry interval")
@@ -125,13 +133,13 @@ func main() {
 			"docker.io": "registry-1.docker.io",
 			"ollama.ai": "registry.ollama.ai",
 		}),
-		crproxy.WithPathInfoModifyFunc(func(info *crproxy.ModifyInfo) *crproxy.ModifyInfo {
+		crproxy.WithPathInfoModifyFunc(func(info *crproxy.ImageInfo) *crproxy.ImageInfo {
 			// docker.io/busybox => docker.io/library/busybox
-			if info.Host == "registry-1.docker.io" && !strings.Contains(info.Image, "/") {
-				info.Image = "library/" + info.Image
+			if info.Host == "registry-1.docker.io" && !strings.Contains(info.Name, "/") {
+				info.Name = "library/" + info.Name
 			}
-			if info.Host == "registry.ollama.ai" && !strings.Contains(info.Image, "/") {
-				info.Image = "library/" + info.Image
+			if info.Host == "registry.ollama.ai" && !strings.Contains(info.Name, "/") {
+				info.Name = "library/" + info.Name
 			}
 			return info
 		}),
@@ -162,7 +170,31 @@ func main() {
 		}
 	}
 
-	if len(blockImageList) != 0 || len(allowHostList) != 0 {
+	if allowImageListFromFile != "" {
+		f, err := os.ReadFile(allowImageListFromFile)
+		if err != nil {
+			logger.Println("can't read allow list file %s", allowImageListFromFile)
+			os.Exit(1)
+		}
+
+		lines := bufio.NewReader(bytes.NewReader(f))
+		hosts := []string{}
+		for {
+			line, _, err := lines.ReadLine()
+			if err == io.EOF {
+				break
+			}
+			h := strings.TrimSpace(string(line))
+			if len(h) == 0 {
+				continue
+			}
+			hosts = append(hosts, h)
+		}
+		matcher := hostmatcher.NewMatcher(hosts)
+		opts = append(opts, crproxy.WithBlockFunc(func(info *crproxy.ImageInfo) bool {
+			return !matcher.Match(info.Host + "/" + info.Name)
+		}))
+	} else if len(blockImageList) != 0 || len(allowHostList) != 0 {
 		allowHostMap := map[string]struct{}{}
 		for _, host := range allowHostList {
 			allowHostMap[host] = struct{}{}
@@ -171,7 +203,7 @@ func main() {
 		for _, image := range blockImageList {
 			blockImageMap[image] = struct{}{}
 		}
-		opts = append(opts, crproxy.WithBlockFunc(func(info *crproxy.PathInfo) bool {
+		opts = append(opts, crproxy.WithBlockFunc(func(info *crproxy.ImageInfo) bool {
 			if len(allowHostMap) != 0 {
 				_, ok := allowHostMap[info.Host]
 				if !ok {
@@ -180,7 +212,7 @@ func main() {
 			}
 
 			if len(blockImageMap) != 0 {
-				image := info.Host + "/" + info.Image
+				image := info.Host + "/" + info.Name
 				_, ok := blockImageMap[image]
 				if ok {
 					return true
@@ -191,6 +223,10 @@ func main() {
 		}))
 	}
 
+	if blockMessage != "" {
+		opts = append(opts, crproxy.WithBlockMessage(blockMessage))
+	}
+
 	if len(privilegedIPList) != 0 {
 		opts = append(opts, crproxy.WithPrivilegedIPs(privilegedIPList))
 	}

diff --git a/crproxy.go b/crproxy.go
@@ -23,6 +23,7 @@ import (
 	"github.com/distribution/distribution/v3/registry/client/transport"
 	storagedriver "github.com/distribution/distribution/v3/registry/storage/driver"
 	"github.com/wzshiming/geario"
+	"github.com/wzshiming/hostmatcher"
 	"github.com/wzshiming/httpseek"
 	"github.com/wzshiming/lru"
 )
@@ -36,17 +37,17 @@ type Logger interface {
 	Println(v ...interface{})
 }
 
-type ModifyInfo struct {
-	Host  string
-	Image string
+type ImageInfo struct {
+	Host string
+	Name string
 }
 
 type CRProxy struct {
 	baseClient              *http.Client
 	challengeManager        challenge.Manager
 	clientset               maps.SyncMap[string, *lru.LRU[string, *http.Client]]
 	clientSize              int
-	modify                  func(info *ModifyInfo) *ModifyInfo
+	modify                  func(info *ImageInfo) *ImageInfo
 	insecureDomain          map[string]struct{}
 	domainDisableKeepAlives map[string]struct{}
 	domainAlias             map[string]string
@@ -61,7 +62,8 @@ type CRProxy struct {
 	blobsSpeedLimitDuration time.Duration
 	ipsSpeedLimit           *geario.B
 	ipsSpeedLimitDuration   time.Duration
-	blockFunc               func(*PathInfo) bool
+	blockFunc               func(*ImageInfo) bool
+	blockMessage            string
 	retry                   int
 	retryInterval           time.Duration
 	storageDriver           storagedriver.StorageDriver
@@ -72,6 +74,7 @@ type CRProxy struct {
 	privilegedIPSet         map[string]struct{}
 	disableTagsList         bool
 	simpleAuth              bool
+	matcher                 hostmatcher.Matcher
 
 	defaultRegistry string
 }
@@ -175,7 +178,7 @@ func WithDomainAlias(domainAlias map[string]string) Option {
 	}
 }
 
-func WithPathInfoModifyFunc(modify func(info *ModifyInfo) *ModifyInfo) Option {
+func WithPathInfoModifyFunc(modify func(info *ImageInfo) *ImageInfo) Option {
 	return func(c *CRProxy) {
 		c.modify = modify
 	}
@@ -196,12 +199,18 @@ func WithDisableKeepAlives(disableKeepAlives []string) Option {
 	}
 }
 
-func WithBlockFunc(blockFunc func(info *PathInfo) bool) Option {
+func WithBlockFunc(blockFunc func(info *ImageInfo) bool) Option {
 	return func(c *CRProxy) {
 		c.blockFunc = blockFunc
 	}
 }
 
+func WithBlockMessage(msg string) Option {
+	return func(c *CRProxy) {
+		c.blockMessage = msg
+	}
+}
+
 func WithRetry(retry int, retryInterval time.Duration) Option {
 	return func(c *CRProxy) {
 		c.retry = retry
@@ -469,16 +478,23 @@ func (c *CRProxy) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 	info.Host = c.getDomainAlias(info.Host)
 
 	if c.modify != nil {
-		n := c.modify(&ModifyInfo{
-			Host:  info.Host,
-			Image: info.Image,
+		n := c.modify(&ImageInfo{
+			Host: info.Host,
+			Name: info.Image,
 		})
 		info.Host = n.Host
-		info.Image = n.Image
+		info.Image = n.Name
 	}
 
-	if c.blockFunc != nil && c.blockFunc(info) {
-		errcode.ServeJSON(rw, errcode.ErrorCodeDenied)
+	if c.blockFunc != nil && c.blockFunc(&ImageInfo{
+		Host: info.Host,
+		Name: info.Image,
+	}) {
+		if c.blockMessage != "" {
+			errcode.ServeJSON(rw, errcode.ErrorCodeDenied.WithMessage(c.blockMessage))
+		} else {
+			errcode.ServeJSON(rw, errcode.ErrorCodeDenied)
+		}
 		return
 	}
 

diff --git a/go.mod b/go.mod
@@ -9,6 +9,7 @@ require (
 	github.com/opencontainers/go-digest v1.0.0
 	github.com/spf13/pflag v1.0.3
 	github.com/wzshiming/geario v0.0.0-20220512063849-ddc203ba7487
+	github.com/wzshiming/hostmatcher v0.0.3
 	github.com/wzshiming/httpseek v0.1.0
 	github.com/wzshiming/lru v0.1.0
 )
@@ -50,11 +51,11 @@ require (
 	github.com/sirupsen/logrus v1.9.3 // indirect
 	github.com/stretchr/testify v1.8.4 // indirect
 	go.opencensus.io v0.24.0 // indirect
-	golang.org/x/crypto v0.17.0 // indirect
-	golang.org/x/net v0.18.0 // indirect
+	golang.org/x/crypto v0.24.0 // indirect
+	golang.org/x/net v0.26.0 // indirect
 	golang.org/x/oauth2 v0.11.0 // indirect
-	golang.org/x/sys v0.15.0 // indirect
-	golang.org/x/text v0.14.0 // indirect
+	golang.org/x/sys v0.21.0 // indirect
+	golang.org/x/text v0.16.0 // indirect
 	google.golang.org/api v0.126.0 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/cloud v0.0.0-20151119220103-975617b05ea8 // indirect

diff --git a/go.sum b/go.sum
@@ -203,6 +203,8 @@ github.com/uber/jaeger-lib v2.4.1+incompatible h1:td4jdvLcExb4cBISKIpHuGoVXh+dVK
 github.com/uber/jaeger-lib v2.4.1+incompatible/go.mod h1:ComeNDZlWwrWnDv8aPp0Ba6+uUTzImX/AauajbLI56U=
 github.com/wzshiming/geario v0.0.0-20220512063849-ddc203ba7487 h1:5YN2BY58L2/5N7olluOPqJmmybcio0+DFY6nwH5gc8c=
 github.com/wzshiming/geario v0.0.0-20220512063849-ddc203ba7487/go.mod h1:Fodw3HJvNUS+/MgqXCRp9iYLQfynAu/LKXGOWoX+D/Q=
+github.com/wzshiming/hostmatcher v0.0.3 h1:+JYAq6vUZXDEQ1Ipfdc/D7HmaIMngcc71ftonyCQVQk=
+github.com/wzshiming/hostmatcher v0.0.3/go.mod h1:F04RIvIWEvOIrIKOlQlMuR8vQMKAVf2YhpU6l31Wwz4=
 github.com/wzshiming/httpseek v0.1.0 h1:lEgL7EBELT/VV9UaTp+m3kw5Pe1KOUdY+IPnKkag6tI=
 github.com/wzshiming/httpseek v0.1.0/go.mod h1:YoZhlLIwNjTBDXIT8NpK5zRjOgZouRXPaBfjVXdqMMs=
 github.com/wzshiming/lru v0.1.0 h1:937SBBo9lDBRMivJqF46eVK2Q3omRWnN7hWpr2i3xZg=
@@ -219,8 +221,8 @@ golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPh
 golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
 golang.org/x/crypto v0.0.0-20211215153901-e495a2d5b3d3/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
 golang.org/x/crypto v0.0.0-20220314234659-1baeb1ce4c0b/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
-golang.org/x/crypto v0.17.0 h1:r8bRNjWL3GshPW3gkd+RpvzWrZAwPS49OmTGZ/uhM4k=
-golang.org/x/crypto v0.17.0/go.mod h1:gCAAfMLgwOJRpTjQ2zCCt2OcSfYMTeZVSRtQlPC7Nq4=
+golang.org/x/crypto v0.24.0 h1:mnl8DM0o513X8fdIkmyFE/5hTYxbwYOjDS/+rK6qpRI=
+golang.org/x/crypto v0.24.0/go.mod h1:Z1PMYSOR5nyMcyAVAIQSKCDwalqy85Aqn1x3Ws4L5DM=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
 golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
@@ -241,8 +243,8 @@ golang.org/x/net v0.0.0-20201110031124-69a78807bb2b/go.mod h1:sp8m0HH+o8qH0wwXwY
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
-golang.org/x/net v0.18.0 h1:mIYleuAkSbHh0tCv7RvjL3F6ZVbLjq4+R7zbOn3Kokg=
-golang.org/x/net v0.18.0/go.mod h1:/czyP5RqHAH4odGYxBJ1qz0+CE5WZ+2j1YgoEo8F2jQ=
+golang.org/x/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ=
+golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/oauth2 v0.0.0-20200107190931-bf48bf16ab8d/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
 golang.org/x/oauth2 v0.11.0 h1:vPL4xzxBM4niKCW6g9whtaWVXTJf1U5e4aZxxFx/gbU=
@@ -252,8 +254,8 @@ golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.3.0 h1:ftCYgMx6zT/asHUrPw8BLLscYtGznsLAnjq5RH9P66E=
-golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
+golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M=
+golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
@@ -268,8 +270,8 @@ golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.15.0 h1:h48lPFYpsTvQJZF4EKyI4aLHaev3CxivZmv7yZig9pc=
-golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws=
+golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
@@ -278,8 +280,8 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
 golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ=
-golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
-golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
+golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4=
+golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=

diff --git a/sync.go b/sync.go
@@ -97,15 +97,23 @@ func (c *CRProxy) SyncImageLayer(ctx context.Context, image string, filter func(
 	host = c.getDomainAlias(host)
 	var name reference.Named
 
+	info := &ImageInfo{
+		Host: host,
+		Name: reference.Path(named),
+	}
 	if c.modify != nil {
-		n := c.modify(&ModifyInfo{
-			Host:  host,
-			Image: reference.Path(named),
-		})
-		host = n.Host
-		name = newNameWithoutDomain(named, n.Image)
+		info = c.modify(info)
+		name = newNameWithoutDomain(named, info.Name)
 	} else {
-		name = newNameWithoutDomain(named, reference.Path(named))
+		name = newNameWithoutDomain(named, info.Name)
+	}
+
+	if c.blockFunc != nil && c.blockFunc(info) {
+		if c.blockMessage != "" {
+			return errcode.ErrorCodeDenied.WithMessage(c.blockMessage)
+		} else {
+			return errcode.ErrorCodeDenied
+		}
 	}
 
 	err = c.ping(host)