| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | package server | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | import ( | 
					
						
							|  |  |  | 	"bytes" | 
					
						
							|  |  |  | 	"crypto/sha256" | 
					
						
							|  |  |  | 	"encoding/json" | 
					
						
							|  |  |  | 	"errors" | 
					
						
							|  |  |  | 	"fmt" | 
					
						
							|  |  |  | 	"io" | 
					
						
							|  |  |  | 	"io/ioutil" | 
					
						
							|  |  |  | 	"log" | 
					
						
							|  |  |  | 	"net/http" | 
					
						
							|  |  |  | 	"os" | 
					
						
							|  |  |  | 	"path" | 
					
						
							|  |  |  | 	"path/filepath" | 
					
						
							| 
									
										
										
										
											2023-07-18 03:08:10 +08:00
										 |  |  | 	"reflect" | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	"strconv" | 
					
						
							|  |  |  | 	"strings" | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 	"text/template" | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 	"github.com/jmorganca/ollama/api" | 
					
						
							|  |  |  | 	"github.com/jmorganca/ollama/parser" | 
					
						
							|  |  |  | ) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | type RegistryOptions struct { | 
					
						
							|  |  |  | 	Insecure bool | 
					
						
							|  |  |  | 	Username string | 
					
						
							|  |  |  | 	Password string | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | type Model struct { | 
					
						
							|  |  |  | 	Name      string `json:"name"` | 
					
						
							|  |  |  | 	ModelPath string | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 	Template  string | 
					
						
							|  |  |  | 	System    string | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	Options   api.Options | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | func (m *Model) Prompt(request api.GenerateRequest) (string, error) { | 
					
						
							|  |  |  | 	tmpl, err := template.New("").Parse(m.Template) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return "", err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	var vars struct { | 
					
						
							| 
									
										
										
										
											2023-07-20 14:22:19 +08:00
										 |  |  | 		First  bool | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 		System string | 
					
						
							|  |  |  | 		Prompt string | 
					
						
							| 
									
										
										
										
											2023-07-20 10:43:00 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 		// deprecated: versions <= 0.0.7 used this to omit the system prompt
 | 
					
						
							|  |  |  | 		Context []int | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 11:45:32 +08:00
										 |  |  | 	vars.First = len(request.Context) == 0 | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 	vars.System = m.System | 
					
						
							|  |  |  | 	vars.Prompt = request.Prompt | 
					
						
							| 
									
										
										
										
											2023-07-20 14:22:19 +08:00
										 |  |  | 	vars.Context = request.Context | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 	var sb strings.Builder | 
					
						
							|  |  |  | 	if err := tmpl.Execute(&sb, vars); err != nil { | 
					
						
							|  |  |  | 		return "", err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	return sb.String(), nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | type ManifestV2 struct { | 
					
						
							|  |  |  | 	SchemaVersion int      `json:"schemaVersion"` | 
					
						
							|  |  |  | 	MediaType     string   `json:"mediaType"` | 
					
						
							|  |  |  | 	Config        Layer    `json:"config"` | 
					
						
							|  |  |  | 	Layers        []*Layer `json:"layers"` | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | type Layer struct { | 
					
						
							|  |  |  | 	MediaType string `json:"mediaType"` | 
					
						
							|  |  |  | 	Digest    string `json:"digest"` | 
					
						
							|  |  |  | 	Size      int    `json:"size"` | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | type LayerReader struct { | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	Layer | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | 	io.Reader | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | type ConfigV2 struct { | 
					
						
							|  |  |  | 	Architecture string `json:"architecture"` | 
					
						
							|  |  |  | 	OS           string `json:"os"` | 
					
						
							|  |  |  | 	RootFS       RootFS `json:"rootfs"` | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | type RootFS struct { | 
					
						
							|  |  |  | 	Type    string   `json:"type"` | 
					
						
							|  |  |  | 	DiffIDs []string `json:"diff_ids"` | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 00:09:45 +08:00
										 |  |  | func (m *ManifestV2) GetTotalSize() int { | 
					
						
							|  |  |  | 	var total int | 
					
						
							|  |  |  | 	for _, layer := range m.Layers { | 
					
						
							|  |  |  | 		total += layer.Size | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	total += m.Config.Size | 
					
						
							|  |  |  | 	return total | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | func GetManifest(mp ModelPath) (*ManifestV2, error) { | 
					
						
							|  |  |  | 	fp, err := mp.GetManifestPath(false) | 
					
						
							| 
									
										
										
										
											2023-07-18 02:03:55 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return nil, err | 
					
						
							|  |  |  | 	} | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 14:02:12 +08:00
										 |  |  | 	if _, err = os.Stat(fp); err != nil { | 
					
						
							|  |  |  | 		return nil, err | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	var manifest *ManifestV2 | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 	bts, err := os.ReadFile(fp) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return nil, fmt.Errorf("couldn't open file '%s'", fp) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 	if err := json.Unmarshal(bts, &manifest); err != nil { | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		return nil, err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	return manifest, nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | func GetModel(name string) (*Model, error) { | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 	mp := ParseModelPath(name) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	manifest, err := GetManifest(mp) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return nil, err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	model := &Model{ | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 		Name: mp.GetFullTagname(), | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	for _, layer := range manifest.Layers { | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 		filename, err := GetBlobsPath(layer.Digest) | 
					
						
							| 
									
										
										
										
											2023-07-18 02:03:55 +08:00
										 |  |  | 		if err != nil { | 
					
						
							|  |  |  | 			return nil, err | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		switch layer.MediaType { | 
					
						
							|  |  |  | 		case "application/vnd.ollama.image.model": | 
					
						
							|  |  |  | 			model.ModelPath = filename | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 		case "application/vnd.ollama.image.template": | 
					
						
							|  |  |  | 			bts, err := os.ReadFile(filename) | 
					
						
							|  |  |  | 			if err != nil { | 
					
						
							|  |  |  | 				return nil, err | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 			model.Template = string(bts) | 
					
						
							|  |  |  | 		case "application/vnd.ollama.image.system": | 
					
						
							|  |  |  | 			bts, err := os.ReadFile(filename) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 			if err != nil { | 
					
						
							|  |  |  | 				return nil, err | 
					
						
							|  |  |  | 			} | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 			model.System = string(bts) | 
					
						
							| 
									
										
										
										
											2023-07-20 10:43:00 +08:00
										 |  |  | 		case "application/vnd.ollama.image.prompt": | 
					
						
							|  |  |  | 			bts, err := os.ReadFile(filename) | 
					
						
							|  |  |  | 			if err != nil { | 
					
						
							|  |  |  | 				return nil, err | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 			model.Template = string(bts) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		case "application/vnd.ollama.image.params": | 
					
						
							| 
									
										
										
										
											2023-07-18 03:08:10 +08:00
										 |  |  | 			params, err := os.Open(filename) | 
					
						
							|  |  |  | 			if err != nil { | 
					
						
							|  |  |  | 				return nil, err | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 			defer params.Close() | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 			var opts api.Options | 
					
						
							| 
									
										
										
										
											2023-07-18 03:08:10 +08:00
										 |  |  | 			if err = json.NewDecoder(params).Decode(&opts); err != nil { | 
					
						
							|  |  |  | 				return nil, err | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 			model.Options = opts | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	return model, nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-20 12:55:15 +08:00
										 |  |  | func CreateModel(name string, path string, fn func(status string)) error { | 
					
						
							|  |  |  | 	mf, err := os.Open(path) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		fn(fmt.Sprintf("couldn't open modelfile '%s'", path)) | 
					
						
							|  |  |  | 		return fmt.Errorf("failed to open file: %w", err) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	} | 
					
						
							| 
									
										
										
										
											2023-07-20 12:55:15 +08:00
										 |  |  | 	defer mf.Close() | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 	fn("parsing modelfile") | 
					
						
							|  |  |  | 	commands, err := parser.Parse(mf) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | 	var layers []*LayerReader | 
					
						
							| 
									
										
										
										
											2023-07-18 03:08:10 +08:00
										 |  |  | 	params := make(map[string]string) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 	for _, c := range commands { | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 		log.Printf("[%s] - %s\n", c.Name, c.Args) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		switch c.Name { | 
					
						
							|  |  |  | 		case "model": | 
					
						
							|  |  |  | 			fn("looking for model") | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 			mf, err := GetManifest(ParseModelPath(c.Args)) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 			if err != nil { | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 				fp := c.Args | 
					
						
							| 
									
										
										
										
											2023-07-20 12:55:15 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 				// If filePath starts with ~/, replace it with the user's home directory.
 | 
					
						
							|  |  |  | 				if strings.HasPrefix(fp, "~/") { | 
					
						
							|  |  |  | 					parts := strings.Split(fp, "/") | 
					
						
							|  |  |  | 					home, err := os.UserHomeDir() | 
					
						
							|  |  |  | 					if err != nil { | 
					
						
							|  |  |  | 						return fmt.Errorf("failed to open file: %v", err) | 
					
						
							|  |  |  | 					} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 					fp = filepath.Join(home, filepath.Join(parts[1:]...)) | 
					
						
							|  |  |  | 				} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 				// If filePath is not an absolute path, make it relative to the modelfile path
 | 
					
						
							|  |  |  | 				if !filepath.IsAbs(fp) { | 
					
						
							|  |  |  | 					fp = filepath.Join(filepath.Dir(path), fp) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 				} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 				fn("creating model layer") | 
					
						
							|  |  |  | 				file, err := os.Open(fp) | 
					
						
							|  |  |  | 				if err != nil { | 
					
						
							|  |  |  | 					return fmt.Errorf("failed to open file: %v", err) | 
					
						
							|  |  |  | 				} | 
					
						
							|  |  |  | 				defer file.Close() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 				l, err := CreateLayer(file) | 
					
						
							|  |  |  | 				if err != nil { | 
					
						
							|  |  |  | 					return fmt.Errorf("failed to create layer: %v", err) | 
					
						
							|  |  |  | 				} | 
					
						
							|  |  |  | 				l.MediaType = "application/vnd.ollama.image.model" | 
					
						
							|  |  |  | 				layers = append(layers, l) | 
					
						
							|  |  |  | 			} else { | 
					
						
							|  |  |  | 				log.Printf("manifest = %#v", mf) | 
					
						
							|  |  |  | 				for _, l := range mf.Layers { | 
					
						
							|  |  |  | 					newLayer, err := GetLayerWithBufferFromLayer(l) | 
					
						
							|  |  |  | 					if err != nil { | 
					
						
							|  |  |  | 						return err | 
					
						
							|  |  |  | 					} | 
					
						
							|  |  |  | 					layers = append(layers, newLayer) | 
					
						
							|  |  |  | 				} | 
					
						
							|  |  |  | 			} | 
					
						
							| 
									
										
										
										
											2023-07-20 16:13:30 +08:00
										 |  |  | 		case "license", "template", "system", "prompt": | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 			fn(fmt.Sprintf("creating %s layer", c.Name)) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 			// remove the prompt layer if one exists
 | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 			mediaType := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name) | 
					
						
							|  |  |  | 			layers = removeLayerFromLayers(layers, mediaType) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 			layer, err := CreateLayer(strings.NewReader(c.Args)) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 			if err != nil { | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 				return err | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 			} | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 			layer.MediaType = mediaType | 
					
						
							|  |  |  | 			layers = append(layers, layer) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		default: | 
					
						
							| 
									
										
										
										
											2023-07-18 05:21:27 +08:00
										 |  |  | 			params[c.Name] = c.Args | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		} | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	// Create a single layer for the parameters
 | 
					
						
							| 
									
										
										
										
											2023-07-18 03:08:10 +08:00
										 |  |  | 	if len(params) > 0 { | 
					
						
							|  |  |  | 		fn("creating parameter layer") | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		layers = removeLayerFromLayers(layers, "application/vnd.ollama.image.params") | 
					
						
							| 
									
										
										
										
											2023-07-18 03:08:10 +08:00
										 |  |  | 		paramData, err := paramsToReader(params) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		if err != nil { | 
					
						
							|  |  |  | 			return fmt.Errorf("couldn't create params json: %v", err) | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 		l, err := CreateLayer(paramData) | 
					
						
							|  |  |  | 		if err != nil { | 
					
						
							|  |  |  | 			return fmt.Errorf("failed to create layer: %v", err) | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 		l.MediaType = "application/vnd.ollama.image.params" | 
					
						
							|  |  |  | 		layers = append(layers, l) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	digests, err := getLayerDigests(layers) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	var manifestLayers []*Layer | 
					
						
							|  |  |  | 	for _, l := range layers { | 
					
						
							|  |  |  | 		manifestLayers = append(manifestLayers, &l.Layer) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	// Create a layer for the config object
 | 
					
						
							|  |  |  | 	fn("creating config layer") | 
					
						
							|  |  |  | 	cfg, err := createConfigLayer(digests) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	layers = append(layers, cfg) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	err = SaveLayers(layers, fn, false) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	// Create the manifest
 | 
					
						
							|  |  |  | 	fn("writing manifest") | 
					
						
							|  |  |  | 	err = CreateManifest(name, cfg, manifestLayers) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	fn("success") | 
					
						
							|  |  |  | 	return nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | func removeLayerFromLayers(layers []*LayerReader, mediaType string) []*LayerReader { | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	j := 0 | 
					
						
							|  |  |  | 	for _, l := range layers { | 
					
						
							|  |  |  | 		if l.MediaType != mediaType { | 
					
						
							|  |  |  | 			layers[j] = l | 
					
						
							|  |  |  | 			j++ | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	return layers[:j] | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | func SaveLayers(layers []*LayerReader, fn func(status string), force bool) error { | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	// Write each of the layers to disk
 | 
					
						
							|  |  |  | 	for _, layer := range layers { | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 		fp, err := GetBlobsPath(layer.Digest) | 
					
						
							| 
									
										
										
										
											2023-07-18 02:03:55 +08:00
										 |  |  | 		if err != nil { | 
					
						
							|  |  |  | 			return err | 
					
						
							|  |  |  | 		} | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 		_, err = os.Stat(fp) | 
					
						
							|  |  |  | 		if os.IsNotExist(err) || force { | 
					
						
							|  |  |  | 			fn(fmt.Sprintf("writing layer %s", layer.Digest)) | 
					
						
							|  |  |  | 			out, err := os.Create(fp) | 
					
						
							|  |  |  | 			if err != nil { | 
					
						
							|  |  |  | 				log.Printf("couldn't create %s", fp) | 
					
						
							|  |  |  | 				return err | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 			defer out.Close() | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | 			if _, err = io.Copy(out, layer.Reader); err != nil { | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 				return err | 
					
						
							|  |  |  | 			} | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		} else { | 
					
						
							|  |  |  | 			fn(fmt.Sprintf("using already created layer %s", layer.Digest)) | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	return nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | func CreateManifest(name string, cfg *LayerReader, layers []*Layer) error { | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 	mp := ParseModelPath(name) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	manifest := ManifestV2{ | 
					
						
							|  |  |  | 		SchemaVersion: 2, | 
					
						
							|  |  |  | 		MediaType:     "application/vnd.docker.distribution.manifest.v2+json", | 
					
						
							|  |  |  | 		Config: Layer{ | 
					
						
							|  |  |  | 			MediaType: cfg.MediaType, | 
					
						
							|  |  |  | 			Size:      cfg.Size, | 
					
						
							|  |  |  | 			Digest:    cfg.Digest, | 
					
						
							|  |  |  | 		}, | 
					
						
							|  |  |  | 		Layers: layers, | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	manifestJSON, err := json.Marshal(manifest) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 	fp, err := mp.GetManifestPath(true) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							| 
									
										
										
										
											2023-07-18 02:03:55 +08:00
										 |  |  | 	return os.WriteFile(fp, manifestJSON, 0o644) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | func GetLayerWithBufferFromLayer(layer *Layer) (*LayerReader, error) { | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 	fp, err := GetBlobsPath(layer.Digest) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return nil, err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	file, err := os.Open(fp) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return nil, fmt.Errorf("could not open blob: %w", err) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	defer file.Close() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	newLayer, err := CreateLayer(file) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return nil, err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	newLayer.MediaType = layer.MediaType | 
					
						
							|  |  |  | 	return newLayer, nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | func paramsToReader(params map[string]string) (io.ReadSeeker, error) { | 
					
						
							| 
									
										
										
										
											2023-07-18 03:08:10 +08:00
										 |  |  | 	opts := api.DefaultOptions() | 
					
						
							|  |  |  | 	typeOpts := reflect.TypeOf(opts) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	// build map of json struct tags
 | 
					
						
							|  |  |  | 	jsonOpts := make(map[string]reflect.StructField) | 
					
						
							|  |  |  | 	for _, field := range reflect.VisibleFields(typeOpts) { | 
					
						
							|  |  |  | 		jsonTag := strings.Split(field.Tag.Get("json"), ",")[0] | 
					
						
							|  |  |  | 		if jsonTag != "" { | 
					
						
							|  |  |  | 			jsonOpts[jsonTag] = field | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	valueOpts := reflect.ValueOf(&opts).Elem() | 
					
						
							|  |  |  | 	// iterate params and set values based on json struct tags
 | 
					
						
							|  |  |  | 	for key, val := range params { | 
					
						
							|  |  |  | 		if opt, ok := jsonOpts[key]; ok { | 
					
						
							|  |  |  | 			field := valueOpts.FieldByName(opt.Name) | 
					
						
							|  |  |  | 			if field.IsValid() && field.CanSet() { | 
					
						
							|  |  |  | 				switch field.Kind() { | 
					
						
							|  |  |  | 				case reflect.Float32: | 
					
						
							|  |  |  | 					floatVal, err := strconv.ParseFloat(val, 32) | 
					
						
							|  |  |  | 					if err != nil { | 
					
						
							|  |  |  | 						return nil, fmt.Errorf("invalid float value %s", val) | 
					
						
							|  |  |  | 					} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 					field.SetFloat(floatVal) | 
					
						
							|  |  |  | 				case reflect.Int: | 
					
						
							|  |  |  | 					intVal, err := strconv.ParseInt(val, 10, 0) | 
					
						
							|  |  |  | 					if err != nil { | 
					
						
							|  |  |  | 						return nil, fmt.Errorf("invalid int value %s", val) | 
					
						
							|  |  |  | 					} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 					field.SetInt(intVal) | 
					
						
							|  |  |  | 				case reflect.Bool: | 
					
						
							|  |  |  | 					boolVal, err := strconv.ParseBool(val) | 
					
						
							|  |  |  | 					if err != nil { | 
					
						
							|  |  |  | 						return nil, fmt.Errorf("invalid bool value %s", val) | 
					
						
							|  |  |  | 					} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 					field.SetBool(boolVal) | 
					
						
							|  |  |  | 				case reflect.String: | 
					
						
							|  |  |  | 					field.SetString(val) | 
					
						
							|  |  |  | 				default: | 
					
						
							|  |  |  | 					return nil, fmt.Errorf("unknown type %s for %s", field.Kind(), key) | 
					
						
							|  |  |  | 				} | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	bts, err := json.Marshal(opts) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return nil, err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-18 03:08:10 +08:00
										 |  |  | 	return bytes.NewReader(bts), nil | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | func getLayerDigests(layers []*LayerReader) ([]string, error) { | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	var digests []string | 
					
						
							|  |  |  | 	for _, l := range layers { | 
					
						
							|  |  |  | 		if l.Digest == "" { | 
					
						
							|  |  |  | 			return nil, fmt.Errorf("layer is missing a digest") | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 		digests = append(digests, l.Digest) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	return digests, nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | // CreateLayer creates a Layer object from a given file
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | func CreateLayer(f io.ReadSeeker) (*LayerReader, error) { | 
					
						
							|  |  |  | 	digest, size := GetSHA256Digest(f) | 
					
						
							|  |  |  | 	f.Seek(0, 0) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | 	layer := &LayerReader{ | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		Layer: Layer{ | 
					
						
							|  |  |  | 			MediaType: "application/vnd.docker.image.rootfs.diff.tar", | 
					
						
							|  |  |  | 			Digest:    digest, | 
					
						
							|  |  |  | 			Size:      size, | 
					
						
							|  |  |  | 		}, | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | 		Reader: f, | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	return layer, nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 14:02:12 +08:00
										 |  |  | func DeleteModel(name string) error { | 
					
						
							| 
									
										
										
										
											2023-07-21 07:09:23 +08:00
										 |  |  | 	mp := ParseModelPath(name) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	manifest, err := GetManifest(mp) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	deleteMap := make(map[string]bool) | 
					
						
							|  |  |  | 	for _, layer := range manifest.Layers { | 
					
						
							|  |  |  | 		deleteMap[layer.Digest] = true | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	deleteMap[manifest.Config.Digest] = true | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	fp, err := GetManifestPath() | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	err = filepath.Walk(fp, func(path string, info os.FileInfo, err error) error { | 
					
						
							|  |  |  | 		if err != nil { | 
					
						
							|  |  |  | 			return err | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 		if !info.IsDir() { | 
					
						
							|  |  |  | 			path := path[len(fp)+1:] | 
					
						
							|  |  |  | 			slashIndex := strings.LastIndex(path, "/") | 
					
						
							|  |  |  | 			if slashIndex == -1 { | 
					
						
							|  |  |  | 				return nil | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 			tag := path[:slashIndex] + ":" + path[slashIndex+1:] | 
					
						
							|  |  |  | 			fmp := ParseModelPath(tag) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 			// skip the manifest we're trying to delete
 | 
					
						
							|  |  |  | 			if mp.GetFullTagname() == fmp.GetFullTagname() { | 
					
						
							|  |  |  | 				return nil | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 			// save (i.e. delete from the deleteMap) any files used in other manifests
 | 
					
						
							|  |  |  | 			manifest, err := GetManifest(fmp) | 
					
						
							|  |  |  | 			if err != nil { | 
					
						
							|  |  |  | 				log.Printf("skipping file: %s", fp) | 
					
						
							|  |  |  | 				return nil | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 			for _, layer := range manifest.Layers { | 
					
						
							|  |  |  | 				delete(deleteMap, layer.Digest) | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 			delete(deleteMap, manifest.Config.Digest) | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 		return nil | 
					
						
							|  |  |  | 	}) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	// only delete the files which are still in the deleteMap
 | 
					
						
							|  |  |  | 	for k, v := range deleteMap { | 
					
						
							|  |  |  | 		if v { | 
					
						
							| 
									
										
										
										
											2023-07-22 08:30:40 +08:00
										 |  |  | 			fp, err := GetBlobsPath(k) | 
					
						
							| 
									
										
										
										
											2023-07-21 07:09:23 +08:00
										 |  |  | 			if err != nil { | 
					
						
							| 
									
										
										
										
											2023-07-22 08:30:40 +08:00
										 |  |  | 				log.Printf("couldn't get file path for '%s': %v", k, err) | 
					
						
							|  |  |  | 				continue | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 			if err := os.Remove(fp); err != nil { | 
					
						
							|  |  |  | 				log.Printf("couldn't remove file '%s': %v", fp, err) | 
					
						
							| 
									
										
										
										
											2023-07-21 07:09:23 +08:00
										 |  |  | 				continue | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	fp, err = mp.GetManifestPath(false) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	err = os.Remove(fp) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		log.Printf("couldn't remove manifest file '%s': %v", fp, err) | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	return nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | func PushModel(name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error { | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 	mp := ParseModelPath(name) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 	fn(api.ProgressResponse{Status: "retrieving manifest"}) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 	manifest, err := GetManifest(mp) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"}) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	var layers []*Layer | 
					
						
							|  |  |  | 	for _, layer := range manifest.Layers { | 
					
						
							|  |  |  | 		layers = append(layers, layer) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	layers = append(layers, &manifest.Config) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	for _, layer := range layers { | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | 		exists, err := checkBlobExistence(mp, layer.Digest, regOpts) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		if err != nil { | 
					
						
							|  |  |  | 			return err | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 		if exists { | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 			fn(api.ProgressResponse{ | 
					
						
							|  |  |  | 				Status:    "using existing layer", | 
					
						
							|  |  |  | 				Digest:    layer.Digest, | 
					
						
							| 
									
										
										
										
											2023-07-23 08:31:26 +08:00
										 |  |  | 				Total:     layer.Size, | 
					
						
							|  |  |  | 				Completed: layer.Size, | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 			}) | 
					
						
							| 
									
										
										
										
											2023-07-23 08:31:26 +08:00
										 |  |  | 			log.Printf("Layer %s already exists", layer.Digest) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 			continue | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 		fn(api.ProgressResponse{ | 
					
						
							| 
									
										
										
										
											2023-07-23 08:31:26 +08:00
										 |  |  | 			Status: "starting upload", | 
					
						
							|  |  |  | 			Digest: layer.Digest, | 
					
						
							|  |  |  | 			Total:  layer.Size, | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 		}) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | 		location, err := startUpload(mp, regOpts) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		if err != nil { | 
					
						
							|  |  |  | 			log.Printf("couldn't start upload: %v", err) | 
					
						
							|  |  |  | 			return err | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-23 08:31:26 +08:00
										 |  |  | 		err = uploadBlobChunked(mp, location, layer, regOpts, fn) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		if err != nil { | 
					
						
							|  |  |  | 			log.Printf("error uploading blob: %v", err) | 
					
						
							|  |  |  | 			return err | 
					
						
							|  |  |  | 		} | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-23 08:31:26 +08:00
										 |  |  | 	fn(api.ProgressResponse{Status: "pushing manifest"}) | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | 	url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	headers := map[string]string{ | 
					
						
							|  |  |  | 		"Content-Type": "application/vnd.docker.distribution.manifest.v2+json", | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	manifestJSON, err := json.Marshal(manifest) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | 	resp, err := makeRequest("PUT", url, headers, bytes.NewReader(manifestJSON), regOpts) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	defer resp.Body.Close() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	// Check for success: For a successful upload, the Docker registry will respond with a 201 Created
 | 
					
						
							|  |  |  | 	if resp.StatusCode != http.StatusCreated { | 
					
						
							|  |  |  | 		body, _ := io.ReadAll(resp.Body) | 
					
						
							|  |  |  | 		return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body)) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-23 08:31:26 +08:00
										 |  |  | 	fn(api.ProgressResponse{Status: "success"}) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 	return nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | func PullModel(name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error { | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 	mp := ParseModelPath(name) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 	fn(api.ProgressResponse{Status: "pulling manifest"}) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | 	manifest, err := pullModelManifest(mp, regOpts) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return fmt.Errorf("pull model manifest: %q", err) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	var layers []*Layer | 
					
						
							| 
									
										
										
										
											2023-07-21 02:18:00 +08:00
										 |  |  | 	layers = append(layers, manifest.Layers...) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	layers = append(layers, &manifest.Config) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	for _, layer := range layers { | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | 		if err := downloadBlob(mp, layer.Digest, regOpts, fn); err != nil { | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 			return err | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-21 02:44:05 +08:00
										 |  |  | 	fn(api.ProgressResponse{Status: "verifying sha256 digest"}) | 
					
						
							|  |  |  | 	for _, layer := range layers { | 
					
						
							|  |  |  | 		if err := verifyBlob(layer.Digest); err != nil { | 
					
						
							|  |  |  | 			return err | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 	fn(api.ProgressResponse{Status: "writing manifest"}) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-18 02:03:55 +08:00
										 |  |  | 	manifestJSON, err := json.Marshal(manifest) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 	fp, err := mp.GetManifestPath(true) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-21 02:18:00 +08:00
										 |  |  | 	err = os.WriteFile(fp, manifestJSON, 0o644) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		log.Printf("couldn't write to %s", fp) | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 	fn(api.ProgressResponse{Status: "success"}) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 	return nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | func pullModelManifest(mp ModelPath, regOpts *RegistryOptions) (*ManifestV2, error) { | 
					
						
							|  |  |  | 	url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	headers := map[string]string{ | 
					
						
							|  |  |  | 		"Accept": "application/vnd.docker.distribution.manifest.v2+json", | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | 	resp, err := makeRequest("GET", url, headers, nil, regOpts) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		log.Printf("couldn't get manifest: %v", err) | 
					
						
							|  |  |  | 		return nil, err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	defer resp.Body.Close() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	// Check for success: For a successful upload, the Docker registry will respond with a 201 Created
 | 
					
						
							|  |  |  | 	if resp.StatusCode != http.StatusOK { | 
					
						
							|  |  |  | 		body, _ := io.ReadAll(resp.Body) | 
					
						
							| 
									
										
										
										
											2023-07-18 03:33:39 +08:00
										 |  |  | 		return nil, fmt.Errorf("registry responded with code %d: %s", resp.StatusCode, body) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	var m *ManifestV2 | 
					
						
							|  |  |  | 	if err := json.NewDecoder(resp.Body).Decode(&m); err != nil { | 
					
						
							|  |  |  | 		return nil, err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	return m, err | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | func createConfigLayer(layers []string) (*LayerReader, error) { | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	// TODO change architecture and OS
 | 
					
						
							|  |  |  | 	config := ConfigV2{ | 
					
						
							|  |  |  | 		Architecture: "arm64", | 
					
						
							|  |  |  | 		OS:           "linux", | 
					
						
							|  |  |  | 		RootFS: RootFS{ | 
					
						
							|  |  |  | 			Type:    "layers", | 
					
						
							|  |  |  | 			DiffIDs: layers, | 
					
						
							|  |  |  | 		}, | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	configJSON, err := json.Marshal(config) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return nil, err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 15:47:55 +08:00
										 |  |  | 	digest, size := GetSHA256Digest(bytes.NewBuffer(configJSON)) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | 	layer := &LayerReader{ | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		Layer: Layer{ | 
					
						
							|  |  |  | 			MediaType: "application/vnd.docker.container.image.v1+json", | 
					
						
							|  |  |  | 			Digest:    digest, | 
					
						
							|  |  |  | 			Size:      size, | 
					
						
							|  |  |  | 		}, | 
					
						
							| 
									
										
										
										
											2023-07-19 15:47:55 +08:00
										 |  |  | 		Reader: bytes.NewBuffer(configJSON), | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	} | 
					
						
							|  |  |  | 	return layer, nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | // GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
 | 
					
						
							| 
									
										
										
										
											2023-07-19 08:14:12 +08:00
										 |  |  | func GetSHA256Digest(r io.Reader) (string, int) { | 
					
						
							|  |  |  | 	h := sha256.New() | 
					
						
							|  |  |  | 	n, err := io.Copy(h, r) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		log.Fatal(err) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	return fmt.Sprintf("sha256:%x", h.Sum(nil)), int(n) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | func startUpload(mp ModelPath, regOpts *RegistryOptions) (string, error) { | 
					
						
							|  |  |  | 	url := fmt.Sprintf("%s/v2/%s/blobs/uploads/", mp.Registry, mp.GetNamespaceRepository()) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | 	resp, err := makeRequest("POST", url, nil, nil, regOpts) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		log.Printf("couldn't start upload: %v", err) | 
					
						
							|  |  |  | 		return "", err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	defer resp.Body.Close() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	// Check for success
 | 
					
						
							|  |  |  | 	if resp.StatusCode != http.StatusAccepted { | 
					
						
							|  |  |  | 		body, _ := io.ReadAll(resp.Body) | 
					
						
							| 
									
										
										
										
											2023-07-18 03:33:39 +08:00
										 |  |  | 		return "", fmt.Errorf("registry responded with code %d: %s", resp.StatusCode, body) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	// Extract UUID location from header
 | 
					
						
							|  |  |  | 	location := resp.Header.Get("Location") | 
					
						
							|  |  |  | 	if location == "" { | 
					
						
							|  |  |  | 		return "", fmt.Errorf("location header is missing in response") | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	return location, nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | // Function to check if a blob already exists in the Docker registry
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | func checkBlobExistence(mp ModelPath, digest string, regOpts *RegistryOptions) (bool, error) { | 
					
						
							|  |  |  | 	url := fmt.Sprintf("%s/v2/%s/blobs/%s", mp.Registry, mp.GetNamespaceRepository(), digest) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | 	resp, err := makeRequest("HEAD", url, nil, nil, regOpts) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		log.Printf("couldn't check for blob: %v", err) | 
					
						
							|  |  |  | 		return false, err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	defer resp.Body.Close() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	// Check for success: If the blob exists, the Docker registry will respond with a 200 OK
 | 
					
						
							|  |  |  | 	return resp.StatusCode == http.StatusOK, nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-23 08:31:26 +08:00
										 |  |  | func uploadBlobChunked(mp ModelPath, location string, layer *Layer, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error { | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	// TODO allow resumability
 | 
					
						
							|  |  |  | 	// TODO allow canceling uploads via DELETE
 | 
					
						
							|  |  |  | 	// TODO allow cross repo blob mount
 | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-23 08:31:26 +08:00
										 |  |  | 	// Create URL
 | 
					
						
							|  |  |  | 	url := fmt.Sprintf("%s", location) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 	fp, err := GetBlobsPath(layer.Digest) | 
					
						
							| 
									
										
										
										
											2023-07-18 02:03:55 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	f, err := os.Open(fp) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-23 08:31:26 +08:00
										 |  |  | 	headers := make(map[string]string) | 
					
						
							|  |  |  | 	headers["Content-Type"] = "application/octet-stream" | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-23 08:31:26 +08:00
										 |  |  | 	chunkSize := 1 << 20 | 
					
						
							|  |  |  | 	buf := make([]byte, chunkSize) | 
					
						
							|  |  |  | 	var totalUploaded int | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-23 08:31:26 +08:00
										 |  |  | 	for { | 
					
						
							|  |  |  | 		n, err := f.Read(buf) | 
					
						
							|  |  |  | 		if err != nil { | 
					
						
							|  |  |  | 			return err | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 		headers["Content-Length"] = fmt.Sprintf("%d", n) | 
					
						
							|  |  |  | 		headers["Content-Range"] = fmt.Sprintf("%d-%d", totalUploaded, totalUploaded+n-1) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 		fn(api.ProgressResponse{ | 
					
						
							|  |  |  | 			Status:    fmt.Sprintf("uploading %s", layer.Digest), | 
					
						
							|  |  |  | 			Digest:    layer.Digest, | 
					
						
							|  |  |  | 			Total:     int(layer.Size), | 
					
						
							|  |  |  | 			Completed: int(totalUploaded), | 
					
						
							|  |  |  | 		}) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 		// change the buffersize for the last chunk
 | 
					
						
							|  |  |  | 		if n < chunkSize { | 
					
						
							|  |  |  | 			buf = buf[:n] | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 		resp, err := makeRequest("PATCH", url, headers, bytes.NewReader(buf), regOpts) | 
					
						
							|  |  |  | 		if err != nil { | 
					
						
							|  |  |  | 			log.Printf("couldn't upload blob: %v", err) | 
					
						
							|  |  |  | 			return err | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 		defer resp.Body.Close() | 
					
						
							|  |  |  | 		url = resp.Header.Get("Location") | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 		// Check for success: For a successful upload, the Docker registry will respond with a 201 Created
 | 
					
						
							|  |  |  | 		if resp.StatusCode != http.StatusAccepted { | 
					
						
							|  |  |  | 			fn(api.ProgressResponse{ | 
					
						
							|  |  |  | 				Status:    fmt.Sprintf("error uploading layer"), | 
					
						
							|  |  |  | 				Digest:    layer.Digest, | 
					
						
							|  |  |  | 				Total:     int(layer.Size), | 
					
						
							|  |  |  | 				Completed: int(totalUploaded), | 
					
						
							|  |  |  | 			}) | 
					
						
							|  |  |  | 			body, _ := io.ReadAll(resp.Body) | 
					
						
							|  |  |  | 			return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body)) | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 		totalUploaded += n | 
					
						
							|  |  |  | 		if totalUploaded >= layer.Size { | 
					
						
							|  |  |  | 			url = fmt.Sprintf("%s&digest=%s", url, layer.Digest) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 			// finish the upload
 | 
					
						
							|  |  |  | 			resp, err := makeRequest("PUT", url, nil, nil, regOpts) | 
					
						
							|  |  |  | 			if err != nil { | 
					
						
							|  |  |  | 				log.Printf("couldn't finish upload: %v", err) | 
					
						
							|  |  |  | 				return err | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 			defer resp.Body.Close() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 			if resp.StatusCode != http.StatusCreated { | 
					
						
							|  |  |  | 				body, _ := io.ReadAll(resp.Body) | 
					
						
							|  |  |  | 				return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body)) | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 			break | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 	} | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	return nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | func downloadBlob(mp ModelPath, digest string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error { | 
					
						
							| 
									
										
										
										
											2023-07-18 13:44:21 +08:00
										 |  |  | 	fp, err := GetBlobsPath(digest) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 	if fi, _ := os.Stat(fp); fi != nil { | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		// we already have the file, so return
 | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 		fn(api.ProgressResponse{ | 
					
						
							|  |  |  | 			Digest:    digest, | 
					
						
							|  |  |  | 			Total:     int(fi.Size()), | 
					
						
							|  |  |  | 			Completed: int(fi.Size()), | 
					
						
							|  |  |  | 		}) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		return nil | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	var size int64 | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	fi, err := os.Stat(fp + "-partial") | 
					
						
							|  |  |  | 	switch { | 
					
						
							|  |  |  | 	case errors.Is(err, os.ErrNotExist): | 
					
						
							|  |  |  | 		// noop, file doesn't exist so create it
 | 
					
						
							|  |  |  | 	case err != nil: | 
					
						
							|  |  |  | 		return fmt.Errorf("stat: %w", err) | 
					
						
							|  |  |  | 	default: | 
					
						
							|  |  |  | 		size = fi.Size() | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | 	url := fmt.Sprintf("%s/v2/%s/blobs/%s", mp.Registry, mp.GetNamespaceRepository(), digest) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	headers := map[string]string{ | 
					
						
							|  |  |  | 		"Range": fmt.Sprintf("bytes=%d-", size), | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | 	resp, err := makeRequest("GET", url, headers, nil, regOpts) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		log.Printf("couldn't download blob: %v", err) | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	defer resp.Body.Close() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent { | 
					
						
							|  |  |  | 		body, _ := ioutil.ReadAll(resp.Body) | 
					
						
							|  |  |  | 		return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body)) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	err = os.MkdirAll(path.Dir(fp), 0o700) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return fmt.Errorf("make blobs directory: %w", err) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	out, err := os.OpenFile(fp+"-partial", os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		panic(err) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	defer out.Close() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	remaining, _ := strconv.ParseInt(resp.Header.Get("Content-Length"), 10, 64) | 
					
						
							|  |  |  | 	completed := size | 
					
						
							|  |  |  | 	total := remaining + completed | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	for { | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 		fn(api.ProgressResponse{ | 
					
						
							|  |  |  | 			Status:    fmt.Sprintf("downloading %s", digest), | 
					
						
							|  |  |  | 			Digest:    digest, | 
					
						
							|  |  |  | 			Total:     int(total), | 
					
						
							|  |  |  | 			Completed: int(completed), | 
					
						
							|  |  |  | 		}) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 		if completed >= total { | 
					
						
							| 
									
										
										
										
											2023-07-21 02:23:43 +08:00
										 |  |  | 			if err := out.Close(); err != nil { | 
					
						
							|  |  |  | 				return err | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-18 03:39:55 +08:00
										 |  |  | 			if err := os.Rename(fp+"-partial", fp); err != nil { | 
					
						
							| 
									
										
										
										
											2023-07-19 09:51:30 +08:00
										 |  |  | 				fn(api.ProgressResponse{ | 
					
						
							|  |  |  | 					Status:    fmt.Sprintf("error renaming file: %v", err), | 
					
						
							|  |  |  | 					Digest:    digest, | 
					
						
							|  |  |  | 					Total:     int(total), | 
					
						
							|  |  |  | 					Completed: int(completed), | 
					
						
							|  |  |  | 				}) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 				return err | 
					
						
							|  |  |  | 			} | 
					
						
							| 
									
										
										
										
											2023-07-18 03:39:55 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 			break | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 		n, err := io.CopyN(out, resp.Body, 8192) | 
					
						
							|  |  |  | 		if err != nil && !errors.Is(err, io.EOF) { | 
					
						
							|  |  |  | 			return err | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 		completed += n | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	log.Printf("success getting %s\n", digest) | 
					
						
							|  |  |  | 	return nil | 
					
						
							|  |  |  | } | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | func makeRequest(method, url string, headers map[string]string, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) { | 
					
						
							|  |  |  | 	if !strings.HasPrefix(url, "http") { | 
					
						
							|  |  |  | 		if regOpts.Insecure { | 
					
						
							|  |  |  | 			url = "http://" + url | 
					
						
							|  |  |  | 		} else { | 
					
						
							|  |  |  | 			url = "https://" + url | 
					
						
							|  |  |  | 		} | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	req, err := http.NewRequest(method, url, body) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return nil, err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	for k, v := range headers { | 
					
						
							|  |  |  | 		req.Header.Set(k, v) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	// TODO: better auth
 | 
					
						
							| 
									
										
										
										
											2023-07-22 06:42:19 +08:00
										 |  |  | 	if regOpts.Username != "" && regOpts.Password != "" { | 
					
						
							|  |  |  | 		req.SetBasicAuth(regOpts.Username, regOpts.Password) | 
					
						
							| 
									
										
										
										
											2023-07-17 08:02:22 +08:00
										 |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	client := &http.Client{ | 
					
						
							|  |  |  | 		CheckRedirect: func(req *http.Request, via []*http.Request) error { | 
					
						
							|  |  |  | 			if len(via) >= 10 { | 
					
						
							|  |  |  | 				return fmt.Errorf("too many redirects") | 
					
						
							|  |  |  | 			} | 
					
						
							|  |  |  | 			log.Printf("redirected to: %s\n", req.URL) | 
					
						
							|  |  |  | 			return nil | 
					
						
							|  |  |  | 		}, | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	resp, err := client.Do(req) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return nil, err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	return resp, nil | 
					
						
							|  |  |  | } | 
					
						
							| 
									
										
										
										
											2023-07-21 02:44:05 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | func verifyBlob(digest string) error { | 
					
						
							|  |  |  | 	fp, err := GetBlobsPath(digest) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	f, err := os.Open(fp) | 
					
						
							|  |  |  | 	if err != nil { | 
					
						
							|  |  |  | 		return err | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 	defer f.Close() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	fileDigest, _ := GetSHA256Digest(f) | 
					
						
							|  |  |  | 	if digest != fileDigest { | 
					
						
							|  |  |  | 		return fmt.Errorf("digest mismatch: want %s, got %s", digest, fileDigest) | 
					
						
							|  |  |  | 	} | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 	return nil | 
					
						
							|  |  |  | } |