changes

auth/auth.go

@@ -7,7 +7,6 @@ import (
 	"crypto/rand"
 	"encoding/base64"
 	"encoding/pem"
-	"errors"
 	"fmt"
 	"io"
 	"log/slog"
@@ -27,7 +26,7 @@ func privateKey() (ssh.Signer, error) {
 
 	keyPath := filepath.Join(home, ".ollama", defaultPrivateKey)
 	privateKeyFile, err := os.ReadFile(keyPath)
-	if errors.Is(err, os.ErrNotExist) {
+	if os.IsNotExist(err) {
 		err := initializeKeypair()
 		if err != nil {
 			return nil, err
@@ -51,7 +50,7 @@ func GetPublicKey() (ssh.PublicKey, error) {
 
 	pubkeyPath := filepath.Join(home, ".ollama", defaultPrivateKey+".pub")
 	pubKeyFile, err := os.ReadFile(pubkeyPath)
-	if errors.Is(err, os.ErrNotExist) {
+	if os.IsNotExist(err) {
 		// try from privateKey
 		privateKey, err := privateKey()
 		if err != nil {
@@ -114,7 +113,7 @@ func initializeKeypair() error {
 	pubKeyPath := filepath.Join(home, ".ollama", "id_ed25519.pub")
 
 	_, err = os.Stat(privKeyPath)
-	if errors.Is(err, os.ErrNotExist) {
+	if os.IsNotExist(err) {
 		fmt.Printf("Couldn't find '%s'. Generating new private key.\n", privKeyPath)
 		cryptoPublicKey, cryptoPrivateKey, err := ed25519.GenerateKey(rand.Reader)
 		if err != nil {

cmd/cmd.go

@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"context"
 	"crypto/sha256"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -110,7 +111,7 @@ func CreateHandler(cmd *cobra.Command, args []string) error {
 				path = tempfile
 			}
 
-			digest, err := createBlob(cmd, path)
+			digest, err := createBlob(cmd, client, path)
 			if err != nil {
 				return err
 			}
@@ -263,7 +264,7 @@ func tempZipFiles(path string) (string, error) {
 
 var ErrBlobExists = errors.New("blob exists")
 
-func createBlob(cmd *cobra.Command, path string) (string, error) {
+func createBlob(cmd *cobra.Command, client *api.Client, path string) (string, error) {
 	bin, err := os.Open(path)
 	if err != nil {
 		return "", err
@@ -281,23 +282,40 @@ func createBlob(cmd *cobra.Command, path string) (string, error) {
 
 	digest := fmt.Sprintf("sha256:%x", hash.Sum(nil))
 
-	// Use our new CreateBlob request which will include the file path
-	// The server checks for that file and if the server is local, it will copy the file over
-	// If the local copy fails, the server will continue to the default local copy
-	// If that fails, it will continue with the server POST
-	err = CreateBlob(cmd.Context(), path, digest, bin)
+	// We check if we can find the models directory locally
+	// If we can, we return the path to the directory
+	// If we can't, we return an error
+	// If the blob exists already, we return the digest
+	dest, err := getLocalPath(cmd.Context(), digest)
+
 	if errors.Is(err, ErrBlobExists) {
 		return digest, nil
 	}
 
-	if err != nil {
-		return "", err
+	// Successfully found the model directory
+	if err == nil {
+		// Copy blob in via OS specific copy
+		// Linux errors out to use io.copy
+		err = localCopy(path, dest)
+		if err == nil {
+			return digest, nil
+		}
+
+		// Default copy using io.copy
+		err = defaultCopy(path, dest)
+		if err == nil {
+			return digest, nil
+		}
 	}
 
+	// If at any point copying the blob over locally fails, we default to the copy through the server
+	if err = client.CreateBlob(cmd.Context(), digest, bin); err != nil {
+		return "", err
+	}
 	return digest, nil
 }
 
-func CreateBlob(ctx context.Context, src, digest string, r *os.File) (error) {
+func getLocalPath(ctx context.Context, digest string) (string, error) {
 	ollamaHost := envconfig.Host
 
 	client := http.DefaultClient
@@ -306,37 +324,75 @@ func CreateBlob(ctx context.Context, src, digest string, r *os.File) (error) {
 		Host:   net.JoinHostPort(ollamaHost.Host, ollamaHost.Port),
 	}
 
+	data, err := json.Marshal(digest)
+	if err != nil {
+		return "", err
+	}
+
+	reqBody := bytes.NewReader(data)
 	path := fmt.Sprintf("/api/blobs/%s", digest)
 	requestURL := base.JoinPath(path)
-	request, err := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), r)
+	request, err := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), reqBody)
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	authz, err := api.Authorization(ctx, request)
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	request.Header.Set("Authorization", authz)
 	request.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
-	request.Header.Set("X-Ollama-File", src)
+	request.Header.Set("X-Redirect-Create", "1")
 
 	resp, err := client.Do(request)
 	if err != nil {
-		return err
+		return "", err
 	}
 	defer resp.Body.Close()
 
-	if resp.StatusCode == http.StatusCreated {
-		return nil
+	if resp.StatusCode == http.StatusTemporaryRedirect {
+		dest := resp.Header.Get("LocalLocation")
+
+		return dest, nil
+	}
+	return "", ErrBlobExists
+}
+
+func defaultCopy(path string, dest string) error {
+	// This function should be called if the server is local
+	// It should find the model directory, copy the blob over, and return the digest
+	dirPath := filepath.Dir(dest)
+
+	if err := os.MkdirAll(dirPath, 0o755); err != nil {
+		return err
 	}
 
-	if resp.StatusCode == http.StatusOK {
-		return ErrBlobExists
+	// Copy blob over
+	sourceFile, err := os.Open(path)
+	if err != nil {
+		return fmt.Errorf("could not open source file: %v", err)
+	}
+	defer sourceFile.Close()
+
+	destFile, err := os.Create(dest)
+	if err != nil {
+		return fmt.Errorf("could not create destination file: %v", err)
+	}
+	defer destFile.Close()
+
+	_, err = io.CopyBuffer(destFile, sourceFile, make([]byte, 4*1024*1024))
+	if err != nil {
+		return fmt.Errorf("error copying file: %v", err)
 	}
 
-	return err
+	err = destFile.Sync()
+	if err != nil {
+		return fmt.Errorf("error flushing file: %v", err)
+	}
+
+	return nil
 }
 
 func RunHandler(cmd *cobra.Command, args []string) error {

cmd/copy_darwin.go

@@ -1,4 +1,4 @@
-package server
+package cmd
 
 import (
 	"os"

cmd/copy_linux.go

@@ -1,4 +1,4 @@
-package server
+package cmd
 
 import "errors"
 

cmd/copy_windows.go

@@ -1,7 +1,7 @@
 //go:build windows
 // +build windows
 
-package server
+package cmd
 
 import (
 	"os"

server/routes.go

@@ -942,13 +942,10 @@ func (s *Server) CreateBlobHandler(c *gin.Context) {
 		c.Status(http.StatusOK)
 		return
 	}
-
-	if c.GetHeader("X-Ollama-File") != "" && s.isLocal(c) {
-		err = localBlobCopy(c.GetHeader("X-Ollama-File"), path)
-		if err == nil {
-			c.Status(http.StatusCreated)
-			return
-		}
+	if c.GetHeader("X-Redirect-Create") == "1" && s.isLocal(c) {
+		c.Header("LocalLocation", path)
+		c.Status(http.StatusTemporaryRedirect)
+		return
 	}
 
 	layer, err := NewLayer(c.Request.Body, "")
@@ -965,25 +962,6 @@ func (s *Server) CreateBlobHandler(c *gin.Context) {
 	c.Status(http.StatusCreated)
 }
 
-func localBlobCopy (src, dest string) error {
-	_, err := os.Stat(src)
-	if err != nil {
-		return err
-	}
-
-	err = localCopy(src, dest)
-	if err == nil {
-		return nil
-	}
-
-	err = defaultCopy(src, dest)
-	if err == nil {
-		return nil
-	}
-
-	return fmt.Errorf("failed to copy blob")
-}
-
 func (s *Server) isLocal(c *gin.Context) bool {
 	if authz := c.GetHeader("Authorization"); authz != "" {
 		parts := strings.Split(authz, ":")
@@ -1032,41 +1010,6 @@ func (s *Server) isLocal(c *gin.Context) bool {
 	return false
 }
 
-func defaultCopy(path string, dest string) error {
-	// This function should be called if the server is local
-	// It should find the model directory, copy the blob over, and return the digest
-	dirPath := filepath.Dir(dest)
-
-	if err := os.MkdirAll(dirPath, 0o755); err != nil {
-		return err
-	}
-
-	// Copy blob over
-	sourceFile, err := os.Open(path)
-	if err != nil {
-		return fmt.Errorf("could not open source file: %v", err)
-	}
-	defer sourceFile.Close()
-
-	destFile, err := os.Create(dest)
-	if err != nil {
-		return fmt.Errorf("could not create destination file: %v", err)
-	}
-	defer destFile.Close()
-
-	_, err = io.CopyBuffer(destFile, sourceFile, make([]byte, 4*1024*1024))
-	if err != nil {
-		return fmt.Errorf("error copying file: %v", err)
-	}
-
-	err = destFile.Sync()
-	if err != nil {
-		return fmt.Errorf("error flushing file: %v", err)
-	}
-
-	return nil
-}
-
 func isLocalIP(ip netip.Addr) bool {
 	if interfaces, err := net.Interfaces(); err == nil {
 		for _, iface := range interfaces {

server/routes_test.go

@@ -535,7 +535,6 @@ func TestIsLocalReal(t *testing.T) {
 	gin.SetMode(gin.TestMode)
 	clientPubLoc := t.TempDir()
 	t.Setenv("HOME", clientPubLoc)
-	t.Setenv("USERPROFILE", clientPubLoc)
 
 	_, err := auth.GetPublicKey()
 	if err != nil {
@@ -573,7 +572,6 @@ func TestIsLocalReal(t *testing.T) {
 	t.Run("different server pubkey", func(t *testing.T) {
 		serverPubLoc := t.TempDir()
 		t.Setenv("HOME", serverPubLoc)
-		t.Setenv("USERPROFILE", serverPubLoc)
 		_, err := auth.GetPublicKey()
 		if err != nil {
 			t.Fatal(err)