LocalAI/pkg/downloader/uri_test.go

package downloader_test

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"regexp"
	"strconv"

	. "github.com/mudler/LocalAI/pkg/downloader"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

var _ = Describe("Gallery API tests", func() {
	Context("URI", func() {
		It("parses github with a branch", func() {
			uri := URI("github:go-skynet/model-gallery/gpt4all-j.yaml")
			Expect(
				uri.ReadWithCallback("", func(url string, i []byte) error {
					Expect(url).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml"))
					return nil
				}),
			).ToNot(HaveOccurred())
		})
		It("parses github without a branch", func() {
			uri := URI("github:go-skynet/model-gallery/gpt4all-j.yaml@main")

			Expect(
				uri.ReadWithCallback("", func(url string, i []byte) error {
					Expect(url).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml"))
					return nil
				}),
			).ToNot(HaveOccurred())
		})
		It("parses github with urls", func() {
			uri := URI("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml")
			Expect(
				uri.ReadWithCallback("", func(url string, i []byte) error {
					Expect(url).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml"))
					return nil
				}),
			).ToNot(HaveOccurred())
		})
	})

	Context("HuggingFace mirror", func() {
		var originalEndpoint string

		BeforeEach(func() {
			originalEndpoint = HF_ENDPOINT
		})

		AfterEach(func() {
			HF_ENDPOINT = originalEndpoint
		})

		It("rewrites direct https://huggingface.co URLs when mirror is set", func() {
			HF_ENDPOINT = "https://hf-mirror.com"
			uri := URI("https://huggingface.co/TheBloke/model-GGUF/resolve/main/model.Q4_K_M.gguf")
			Expect(uri.ResolveURL()).To(Equal("https://hf-mirror.com/TheBloke/model-GGUF/resolve/main/model.Q4_K_M.gguf"))
		})

		It("does not rewrite direct https://huggingface.co URLs when no mirror is set", func() {
			HF_ENDPOINT = "https://huggingface.co"
			uri := URI("https://huggingface.co/TheBloke/model-GGUF/resolve/main/model.Q4_K_M.gguf")
			Expect(uri.ResolveURL()).To(Equal("https://huggingface.co/TheBloke/model-GGUF/resolve/main/model.Q4_K_M.gguf"))
		})

		It("rewrites hf:// URIs when mirror is set", func() {
			HF_ENDPOINT = "https://hf-mirror.com"
			uri := URI("hf://TheBloke/model-GGUF/model.Q4_K_M.gguf")
			Expect(uri.ResolveURL()).To(Equal("https://hf-mirror.com/TheBloke/model-GGUF/resolve/main/model.Q4_K_M.gguf"))
		})

		It("does not rewrite non-huggingface URLs", func() {
			HF_ENDPOINT = "https://hf-mirror.com"
			uri := URI("https://example.com/some/file.gguf")
			Expect(uri.ResolveURL()).To(Equal("https://example.com/some/file.gguf"))
		})
	})
})

var _ = Describe("OCIReference", func() {
	DescribeTable("returns the registry reference without the oci:// scheme",
		func(uri, want string) {
			Expect(URI(uri).OCIReference()).To(Equal(want))
		},
		Entry("tag form", "oci://registry.example.com/acme/backend:v1", "registry.example.com/acme/backend:v1"),
		Entry("digest form", "oci://registry.example.com:5000/acme/backend@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", "registry.example.com:5000/acme/backend@sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"),
		Entry("reference without a scheme", "quay.io/acme/backend:latest", "quay.io/acme/backend:latest"),
	)
})

var _ = Describe("LooksLikeRegistryOCI", func() {
	DescribeTable("accepts only URIs that name an image in a registry",
		func(uri string, want bool) {
			Expect(URI(uri).LooksLikeRegistryOCI()).To(Equal(want))
		},
		Entry("oci:// reference", "oci://registry.example.com/acme/backend:v1", true),
		Entry("bare quay.io reference", "quay.io/acme/backend:latest", true),
		Entry("ollama:// model", "ollama://gemma:2b", false),
		Entry("ocifile:// tarball", "ocifile:///srv/backend.tar", false),
		Entry("https URL", "https://example.com/backend.tar", false),
	)
})

var _ = Describe("ContentLength", func() {
	Context("local file", func() {
		It("returns file size for existing file", func() {
			dir, err := os.MkdirTemp("", "contentlength-*")
			Expect(err).ToNot(HaveOccurred())
			defer os.RemoveAll(dir)
			fpath := filepath.Join(dir, "model.gguf")
			err = os.WriteFile(fpath, make([]byte, 1234), 0644)
			Expect(err).ToNot(HaveOccurred())
			uri := URI("file://" + fpath)
			ctx := context.Background()
			size, err := uri.ContentLength(ctx)
			Expect(err).ToNot(HaveOccurred())
			Expect(size).To(Equal(int64(1234)))
		})
		It("returns error for missing file", func() {
			uri := URI("file:///nonexistent/path/model.gguf")
			ctx := context.Background()
			_, err := uri.ContentLength(ctx)
			Expect(err).To(HaveOccurred())
		})
	})
	Context("HTTP", func() {
		It("returns Content-Length when present", func() {
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				Expect(r.Method).To(Equal("HEAD"))
				w.Header().Set("Content-Length", "1000")
				w.WriteHeader(http.StatusOK)
			}))
			defer server.Close()
			uri := URI(server.URL)
			ctx := context.Background()
			size, err := uri.ContentLength(ctx)
			Expect(err).ToNot(HaveOccurred())
			Expect(size).To(Equal(int64(1000)))
		})
		It("returns error on 404", func() {
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.WriteHeader(http.StatusNotFound)
			}))
			defer server.Close()
			uri := URI(server.URL)
			ctx := context.Background()
			_, err := uri.ContentLength(ctx)
			Expect(err).To(HaveOccurred())
		})
		It("uses Range when Content-Length missing and Accept-Ranges bytes", func() {
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				if r.Method == "HEAD" {
					w.Header().Set("Accept-Ranges", "bytes")
					w.WriteHeader(http.StatusOK)
					return
				}
				Expect(r.Header.Get("Range")).To(Equal("bytes=0-0"))
				w.Header().Set("Content-Range", "bytes 0-0/5000")
				w.WriteHeader(http.StatusPartialContent)
			}))
			defer server.Close()
			uri := URI(server.URL)
			ctx := context.Background()
			size, err := uri.ContentLength(ctx)
			Expect(err).ToNot(HaveOccurred())
			Expect(size).To(Equal(int64(5000)))
		})
		It("respects context cancellation", func() {
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Length", "1000")
				w.WriteHeader(http.StatusOK)
			}))
			defer server.Close()
			ctx, cancel := context.WithCancel(context.Background())
			cancel()
			uri := URI(server.URL)
			_, err := uri.ContentLength(ctx)
			Expect(err).To(HaveOccurred())
			Expect(errors.Is(err, context.Canceled)).To(BeTrue())
		})
	})
})

type RangeHeaderError struct {
	msg string
}

func (e *RangeHeaderError) Error() string { return e.msg }

var _ = Describe("Download Test", func() {
	var mockData []byte
	var mockDataSha string
	var filePath string

	extractRangeHeader := func(rangeString string) (int, int, error) {
		regex := regexp.MustCompile(`^bytes=(\d+)-(\d+|)$`)
		matches := regex.FindStringSubmatch(rangeString)
		rangeErr := RangeHeaderError{msg: "invalid / ill-formatted range"}
		if matches == nil {
			return -1, -1, &rangeErr
		}
		startPos, err := strconv.Atoi(matches[1])
		if err != nil {
			return -1, -1, err
		}

		endPos := -1
		if matches[2] != "" {
			endPos, err = strconv.Atoi(matches[2])
			if err != nil {
				return -1, -1, err
			}
			endPos += 1 // because range is inclusive in rangeString
		}
		return startPos, endPos, nil
	}

	getMockServer := func(supportsRangeHeader bool) *httptest.Server {
		mockServer := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.Method != "HEAD" && r.Method != "GET" {
				w.WriteHeader(http.StatusNotFound)
				return
			}
			if r.Method == "HEAD" {
				if supportsRangeHeader {
					w.Header().Add("Accept-Ranges", "bytes")
				}
				w.WriteHeader(http.StatusOK)
				return
			}
			// GET method
			startPos := 0
			endPos := len(mockData)
			var err error
			var respData []byte
			rangeString := r.Header.Get("Range")
			if rangeString != "" {
				startPos, endPos, err = extractRangeHeader(rangeString)
				if err != nil {
					if _, ok := err.(*RangeHeaderError); ok {
						w.WriteHeader(http.StatusBadRequest)
						return
					}
					Expect(err).ToNot(HaveOccurred())
				}
				if endPos == -1 {
					endPos = len(mockData)
				}
				if startPos < 0 || startPos >= len(mockData) || endPos < 0 || endPos > len(mockData) || startPos > endPos {
					w.WriteHeader(http.StatusBadRequest)
					return
				}
			}
			respData = mockData[startPos:endPos]
			w.Header().Set("Content-Length", strconv.Itoa(len(respData)))
			if rangeString != "" {
				w.Header().Set(
					"Content-Range",
					fmt.Sprintf("bytes %d-%d/%d", startPos, endPos-1, len(mockData)),
				)
				w.WriteHeader(http.StatusPartialContent)
			} else {
				w.WriteHeader(http.StatusOK)
			}
			_, _ = w.Write(respData)
		}))
		mockServer.EnableHTTP2 = true
		mockServer.Start()
		return mockServer
	}

	BeforeEach(func() {
		mockData = make([]byte, 20000)
		_, err := rand.Read(mockData)
		Expect(err).ToNot(HaveOccurred())
		_mockDataSha := sha256.New()
		_, err = _mockDataSha.Write(mockData)
		Expect(err).ToNot(HaveOccurred())
		mockDataSha = fmt.Sprintf("%x", _mockDataSha.Sum(nil))
		dir, err := os.Getwd()
		filePath = dir + "/my_supercool_model"
		Expect(err).NotTo(HaveOccurred())
	})

	Context("URI DownloadFile", func() {
		It("fetches files from mock server", func() {
			mockServer := getMockServer(true)
			defer mockServer.Close()
			uri := URI(mockServer.URL)
			err := uri.DownloadFile(filePath, mockDataSha, 1, 1, func(s1, s2, s3 string, f float64) {})
			Expect(err).ToNot(HaveOccurred())
		})

		It("resumes partially downloaded files", func() {
			mockServer := getMockServer(true)
			defer mockServer.Close()
			uri := URI(mockServer.URL)
			// Create a partial file
			tmpFilePath := filePath + ".partial"
			file, err := os.OpenFile(tmpFilePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
			Expect(err).ToNot(HaveOccurred())
			_, err = file.Write(mockData[0:10000])
			Expect(err).ToNot(HaveOccurred())
			err = uri.DownloadFile(filePath, mockDataSha, 1, 1, func(s1, s2, s3 string, f float64) {})
			Expect(err).ToNot(HaveOccurred())
		})

		It("restarts download from 0 if server doesn't support Range header", func() {
			mockServer := getMockServer(false)
			defer mockServer.Close()
			uri := URI(mockServer.URL)
			// Create a partial file
			tmpFilePath := filePath + ".partial"
			file, err := os.OpenFile(tmpFilePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
			Expect(err).ToNot(HaveOccurred())
			_, err = file.Write(mockData[0:10000])
			Expect(err).ToNot(HaveOccurred())
			err = uri.DownloadFile(filePath, mockDataSha, 1, 1, func(s1, s2, s3 string, f float64) {})
			Expect(err).ToNot(HaveOccurred())
		})

		// A file that fails its SHA check must not be usable. The historical
		// implementation renamed the temp file to its final path *before*
		// verifying the hash, so a mismatch returned an error but left a
		// tampered file at the destination — the next caller (e.g. a backend
		// launcher) could pick it up and run with it.
		It("does not leave a corrupted file at the destination on SHA mismatch", func() {
			mockServer := getMockServer(true)
			defer mockServer.Close()
			uri := URI(mockServer.URL)

			// Use a clearly-wrong expected SHA; the server will return real
			// data with a different hash.
			wrongSHA := "0000000000000000000000000000000000000000000000000000000000000000"
			err := uri.DownloadFile(filePath, wrongSHA, 1, 1, func(s1, s2, s3 string, f float64) {})
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("SHA"))

			// The file must not exist at the final destination.
			_, statErr := os.Stat(filePath)
			Expect(os.IsNotExist(statErr)).To(BeTrue(),
				"download with wrong SHA left a file at %s — a subsequent caller could load tampered content", filePath)
		})

		// A download without an expected digest is a supply-chain footgun.
		// The downloader allows it (backend installs pass through here
		// today and don't yet ship a digest) but it is the caller's
		// responsibility to know when integrity is required. The downloader
		// emits a WARN log on every empty-digest download to make this
		// visible at the default log level.
		It("succeeds with empty SHA but emits an integrity warning", func() {
			mockServer := getMockServer(true)
			defer mockServer.Close()
			uri := URI(mockServer.URL)

			// No assertion on logs (we don't capture xlog output here),
			// but the call must succeed so existing backend installs do
			// not regress.
			err := uri.DownloadFile(filePath, "", 1, 1, func(s1, s2, s3 string, f float64) {})
			Expect(err).ToNot(HaveOccurred())
			_, statErr := os.Stat(filePath)
			Expect(statErr).ToNot(HaveOccurred())
		})
	})

	AfterEach(func() {
		os.Remove(filePath) // cleanup, also checks existence of filePath`
		os.Remove(filePath + ".partial")
	})
})