From 703684a82a150ba59b300ca9c1ec678bd9b4ab27 Mon Sep 17 00:00:00 2001 From: Blake Mizerany Date: Thu, 14 Mar 2024 20:18:06 -0700 Subject: [PATCH] server: replace blob prefix separator from ':' to '-' (#3146) This fixes issues with blob file names that contain ':' characters to be rejected by file systems that do not support them. --- server/fixblobs.go | 26 +++++++++++++ server/fixblobs_test.go | 83 +++++++++++++++++++++++++++++++++++++++++ server/images.go | 4 +- server/layers.go | 6 +-- server/modelpath.go | 6 +-- server/routes.go | 8 ++++ 6 files changed, 120 insertions(+), 13 deletions(-) create mode 100644 server/fixblobs.go create mode 100644 server/fixblobs_test.go diff --git a/server/fixblobs.go b/server/fixblobs.go new file mode 100644 index 00000000..6e6b3114 --- /dev/null +++ b/server/fixblobs.go @@ -0,0 +1,26 @@ +package server + +import ( + "os" + "path/filepath" + "strings" +) + +// fixBlobs walks the provided dir and replaces (":") to ("-") in the file +// prefix. (e.g. sha256:1234 -> sha256-1234) +func fixBlobs(dir string) error { + return filepath.Walk(dir, func(path string, info os.FileInfo, err error) error { + if err != nil { + return err + } + baseName := filepath.Base(path) + typ, sha, ok := strings.Cut(baseName, ":") + if ok && typ == "sha256" { + newPath := filepath.Join(filepath.Dir(path), typ+"-"+sha) + if err := os.Rename(path, newPath); err != nil { + return err + } + } + return nil + }) +} diff --git a/server/fixblobs_test.go b/server/fixblobs_test.go new file mode 100644 index 00000000..c2d4851a --- /dev/null +++ b/server/fixblobs_test.go @@ -0,0 +1,83 @@ +package server + +import ( + "io/fs" + "os" + "path/filepath" + "runtime" + "slices" + "strings" + "testing" +) + +func TestFixBlobs(t *testing.T) { + cases := []struct { + path []string + want []string + }{ + {path: []string{"sha256-1234"}, want: []string{"sha256-1234"}}, + {path: []string{"sha256:1234"}, want: []string{"sha256-1234"}}, + {path: []string{"sha259:5678"}, want: []string{"sha259:5678"}}, + {path: []string{"sha256:abcd"}, want: []string{"sha256-abcd"}}, + {path: []string{"x/y/sha256:abcd"}, want: []string{"x/y/sha256-abcd"}}, + {path: []string{"x:y/sha256:abcd"}, want: []string{"x:y/sha256-abcd"}}, + {path: []string{"x:y/sha256:abcd"}, want: []string{"x:y/sha256-abcd"}}, + {path: []string{"x:y/sha256:abcd", "sha256:1234"}, want: []string{"x:y/sha256-abcd", "sha256-1234"}}, + {path: []string{"x:y/sha256:abcd", "sha256-1234"}, want: []string{"x:y/sha256-abcd", "sha256-1234"}}, + } + + for _, tt := range cases { + t.Run(strings.Join(tt.path, "|"), func(t *testing.T) { + hasColon := slices.ContainsFunc(tt.path, func(s string) bool { return strings.Contains(s, ":") }) + if hasColon && runtime.GOOS == "windows" { + t.Skip("skipping test on windows") + } + + rootDir := t.TempDir() + for _, path := range tt.path { + fullPath := filepath.Join(rootDir, path) + fullDir, _ := filepath.Split(fullPath) + + t.Logf("creating dir %s", fullDir) + if err := os.MkdirAll(fullDir, 0o755); err != nil { + t.Fatal(err) + } + + t.Logf("writing file %s", fullPath) + if err := os.WriteFile(fullPath, nil, 0o644); err != nil { + t.Fatal(err) + } + } + + if err := fixBlobs(rootDir); err != nil { + t.Fatal(err) + } + + got := slurpFiles(os.DirFS(rootDir)) + + slices.Sort(tt.want) + slices.Sort(got) + if !slices.Equal(got, tt.want) { + t.Fatalf("got = %v, want %v", got, tt.want) + } + }) + } +} + +func slurpFiles(fsys fs.FS) []string { + var sfs []string + fn := func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if d.IsDir() { + return nil + } + sfs = append(sfs, path) + return nil + } + if err := fs.WalkDir(fsys, ".", fn); err != nil { + panic(err) + } + return sfs +} diff --git a/server/images.go b/server/images.go index 06f8ffd9..11c11745 100644 --- a/server/images.go +++ b/server/images.go @@ -795,9 +795,7 @@ func PruneLayers() error { for _, blob := range blobs { name := blob.Name() - if runtime.GOOS == "windows" { - name = strings.ReplaceAll(name, "-", ":") - } + name = strings.ReplaceAll(name, "-", ":") if strings.HasPrefix(name, "sha256:") { deleteMap[name] = struct{}{} } diff --git a/server/layers.go b/server/layers.go index adbf7175..07787406 100644 --- a/server/layers.go +++ b/server/layers.go @@ -5,7 +5,6 @@ import ( "fmt" "io" "os" - "runtime" "strings" "golang.org/x/exp/slices" @@ -47,10 +46,7 @@ func NewLayer(r io.Reader, mediatype string) (*Layer, error) { return nil, err } - delimiter := ":" - if runtime.GOOS == "windows" { - delimiter = "-" - } + const delimiter = "-" pattern := strings.Join([]string{"sha256", "*-partial"}, delimiter) temp, err := os.CreateTemp(blobs, pattern) diff --git a/server/modelpath.go b/server/modelpath.go index af3f36ab..7d333876 100644 --- a/server/modelpath.go +++ b/server/modelpath.go @@ -6,7 +6,6 @@ import ( "net/url" "os" "path/filepath" - "runtime" "strings" ) @@ -150,10 +149,7 @@ func GetBlobsPath(digest string) (string, error) { return "", err } - if runtime.GOOS == "windows" { - digest = strings.ReplaceAll(digest, ":", "-") - } - + digest = strings.ReplaceAll(digest, ":", "-") path := filepath.Join(dir, "blobs", digest) dirPath := filepath.Dir(path) if digest == "" { diff --git a/server/routes.go b/server/routes.go index f71bbd91..a03f39e7 100644 --- a/server/routes.go +++ b/server/routes.go @@ -1088,6 +1088,14 @@ func Serve(ln net.Listener) error { slog.SetDefault(slog.New(handler)) + blobsDir, err := GetBlobsPath("") + if err != nil { + return err + } + if err := fixBlobs(blobsDir); err != nil { + return err + } + if noprune := os.Getenv("OLLAMA_NOPRUNE"); noprune == "" { // clean up unused layers and manifests if err := PruneLayers(); err != nil {