Add --registry-alias for gateway

DaoCloud · Dec 27, 2024 · e8186f9 · e8186f9
1 parent 71e9838
commit e8186f9
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/cmd/crproxy/cluster/gateway/gateway.go b/cmd/crproxy/cluster/gateway/gateway.go
@@ -47,6 +47,8 @@ type flagpole struct {
 
 	DefaultRegistry         string
 	OverrideDefaultRegistry map[string]string
+
+	RegistryAlias map[string]string
 }
 
 func NewCommand() *cobra.Command {
@@ -86,7 +88,7 @@ func NewCommand() *cobra.Command {
 
 	cmd.Flags().StringVar(&flags.DefaultRegistry, "default-registry", flags.DefaultRegistry, "default registry used for non full-path docker pull, like:docker.io")
 	cmd.Flags().StringToStringVar(&flags.OverrideDefaultRegistry, "override-default-registry", flags.OverrideDefaultRegistry, "override default registry")
-
+	cmd.Flags().StringToStringVar(&flags.RegistryAlias, "registry-alias", flags.RegistryAlias, "registry alias")
 	return cmd
 }
 
@@ -103,6 +105,13 @@ func runE(ctx context.Context, flags *flagpole) error {
 		gateway.WithDefaultRegistry(flags.DefaultRegistry),
 		gateway.WithOverrideDefaultRegistry(flags.OverrideDefaultRegistry),
 		gateway.WithPathInfoModifyFunc(func(info *gateway.ImageInfo) *gateway.ImageInfo {
+			if len(flags.RegistryAlias) != 0 {
+				h, ok := flags.RegistryAlias[info.Host]
+				if ok {
+					info.Host = h
+				}
+			}
+
 			if info.Host == "docker.io" {
 				info.Host = "registry-1.docker.io"
 			} else if info.Host == "ollama.ai" {