Repository: kubernetes/kubernetes
File: test/images/agnhost/fakeregistryserver/registryserver.go (Go)

Code before:
	private     bool
	registryDir = "/var/registry"
)

const (
	privateRegistryUser = "user"
	privateRegistryPass = "password"
)

func init() {
	CmdFakeRegistryServer.Flags().IntVar(&port, "port", 5000, "Port number.")
	CmdFakeRegistryServer.Flags().BoolVar(&private, "private", false, "Enable authentication for the registry.")
}

// CmdFakeRegistryServer is the cobra command for the fake registry server
var CmdFakeRegistryServer = &cobra.Command{
	Use:   "fake-registry-server",
	Short: "Starts a fake registry server for testing",
	Long:  fmt.Sprintf("Starts a fake registry server that serves static OCI image files from %s folder", registryDir),
	Run:   main,
}

func main(cmd *cobra.Command, args []string) {
	// Set registryDir relative to the executable location if not explicitly set
	// This will not affect linux platfrom but will make it easier to run tests on Windows
	if exePath, err := os.Executable(); err == nil {
		exeDir := filepath.Dir(exePath)
		registryDir = filepath.Join(exeDir, "var", "registry")
	}

	log.Printf("Registry server starting with registry directory: %s", registryDir)

	registryMux := NewRegistryServerMux(private)

	addr := fmt.Sprintf(":%d", port)
	log.Printf("HTTP server starting to listen on %s", addr)
	if err := http.ListenAndServe(addr, registryMux); err != nil {
		log.Fatalf("Error while starting the HTTP server: %v", err)
	}
}

func NewRegistryServerMux(isPrivate bool) *http.ServeMux {
	mux := http.NewServeMux()

	var v2Handler http.Handler = http.HandlerFunc(handleV2)
	if isPrivate {
		v2Handler = auth(v2Handler)
	}
	mux.Handle("/v2/", v2Handler)

	return mux

Review comment (suggestion):
In the initial version I've added the `registryDir` cobra arg to set the location of the registry, but then removed it to keep changeset smaller. Should we just add it back as an argument?

cc @stlaz @BenTheElder 

Code after:
	registryDir = "/var/registry"
)

const (
	privateRegistryUser = "user"
	privateRegistryPass = "password"
)

func init() {
	CmdFakeRegistryServer.Flags().IntVar(&port, "port", 5000, "Port number.")
	CmdFakeRegistryServer.Flags().BoolVar(&private, "private", false, "Enable authentication for the registry.")
}

// CmdFakeRegistryServer is the cobra command for the fake registry server
var CmdFakeRegistryServer = &cobra.Command{
	Use:   "fake-registry-server",
	Short: "Starts a fake registry server for testing",
	Long:  fmt.Sprintf("Starts a fake registry server that serves static OCI image files from %s folder", registryDir),
	Run:   main,
}

func main(cmd *cobra.Command, args []string) {
	log.Printf("Registry server starting with registry directory: %s", registryDir)

	registryMux := NewRegistryServerMux(private)

	addr := fmt.Sprintf(":%d", port)
	log.Printf("HTTP server starting to listen on %s", addr)
	if err := http.ListenAndServe(addr, registryMux); err != nil {
		log.Fatalf("Error while starting the HTTP server: %v", err)
	}
}

func NewRegistryServerMux(isPrivate bool) *http.ServeMux {
	mux := http.NewServeMux()

	var v2Handler http.Handler = http.HandlerFunc(handleV2)
	if isPrivate {
		v2Handler = auth(v2Handler)
	}
	mux.Handle("/v2/", v2Handler)

	return mux
}

func auth(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		user, pass, ok := r.BasicAuth()
		if !ok || user != privateRegistryUser || pass != privateRegistryPass {
			w.Header().Set("WWW-Authenticate", `Basic realm="Restricted"`)
			w.WriteHeader(http.StatusUnauthorized)