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.
This commit is contained in:
Blake Mizerany 2024-03-14 20:18:06 -07:00 committed by GitHub
parent 6459377ae0
commit 703684a82a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 120 additions and 13 deletions

26
server/fixblobs.go Normal file
View file

@ -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
})
}

83
server/fixblobs_test.go Normal file
View file

@ -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
}

View file

@ -795,9 +795,7 @@ func PruneLayers() error {
for _, blob := range blobs { for _, blob := range blobs {
name := blob.Name() name := blob.Name()
if runtime.GOOS == "windows" {
name = strings.ReplaceAll(name, "-", ":") name = strings.ReplaceAll(name, "-", ":")
}
if strings.HasPrefix(name, "sha256:") { if strings.HasPrefix(name, "sha256:") {
deleteMap[name] = struct{}{} deleteMap[name] = struct{}{}
} }

View file

@ -5,7 +5,6 @@ import (
"fmt" "fmt"
"io" "io"
"os" "os"
"runtime"
"strings" "strings"
"golang.org/x/exp/slices" "golang.org/x/exp/slices"
@ -47,10 +46,7 @@ func NewLayer(r io.Reader, mediatype string) (*Layer, error) {
return nil, err return nil, err
} }
delimiter := ":" const delimiter = "-"
if runtime.GOOS == "windows" {
delimiter = "-"
}
pattern := strings.Join([]string{"sha256", "*-partial"}, delimiter) pattern := strings.Join([]string{"sha256", "*-partial"}, delimiter)
temp, err := os.CreateTemp(blobs, pattern) temp, err := os.CreateTemp(blobs, pattern)

View file

@ -6,7 +6,6 @@ import (
"net/url" "net/url"
"os" "os"
"path/filepath" "path/filepath"
"runtime"
"strings" "strings"
) )
@ -150,10 +149,7 @@ func GetBlobsPath(digest string) (string, error) {
return "", err return "", err
} }
if runtime.GOOS == "windows" {
digest = strings.ReplaceAll(digest, ":", "-") digest = strings.ReplaceAll(digest, ":", "-")
}
path := filepath.Join(dir, "blobs", digest) path := filepath.Join(dir, "blobs", digest)
dirPath := filepath.Dir(path) dirPath := filepath.Dir(path)
if digest == "" { if digest == "" {

View file

@ -1088,6 +1088,14 @@ func Serve(ln net.Listener) error {
slog.SetDefault(slog.New(handler)) 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 == "" { if noprune := os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
// clean up unused layers and manifests // clean up unused layers and manifests
if err := PruneLayers(); err != nil { if err := PruneLayers(); err != nil {