ollama/types/model/name_test.go

package model

import (
	"bytes"
	"cmp"
	"fmt"
	"log/slog"
	"slices"
	"strings"
	"testing"
)

type fields struct {
	host, namespace, model, tag, build string
	digest                             string
}

func fieldsFromName(p Name) fields {
	return fields{
		host:      p.parts[PartHost],
		namespace: p.parts[PartNamespace],
		model:     p.parts[PartModel],
		tag:       p.parts[PartTag],
		build:     p.parts[PartBuild],
		digest:    p.parts[PartDigest],
	}
}

var testNames = map[string]fields{
	"mistral:latest":                 {model: "mistral", tag: "latest"},
	"mistral":                        {model: "mistral"},
	"mistral:30B":                    {model: "mistral", tag: "30B"},
	"mistral:7b":                     {model: "mistral", tag: "7b"},
	"mistral:7b+Q4_0":                {model: "mistral", tag: "7b", build: "Q4_0"},
	"mistral+KQED":                   {model: "mistral", build: "KQED"},
	"mistral.x-3:7b+Q4_0":            {model: "mistral.x-3", tag: "7b", build: "Q4_0"},
	"mistral:7b+q4_0":                {model: "mistral", tag: "7b", build: "q4_0"},
	"llama2":                         {model: "llama2"},
	"user/model":                     {namespace: "user", model: "model"},
	"example.com/ns/mistral:7b+Q4_0": {host: "example.com", namespace: "ns", model: "mistral", tag: "7b", build: "Q4_0"},
	"example.com/ns/mistral:7b+X":    {host: "example.com", namespace: "ns", model: "mistral", tag: "7b", build: "X"},

	// invalid digest
	"mistral:latest@invalid256-": {},
	"mistral:latest@-123":        {},
	"mistral:latest@!-123":       {},
	"mistral:latest@1-!":         {},
	"mistral:latest@":            {},

	// resolved
	"x@sha123-1": {model: "x", digest: "sha123-1"},
	"@sha456-2":  {digest: "sha456-2"},

	"@@sha123-1": {},

	// preserves case for build
	"x+b": {model: "x", build: "b"},

	// invalid (includes fuzzing trophies)
	" / / : + ": {},
	" / : + ":   {},
	" : + ":     {},
	" + ":       {},
	" : ":       {},
	" / ":       {},
	" /":        {},
	"/ ":        {},
	"/":         {},
	":":         {},
	"+":         {},

	// (".") in namepsace is not allowed
	"invalid.com/7b+x": {},

	"invalid:7b+Q4_0:latest": {},
	"in valid":               {},
	"invalid/y/z/foo":        {},
	"/0":                     {},
	"0 /0":                   {},
	"0 /":                    {},
	"0/":                     {},
	":/0":                    {},
	"+0/00000":               {},
	"0+.\xf2\x80\xf6\x9d00000\xe5\x99\xe6\xd900\xd90\xa60\x91\xdc0\xff\xbf\x99\xe800\xb9\xdc\xd6\xc300\x970\xfb\xfd0\xe0\x8a\xe1\xad\xd40\x9700\xa80\x980\xdd0000\xb00\x91000\xfe0\x89\x9b\x90\x93\x9f0\xe60\xf7\x84\xb0\x87\xa5\xff0\xa000\x9a\x85\xf6\x85\xfe\xa9\xf9\xe9\xde00\xf4\xe0\x8f\x81\xad\xde00\xd700\xaa\xe000000\xb1\xee0\x91": {},
	"0//0":                        {},
	"m+^^^":                       {},
	"file:///etc/passwd":          {},
	"file:///etc/passwd:latest":   {},
	"file:///etc/passwd:latest+u": {},

	":x": {},
	"+x": {},
	"x+": {},

	// Disallow ("\.+") in any part to prevent path traversal anywhere
	// we convert the name to a path.
	"../etc/passwd":  {},
	".../etc/passwd": {},
	"./../passwd":    {},
	"./0+..":         {},

	strings.Repeat("a", MaxNamePartLen):   {model: strings.Repeat("a", MaxNamePartLen)},
	strings.Repeat("a", MaxNamePartLen+1): {},
}

// TestConsecutiveDots tests that consecutive dots are not allowed in any
// part, to avoid path traversal. There also are some tests in testNames, but
// this test is more exhaustive and exists to emphasize the importance of
// preventing path traversal.
func TestNameConsecutiveDots(t *testing.T) {
	for i := 1; i < 10; i++ {
		s := strings.Repeat(".", i)
		if i > 1 {
			if g := ParseNameFill(s, "").String(); g != "" {
				t.Errorf("ParseName(%q) = %q; want empty string", s, g)
			}
		} else {
			if g := ParseNameFill(s, "").String(); g != s {
				t.Errorf("ParseName(%q) = %q; want %q", s, g, s)
			}
		}
	}
}

func TestNameParts(t *testing.T) {
	var p Name
	if w, g := int(PartDigest+1), len(p.Parts()); w != g {
		t.Errorf("Parts() = %d; want %d", g, w)
	}
}

func TestNamePartString(t *testing.T) {
	if g := PartKind(-2).String(); g != "Unknown" {
		t.Errorf("Unknown part = %q; want %q", g, "Unknown")
	}
	for kind, name := range kindNames {
		if g := kind.String(); g != name {
			t.Errorf("%s = %q; want %q", kind, g, name)
		}
	}
}

func TestParseName(t *testing.T) {
	for baseName, want := range testNames {
		for _, prefix := range []string{"", "https://", "http://"} {
			// We should get the same results with or without the
			// http(s) prefixes
			s := prefix + baseName

			t.Run(s, func(t *testing.T) {
				name := ParseNameFill(s, "")
				got := fieldsFromName(name)
				if got != want {
					t.Errorf("ParseName(%q) = %q; want %q", s, got, want)
				}

				// test round-trip
				if !ParseNameFill(name.String(), "").EqualFold(name) {
					t.Errorf("ParseName(%q).String() = %s; want %s", s, name.String(), baseName)
				}
			})
		}
	}
}

func TestCompleteWithAndWithoutBuild(t *testing.T) {
	cases := []struct {
		in              string
		complete        bool
		completeNoBuild bool
	}{
		{"", false, false},
		{"incomplete/mistral:7b+x", false, false},
		{"incomplete/mistral:7b+Q4_0", false, false},
		{"incomplete:7b+x", false, false},
		{"complete.com/x/mistral:latest+Q4_0", true, true},
		{"complete.com/x/mistral:latest", false, true},
	}

	for _, tt := range cases {
		t.Run(tt.in, func(t *testing.T) {
			p := ParseNameFill(tt.in, "")
			t.Logf("ParseName(%q) = %#v", tt.in, p)
			if g := p.IsComplete(); g != tt.complete {
				t.Errorf("Complete(%q) = %v; want %v", tt.in, g, tt.complete)
			}
			if g := p.IsCompleteNoBuild(); g != tt.completeNoBuild {
				t.Errorf("CompleteNoBuild(%q) = %v; want %v", tt.in, g, tt.completeNoBuild)
			}
		})
	}

	// Complete uses Parts which returns a slice, but it should be
	// inlined when used in Complete, preventing any allocations or
	// escaping to the heap.
	allocs := testing.AllocsPerRun(1000, func() {
		keep(ParseNameFill("complete.com/x/mistral:latest+Q4_0", "").IsComplete())
	})
	if allocs > 0 {
		t.Errorf("Complete allocs = %v; want 0", allocs)
	}
}

func TestNameLogValue(t *testing.T) {
	cases := []string{
		"example.com/library/mistral:latest+Q4_0",
		"mistral:latest",
		"mistral:7b+Q4_0",
	}
	for _, s := range cases {
		t.Run(s, func(t *testing.T) {
			var b bytes.Buffer
			log := slog.New(slog.NewTextHandler(&b, nil))
			name := ParseNameFill(s, "")
			log.Info("", "name", name)
			want := fmt.Sprintf("name=%s", name.GoString())
			got := b.String()
			if !strings.Contains(got, want) {
				t.Errorf("expected log output to contain %q; got %q", want, got)
			}
		})
	}
}

func TestNameGoString(t *testing.T) {
	cases := []struct {
		name         string
		in           string
		wantString   string
		wantGoString string // default is tt.in
	}{
		{
			name:         "Complete Name",
			in:           "example.com/library/mistral:latest+Q4_0",
			wantGoString: "example.com/library/mistral:latest+Q4_0@?",
		},
		{
			name:         "Short Name",
			in:           "mistral:latest",
			wantGoString: "?/?/mistral:latest+?@?",
		},
		{
			name:         "Long Name",
			in:           "library/mistral:latest",
			wantGoString: "?/library/mistral:latest+?@?",
		},
		{
			name:         "Case Preserved",
			in:           "Library/Mistral:Latest",
			wantGoString: "?/Library/Mistral:Latest+?@?",
		},
		{
			name:         "With digest",
			in:           "Library/Mistral:Latest@sha256-123456",
			wantGoString: "?/Library/Mistral:Latest+?@sha256-123456",
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			p := ParseNameFill(tt.in, "")
			tt.wantGoString = cmp.Or(tt.wantGoString, tt.in)
			if g := fmt.Sprintf("%#v", p); g != tt.wantGoString {
				t.Errorf("GoString() = %q; want %q", g, tt.wantGoString)
			}
		})
	}
}

func TestDisplayShortest(t *testing.T) {
	cases := []struct {
		in        string
		mask      string
		want      string
		wantPanic bool
	}{
		{"example.com/library/mistral:latest+Q4_0", "example.com/library/_:latest", "mistral", false},
		{"example.com/library/mistral:latest+Q4_0", "example.com/_/_:latest", "library/mistral", false},
		{"example.com/library/mistral:latest+Q4_0", "", "example.com/library/mistral", false},
		{"example.com/library/mistral:latest+Q4_0", "", "example.com/library/mistral", false},

		// case-insensitive
		{"Example.com/library/mistral:latest+Q4_0", "example.com/library/_:latest", "mistral", false},
		{"example.com/Library/mistral:latest+Q4_0", "example.com/library/_:latest", "mistral", false},
		{"example.com/library/Mistral:latest+Q4_0", "example.com/library/_:latest", "Mistral", false},
		{"example.com/library/mistral:Latest+Q4_0", "example.com/library/_:latest", "mistral", false},
		{"example.com/library/mistral:Latest+q4_0", "example.com/library/_:latest", "mistral", false},

		// invalid mask
		{"example.com/library/mistral:latest+Q4_0", "example.com/mistral", "", true},

		// DefaultMask
		{"registry.ollama.ai/library/mistral:latest+Q4_0", DefaultMask, "mistral", false},

		// Auto-Fill
		{"x", "example.com/library/_:latest", "x", false},
		{"x", "example.com/library/_:latest+Q4_0", "x", false},
		{"x/y:z", "a.com/library/_:latest+Q4_0", "x/y:z", false},
		{"x/y:z", "a.com/library/_:latest+Q4_0", "x/y:z", false},
	}

	for _, tt := range cases {
		t.Run("", func(t *testing.T) {
			defer func() {
				if tt.wantPanic {
					if recover() == nil {
						t.Errorf("expected panic")
					}
				}
			}()

			p := ParseNameFill(tt.in, "")
			t.Logf("ParseName(%q) = %#v", tt.in, p)
			if g := p.DisplayShortest(tt.mask); g != tt.want {
				t.Errorf("got = %q; want %q", g, tt.want)
			}
		})
	}
}

func TestParseNameAllocs(t *testing.T) {
	allocs := testing.AllocsPerRun(1000, func() {
		keep(ParseNameFill("example.com/mistral:7b+Q4_0", ""))
	})
	if allocs > 0 {
		t.Errorf("ParseName allocs = %v; want 0", allocs)
	}
}

func BenchmarkParseName(b *testing.B) {
	b.ReportAllocs()

	for range b.N {
		keep(ParseNameFill("example.com/mistral:7b+Q4_0", ""))
	}
}

func FuzzParseName(f *testing.F) {
	f.Add("example.com/mistral:7b+Q4_0")
	f.Add("example.com/mistral:7b+q4_0")
	f.Add("example.com/mistral:7b+x")
	f.Add("x/y/z:8n+I")
	f.Add(":x")
	f.Add("@sha256-123456")
	f.Add("example.com/mistral:latest+Q4_0@sha256-123456")
	f.Add(":@!@")
	f.Add("...")
	f.Fuzz(func(t *testing.T, s string) {
		r0 := ParseNameFill(s, "")

		if strings.Contains(s, "..") && !r0.IsZero() {
			t.Fatalf("non-zero value for path with '..': %q", s)
		}

		if !r0.IsValid() && !r0.IsResolved() {
			if !r0.EqualFold(Name{}) {
				t.Errorf("expected invalid path to be zero value; got %#v", r0)
			}
			t.Skipf("invalid path: %q", s)
		}

		for _, p := range r0.Parts() {
			if len(p) > MaxNamePartLen {
				t.Errorf("part too long: %q", p)
			}
		}

		if !strings.EqualFold(r0.String(), s) {
			t.Errorf("String() did not round-trip with case insensitivity: %q\ngot  = %q\nwant = %q", s, r0.String(), s)
		}

		r1 := ParseNameFill(r0.String(), "")
		if !r0.EqualFold(r1) {
			t.Errorf("round-trip mismatch: %+v != %+v", r0, r1)
		}
	})
}

func TestFill(t *testing.T) {
	cases := []struct {
		dst  string
		src  string
		want string
	}{
		{"mistral", "o.com/library/PLACEHOLDER:latest+Q4_0", "o.com/library/mistral:latest+Q4_0"},
		{"o.com/library/mistral", "PLACEHOLDER:latest+Q4_0", "o.com/library/mistral:latest+Q4_0"},
		{"", "o.com/library/mistral:latest+Q4_0", "o.com/library/mistral:latest+Q4_0"},
	}

	for _, tt := range cases {
		t.Run(tt.dst, func(t *testing.T) {
			r := Fill(ParseNameFill(tt.dst, ""), ParseNameFill(tt.src, ""))
			if r.String() != tt.want {
				t.Errorf("Fill(%q, %q) = %q; want %q", tt.dst, tt.src, r, tt.want)
			}
		})
	}
}

func TestNameStringAllocs(t *testing.T) {
	name := ParseNameFill("example.com/ns/mistral:latest+Q4_0", "")
	allocs := testing.AllocsPerRun(1000, func() {
		keep(name.String())
	})
	if allocs > 1 {
		t.Errorf("String allocs = %v; want 0", allocs)
	}
}

func ExampleFill() {
	defaults := ParseNameFill("registry.ollama.com/library/PLACEHOLDER:latest+Q4_0", "")
	r := Fill(ParseNameFill("mistral", ""), defaults)
	fmt.Println(r)

	// Output:
	// registry.ollama.com/library/mistral:latest+Q4_0
}

func ExampleName_MapHash() {
	m := map[uint64]bool{}

	// key 1
	m[ParseNameFill("mistral:latest+q4", "").MapHash()] = true
	m[ParseNameFill("miSTRal:latest+Q4", "").MapHash()] = true
	m[ParseNameFill("mistral:LATest+Q4", "").MapHash()] = true

	// key 2
	m[ParseNameFill("mistral:LATest", "").MapHash()] = true

	fmt.Println(len(m))
	// Output:
	// 2
}

func ExampleName_CompareFold_sort() {
	names := []Name{
		ParseNameFill("mistral:latest", ""),
		ParseNameFill("mistRal:7b+q4", ""),
		ParseNameFill("MIstral:7b", ""),
	}

	slices.SortFunc(names, Name.CompareFold)

	for _, n := range names {
		fmt.Println(n)
	}

	// Output:
	// MIstral:7b
	// mistRal:7b+q4
	// mistral:latest
}

func ExampleName_completeAndResolved() {
	for _, s := range []string{
		"x/y/z:latest+q4_0@sha123-1",
		"x/y/z:latest+q4_0",
		"@sha123-1",
	} {
		name := ParseNameFill(s, "")
		fmt.Printf("complete:%v resolved:%v  digest:%s\n", name.IsComplete(), name.IsResolved(), name.Digest())
	}

	// Output:
	// complete:true resolved:true  digest:sha123-1
	// complete:true resolved:false  digest:
	// complete:false resolved:true  digest:sha123-1
}

func ExampleName_DisplayShortest() {
	name := ParseNameFill("example.com/jmorganca/mistral:latest+Q4_0", "")

	fmt.Println(name.DisplayShortest("example.com/jmorganca/_:latest"))
	fmt.Println(name.DisplayShortest("example.com/_/_:latest"))
	fmt.Println(name.DisplayShortest("example.com/_/_:_"))
	fmt.Println(name.DisplayShortest("_/_/_:_"))

	// Default
	name = ParseNameFill("registry.ollama.ai/library/mistral:latest+Q4_0", "")
	fmt.Println(name.DisplayShortest(""))

	// Output:
	// mistral
	// jmorganca/mistral
	// jmorganca/mistral:latest
	// example.com/jmorganca/mistral:latest
	// mistral
}

func keep[T any](v T) T { return v }