convert: only extract large files

server/model.go

@@ -81,88 +81,43 @@ func parseFromModel(ctx context.Context, name model.Name, fn func(api.ProgressRe
 	return layers, nil
 }
 
-func extractFromZipFile(p string, file *os.File, fn func(api.ProgressResponse)) error {
-	stat, err := file.Stat()
-	if err != nil {
-		return err
-	}
-
-	r, err := zip.NewReader(file, stat.Size())
-	if err != nil {
-		return err
-	}
-
-	fn(api.ProgressResponse{Status: "unpacking model metadata"})
-	for _, f := range r.File {
-		if !filepath.IsLocal(f.Name) {
-			return fmt.Errorf("%w: %s", zip.ErrInsecurePath, f.Name)
-		}
-
-		n := filepath.Join(p, f.Name)
-		if err := os.MkdirAll(filepath.Dir(n), 0o750); err != nil {
-			return err
-		}
-
-		// TODO(mxyng): this should not write out all files to disk
-		outfile, err := os.Create(n)
-		if err != nil {
-			return err
-		}
-		defer outfile.Close()
-
-		infile, err := f.Open()
-		if err != nil {
-			return err
-		}
-		defer infile.Close()
-
-		if _, err = io.Copy(outfile, infile); err != nil {
-			return err
-		}
-
-		if err := outfile.Close(); err != nil {
-			return err
-		}
-
-		if err := infile.Close(); err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
-func parseFromZipFile(_ context.Context, file *os.File, digest string, fn func(api.ProgressResponse)) (layers []*layerGGML, err error) {
-	tempDir, err := os.MkdirTemp(filepath.Dir(file.Name()), "")
+func parseFromZipFile(_ context.Context, f *os.File, digest string, fn func(api.ProgressResponse)) (layers []*layerGGML, err error) {
+	fi, err := f.Stat()
 	if err != nil {
 		return nil, err
 	}
-	defer os.RemoveAll(tempDir)
 
-	if err := extractFromZipFile(tempDir, file, fn); err != nil {
+	r, err := zip.NewReader(f, fi.Size())
+	if err != nil {
 		return nil, err
 	}
 
+	p, err := os.MkdirTemp(filepath.Dir(f.Name()), "")
+	if err != nil {
+		return nil, err
+	}
+	defer os.RemoveAll(p)
+
 	fn(api.ProgressResponse{Status: "converting model"})
-
 	// TODO(mxyng): this should write directly into a layer
 	// e.g. NewLayer(arch.Reader(), "application/vnd.ollama.image.model")
-	temp, err := os.CreateTemp(tempDir, "fp16")
+	t, err := os.CreateTemp(p, "fp16")
 	if err != nil {
 		return nil, err
 	}
-	defer temp.Close()
-	defer os.Remove(temp.Name())
+	defer t.Close()
+	defer os.Remove(t.Name())
 
-	if err := convert.Convert(tempDir, temp); err != nil {
+	fn(api.ProgressResponse{Status: "converting model"})
+	if err := convert.Convert(convert.NewZipReader(r, p, 32<<20), t); err != nil {
 		return nil, err
 	}
 
-	if _, err := temp.Seek(0, io.SeekStart); err != nil {
+	if _, err := t.Seek(0, io.SeekStart); err != nil {
 		return nil, err
 	}
 
-	layer, err := NewLayer(temp, "application/vnd.ollama.image.model")
+	layer, err := NewLayer(t, "application/vnd.ollama.image.model")
 	if err != nil {
 		return nil, err
 	}

server/model_test.go

@@ -1,16 +1,11 @@
 package server
 
 import (
-	"archive/zip"
 	"bytes"
 	"encoding/json"
-	"errors"
 	"fmt"
-	"io"
 	"os"
 	"path/filepath"
-	"slices"
-	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -18,103 +13,6 @@ import (
 	"github.com/ollama/ollama/template"
 )
 
-func createZipFile(t *testing.T, name string) *os.File {
-	t.Helper()
-
-	f, err := os.CreateTemp(t.TempDir(), "")
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	zf := zip.NewWriter(f)
-	defer zf.Close()
-
-	zh, err := zf.CreateHeader(&zip.FileHeader{Name: name})
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	if _, err := io.Copy(zh, bytes.NewReader([]byte(""))); err != nil {
-		t.Fatal(err)
-	}
-
-	return f
-}
-
-func TestExtractFromZipFile(t *testing.T) {
-	cases := []struct {
-		name   string
-		expect []string
-		err    error
-	}{
-		{
-			name:   "good",
-			expect: []string{"good"},
-		},
-		{
-			name:   strings.Join([]string{"path", "..", "to", "good"}, string(os.PathSeparator)),
-			expect: []string{filepath.Join("to", "good")},
-		},
-		{
-			name:   strings.Join([]string{"path", "..", "to", "..", "good"}, string(os.PathSeparator)),
-			expect: []string{"good"},
-		},
-		{
-			name:   strings.Join([]string{"path", "to", "..", "..", "good"}, string(os.PathSeparator)),
-			expect: []string{"good"},
-		},
-		{
-			name: strings.Join([]string{"..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "..", "bad"}, string(os.PathSeparator)),
-			err:  zip.ErrInsecurePath,
-		},
-		{
-			name: strings.Join([]string{"path", "..", "..", "to", "bad"}, string(os.PathSeparator)),
-			err:  zip.ErrInsecurePath,
-		},
-	}
-
-	for _, tt := range cases {
-		t.Run(tt.name, func(t *testing.T) {
-			f := createZipFile(t, tt.name)
-			defer f.Close()
-
-			tempDir := t.TempDir()
-			if err := extractFromZipFile(tempDir, f, func(api.ProgressResponse) {}); !errors.Is(err, tt.err) {
-				t.Fatal(err)
-			}
-
-			var matches []string
-			if err := filepath.Walk(tempDir, func(p string, fi os.FileInfo, err error) error {
-				if err != nil {
-					return err
-				}
-
-				if !fi.IsDir() {
-					matches = append(matches, p)
-				}
-
-				return nil
-			}); err != nil {
-				t.Fatal(err)
-			}
-
-			var actual []string
-			for _, match := range matches {
-				rel, err := filepath.Rel(tempDir, match)
-				if err != nil {
-					t.Error(err)
-				}
-
-				actual = append(actual, rel)
-			}
-
-			if !slices.Equal(actual, tt.expect) {
-				t.Fatalf("expected %d files, got %d", len(tt.expect), len(matches))
-			}
-		})
-	}
-}
-
 func readFile(t *testing.T, base, name string) *bytes.Buffer {
 	t.Helper()