LocalAI/pkg/oci/image_auth_internal_test.go

package oci

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/google/go-containerregistry/pkg/authn"
	"github.com/google/go-containerregistry/pkg/name"
	"github.com/google/go-containerregistry/pkg/registry"
	"github.com/google/go-containerregistry/pkg/v1/random"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"

	"github.com/mudler/LocalAI/pkg/credentials"
)

const (
	registryUser = "bot"
	registryPass = "s3cret"
)

// requireBasicAuth fronts a registry the way a private registry behaves: every
// request without the right credentials gets a basic-auth challenge.
func requireBasicAuth(inner http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if u, p, ok := r.BasicAuth(); !ok || u != registryUser || p != registryPass {
			w.Header().Set("WWW-Authenticate", `Basic realm="localai-test"`)
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		inner.ServeHTTP(w, r)
	})
}

func useCredentials(doc string) {
	s, err := credentials.Parse([]byte(doc), func(string) (string, bool) { return "", false })
	Expect(err).NotTo(HaveOccurred())
	prev := credentials.SetDefault(s)
	DeferCleanup(func() { credentials.SetDefault(prev) })
}

// isolateDockerConfig stops the test host's docker or podman login from
// satisfying a spec, and returns the directory that holds config.json.
func isolateDockerConfig() string {
	home := GinkgoT().TempDir()
	dockerDir := filepath.Join(home, ".docker")
	Expect(os.MkdirAll(dockerDir, 0o700)).To(Succeed())
	GinkgoT().Setenv("HOME", home)
	GinkgoT().Setenv("DOCKER_CONFIG", dockerDir)
	GinkgoT().Setenv("REGISTRY_AUTH_FILE", filepath.Join(home, "no-such-auth.json"))
	GinkgoT().Setenv("XDG_RUNTIME_DIR", home)
	return dockerDir
}

func pushPrivateImage(serverURL, repoPath string) string {
	imageRef := strings.TrimPrefix(serverURL, "http://") + "/" + repoPath + ":latest"
	ref, err := name.ParseReference(imageRef)
	Expect(err).NotTo(HaveOccurred())
	img, err := random.Image(4096, 1)
	Expect(err).NotTo(HaveOccurred())
	Expect(remote.Write(ref, img, remote.WithAuth(&authn.Basic{Username: registryUser, Password: registryPass}))).To(Succeed())
	return imageRef
}

var _ = Describe("registry credentials", Serial, func() {
	var (
		server    *httptest.Server
		dockerDir string
		imageRef  string
	)

	BeforeEach(func() {
		dockerDir = isolateDockerConfig()
		server = httptest.NewServer(requireBasicAuth(registry.New()))
		DeferCleanup(server.Close)
		imageRef = pushPrivateImage(server.URL, "acme/backend")
	})

	It("pulls a private image with a credentials rule", func() {
		useCredentials(fmt.Sprintf("- match: %s/acme\n  username: %s\n  password: %s\n  allow_insecure: true\n", server.URL, registryUser, registryPass))

		_, err := GetImage(imageRef, "", nil, nil)
		Expect(err).NotTo(HaveOccurred())
		digest, err := GetImageDigest(imageRef, "", nil, nil)
		Expect(err).NotTo(HaveOccurred())
		Expect(digest).To(HavePrefix("sha256:"))
	})

	It("reports that no rule matches when the registry demands auth", func() {
		useCredentials("")

		_, err := GetImage(imageRef, "", nil, nil)
		var authErr *credentials.AuthError
		Expect(errors.As(err, &authErr)).To(BeTrue(), "got %v", err)
		Expect(authErr.Match).To(BeEmpty())
		Expect(authErr.Status).To(Equal(http.StatusUnauthorized))
	})

	It("names the rule the registry rejected", func() {
		rule := server.URL + "/acme"
		useCredentials(fmt.Sprintf("- match: %s\n  username: %s\n  password: not-it\n  allow_insecure: true\n", rule, registryUser))

		_, err := GetImageDigest(imageRef, "", nil, nil)
		var authErr *credentials.AuthError
		Expect(errors.As(err, &authErr)).To(BeTrue(), "got %v", err)
		Expect(authErr.Match).To(Equal(rule))
		Expect(err.Error()).NotTo(ContainSubstring("not-it"))
	})

	It("still authenticates from docker config when no rule matches", func() {
		host := strings.TrimPrefix(server.URL, "http://")
		encoded := base64.StdEncoding.EncodeToString([]byte(registryUser + ":" + registryPass))
		Expect(os.WriteFile(filepath.Join(dockerDir, "config.json"),
			[]byte(fmt.Sprintf(`{"auths":{%q:{"auth":%q}}}`, host, encoded)), 0o600)).To(Succeed())
		useCredentials("")

		_, err := GetImage(imageRef, "", nil, nil)
		Expect(err).NotTo(HaveOccurred())
	})

	It("authenticates the Range requests that resume an interrupted layer", func() {
		reg := &droppingBlobRegistry{inner: requireBasicAuth(registry.New())}
		dropping := httptest.NewServer(reg)
		DeferCleanup(dropping.Close)
		prev := layerRetryBackoff
		layerRetryBackoff = func(int) time.Duration { return 0 }
		DeferCleanup(func() { layerRetryBackoff = prev })

		ref := pushPrivateImage(dropping.URL, "acme/backend")
		useCredentials(fmt.Sprintf("- match: %s/acme\n  username: %s\n  password: %s\n  allow_insecure: true\n", dropping.URL, registryUser, registryPass))

		pulled, err := GetImage(ref, "", nil, nil)
		Expect(err).NotTo(HaveOccurred())
		tarPath := filepath.Join(GinkgoT().TempDir(), "image.tar")
		Expect(DownloadOCIImageTar(context.Background(), pulled, ref, tarPath, nil)).To(Succeed())

		reg.mu.Lock()
		defer reg.mu.Unlock()
		Expect(reg.rangeRequests).NotTo(BeEmpty())
	})
})

var _ = Describe("FetchImageBlob registry credentials", Serial, func() {
	var (
		server *httptest.Server
		repo   string
		digest string
	)

	BeforeEach(func() {
		isolateDockerConfig()
		// oras only speaks HTTPS unless a repository opts into plain HTTP,
		// which FetchImageBlob does not, so the registry has to serve TLS.
		server = httptest.NewTLSServer(requireBasicAuth(registry.New()))
		DeferCleanup(server.Close)

		// FetchImageBlob builds its oras client on retry.DefaultClient, whose
		// transport falls back to http.DefaultTransport on every request.
		// Swapping it for the test server's client is the only way to trust
		// the test certificate without a production seam; the spec is Serial
		// so nothing else sees the swap.
		prev := http.DefaultTransport
		http.DefaultTransport = server.Client().Transport
		DeferCleanup(func() { http.DefaultTransport = prev })

		repo = strings.TrimPrefix(server.URL, "https://") + "/acme/model"
		ref, err := name.ParseReference(repo + ":latest")
		Expect(err).NotTo(HaveOccurred())
		img, err := random.Image(4096, 1)
		Expect(err).NotTo(HaveOccurred())
		Expect(remote.Write(ref, img,
			remote.WithAuth(&authn.Basic{Username: registryUser, Password: registryPass}),
			remote.WithTransport(server.Client().Transport))).To(Succeed())
		layers, err := img.Layers()
		Expect(err).NotTo(HaveOccurred())
		d, err := layers[0].Digest()
		Expect(err).NotTo(HaveOccurred())
		digest = d.String()
	})

	It("fetches a private blob with a credentials rule", func() {
		useCredentials(fmt.Sprintf("- match: %s/acme\n  username: %s\n  password: %s\n", server.URL, registryUser, registryPass))

		dst := filepath.Join(GinkgoT().TempDir(), "blob")
		Expect(FetchImageBlob(context.Background(), repo, digest, dst, nil)).To(Succeed())
		info, err := os.Stat(dst)
		Expect(err).NotTo(HaveOccurred())
		Expect(info.Size()).To(BeNumerically(">", 0))
	})

	It("fails to fetch a private blob when no rule matches", func() {
		useCredentials("")

		dst := filepath.Join(GinkgoT().TempDir(), "blob")
		err := FetchImageBlob(context.Background(), repo, digest, dst, nil)
		Expect(err).To(MatchError(ContainSubstring("basic credential not found")))
	})
})