images.go 29 KB
Newer Older
1
2
3
package server

import (
4
	"archive/zip"
5
	"bytes"
6
	"context"
7
	"crypto/sha256"
Patrick Devine's avatar
Patrick Devine committed
8
	"encoding/hex"
9
10
11
12
	"encoding/json"
	"errors"
	"fmt"
	"io"
13
	"io/fs"
14
	"log"
15
	"log/slog"
16
	"net/http"
Michael Yang's avatar
Michael Yang committed
17
	"net/url"
18
19
	"os"
	"path/filepath"
Michael Yang's avatar
Michael Yang committed
20
	"runtime"
Michael Yang's avatar
Michael Yang committed
21
	"strconv"
22
	"strings"
Quinn Slack's avatar
Quinn Slack committed
23
	"text/template"
24

Michael Yang's avatar
Michael Yang committed
25
26
	"golang.org/x/exp/slices"

27
28
	"github.com/ollama/ollama/api"
	"github.com/ollama/ollama/convert"
Michael Yang's avatar
Michael Yang committed
29
	"github.com/ollama/ollama/format"
30
31
	"github.com/ollama/ollama/llm"
	"github.com/ollama/ollama/parser"
Michael Yang's avatar
Michael Yang committed
32
	"github.com/ollama/ollama/types/model"
33
	"github.com/ollama/ollama/version"
34
35
)

Michael Yang's avatar
Michael Yang committed
36
37
38
39
40
41
42
type registryOptions struct {
	Insecure bool
	Username string
	Password string
	Token    string
}

43
type Model struct {
Michael Yang's avatar
Michael Yang committed
44
	Name           string `json:"name"`
45
	Config         ConfigV2
Michael Yang's avatar
Michael Yang committed
46
47
	ShortName      string
	ModelPath      string
48
	ParentModel    string
Michael Yang's avatar
Michael Yang committed
49
50
51
52
53
54
	AdapterPaths   []string
	ProjectorPaths []string
	Template       string
	System         string
	License        []string
	Digest         string
Patrick Devine's avatar
Patrick Devine committed
55
	Size           int64
Michael Yang's avatar
Michael Yang committed
56
	Options        map[string]interface{}
57
58
59
	Messages       []Message
}

60
61
62
63
func (m *Model) IsEmbedding() bool {
	return slices.Contains(m.Config.ModelFamilies, "bert") || slices.Contains(m.Config.ModelFamilies, "nomic-bert")
}

64
65
66
type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
67
68
69
70
71
}

type ManifestV2 struct {
	SchemaVersion int      `json:"schemaVersion"`
	MediaType     string   `json:"mediaType"`
Michael Yang's avatar
Michael Yang committed
72
	Config        *Layer   `json:"config"`
73
74
75
76
	Layers        []*Layer `json:"layers"`
}

type ConfigV2 struct {
77
78
79
80
81
82
	ModelFormat   string   `json:"model_format"`
	ModelFamily   string   `json:"model_family"`
	ModelFamilies []string `json:"model_families"`
	ModelType     string   `json:"model_type"`
	FileType      string   `json:"file_type"`

83
	// required by spec
84
85
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
86
	RootFS       RootFS `json:"rootfs"`
87
88
}

Michael Yang's avatar
Michael Yang committed
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
func (c *ConfigV2) SetModelFormat(format string) {
	if c.ModelFormat == "" {
		c.ModelFormat = format
	}
}

func (c *ConfigV2) SetModelFamily(families ...string) {
	for _, family := range families {
		if c.ModelFamily == "" {
			c.ModelFamily = family
		}

		if !slices.Contains(c.ModelFamilies, family) {
			c.ModelFamilies = append(c.ModelFamilies, family)
		}
	}
}

func (c *ConfigV2) SetModelType(modelType string) {
	if c.ModelType == "" {
		c.ModelType = modelType
	}
}

func (c *ConfigV2) SetFileType(fileType string) {
	if c.FileType == "" {
		c.FileType = fileType
	}
}

119
120
121
122
123
type RootFS struct {
	Type    string   `json:"type"`
	DiffIDs []string `json:"diff_ids"`
}

Michael Yang's avatar
Michael Yang committed
124
func (m *ManifestV2) GetTotalSize() (total int64) {
Patrick Devine's avatar
Patrick Devine committed
125
126
127
	for _, layer := range m.Layers {
		total += layer.Size
	}
Michael Yang's avatar
Michael Yang committed
128

Patrick Devine's avatar
Patrick Devine committed
129
130
131
132
	total += m.Config.Size
	return total
}

Patrick Devine's avatar
Patrick Devine committed
133
func GetManifest(mp ModelPath) (*ManifestV2, string, error) {
134
	fp, err := mp.GetManifestPath()
135
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
136
		return nil, "", err
137
	}
138

139
	if _, err = os.Stat(fp); err != nil {
Patrick Devine's avatar
Patrick Devine committed
140
		return nil, "", err
141
142
143
144
	}

	var manifest *ManifestV2

145
	bts, err := os.ReadFile(fp)
146
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
147
		return nil, "", fmt.Errorf("couldn't open file '%s'", fp)
148
149
	}

Patrick Devine's avatar
Patrick Devine committed
150
151
152
	shaSum := sha256.Sum256(bts)
	shaStr := hex.EncodeToString(shaSum[:])

153
	if err := json.Unmarshal(bts, &manifest); err != nil {
Patrick Devine's avatar
Patrick Devine committed
154
		return nil, "", err
155
156
	}

Patrick Devine's avatar
Patrick Devine committed
157
	return manifest, shaStr, nil
158
159
160
}

func GetModel(name string) (*Model, error) {
161
	mp := ParseModelPath(name)
Patrick Devine's avatar
Patrick Devine committed
162
	manifest, digest, err := GetManifest(mp)
163
164
165
166
167
	if err != nil {
		return nil, err
	}

	model := &Model{
168
169
170
171
172
		Name:      mp.GetFullTagname(),
		ShortName: mp.GetShortTagname(),
		Digest:    digest,
		Template:  "{{ .Prompt }}",
		License:   []string{},
Patrick Devine's avatar
Patrick Devine committed
173
		Size:      manifest.GetTotalSize(),
174
175
	}

176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
	filename, err := GetBlobsPath(manifest.Config.Digest)
	if err != nil {
		return nil, err
	}

	configFile, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer configFile.Close()

	if err := json.NewDecoder(configFile).Decode(&model.Config); err != nil {
		return nil, err
	}

191
	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
192
		filename, err := GetBlobsPath(layer.Digest)
193
194
195
196
		if err != nil {
			return nil, err
		}

197
198
199
		switch layer.MediaType {
		case "application/vnd.ollama.image.model":
			model.ModelPath = filename
200
			model.ParentModel = layer.From
201
		case "application/vnd.ollama.image.embed":
202
203
			// Deprecated in versions  > 0.1.2
			// TODO: remove this warning in a future version
204
			slog.Info("WARNING: model contains embeddings, but embeddings in modelfiles have been deprecated and will be ignored.")
205
206
		case "application/vnd.ollama.image.adapter":
			model.AdapterPaths = append(model.AdapterPaths, filename)
Michael Yang's avatar
Michael Yang committed
207
208
		case "application/vnd.ollama.image.projector":
			model.ProjectorPaths = append(model.ProjectorPaths, filename)
209
210
211
212
213
214
215
216
217
		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)
218
219
220
			if err != nil {
				return nil, err
			}
221
222

			model.System = string(bts)
223
224
225
226
227
228
229
		case "application/vnd.ollama.image.prompt":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

			model.Template = string(bts)
230
		case "application/vnd.ollama.image.params":
Michael Yang's avatar
Michael Yang committed
231
232
233
234
235
			params, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer params.Close()
236

237
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
238
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
239
240
				return nil, err
			}
241
242
243
244
245
246
247
248
249
250
		case "application/vnd.ollama.image.messages":
			msgs, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer msgs.Close()

			if err = json.NewDecoder(msgs).Decode(&model.Messages); err != nil {
				return nil, err
			}
Patrick Devine's avatar
Patrick Devine committed
251
252
253
254
255
256
		case "application/vnd.ollama.image.license":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}
			model.License = append(model.License, string(bts))
257
258
259
260
261
262
		}
	}

	return model, nil
}

263
264
func realpath(mfDir, from string) string {
	abspath, err := filepath.Abs(from)
Michael Yang's avatar
Michael Yang committed
265
	if err != nil {
266
		return from
267
268
	}

Michael Yang's avatar
Michael Yang committed
269
	home, err := os.UserHomeDir()
270
	if err != nil {
Michael Yang's avatar
Michael Yang committed
271
		return abspath
272
273
	}

274
	if from == "~" {
Michael Yang's avatar
Michael Yang committed
275
		return home
276
277
278
279
280
281
282
	} else if strings.HasPrefix(from, "~/") {
		return filepath.Join(home, from[2:])
	}

	if _, err := os.Stat(filepath.Join(mfDir, from)); err == nil {
		// this is a file relative to the Modelfile
		return filepath.Join(mfDir, from)
283
284
	}

Michael Yang's avatar
Michael Yang committed
285
286
287
	return abspath
}

Michael Yang's avatar
Michael Yang committed
288
func CreateModel(ctx context.Context, name, modelFileDir, quantization string, commands []parser.Command, fn func(resp api.ProgressResponse)) error {
289
290
291
292
293
294
295
	deleteMap := make(map[string]struct{})
	if manifest, _, err := GetManifest(ParseModelPath(name)); err == nil {
		for _, layer := range append(manifest.Layers, manifest.Config) {
			deleteMap[layer.Digest] = struct{}{}
		}
	}

296
297
	config := ConfigV2{
		OS:           "linux",
Michael Yang's avatar
Michael Yang committed
298
		Architecture: "amd64",
Michael Yang's avatar
Michael Yang committed
299
300
301
		RootFS: RootFS{
			Type: "layers",
		},
302
303
	}

Michael Yang's avatar
Michael Yang committed
304
	var layers Layers
305
	messages := []string{}
Michael Yang's avatar
Michael Yang committed
306

307
	params := make(map[string][]string)
Michael Yang's avatar
Michael Yang committed
308
309
	fromParams := make(map[string]any)

310
	for _, c := range commands {
Michael Yang's avatar
Michael Yang committed
311
312
		mediatype := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)

313
314
		switch c.Name {
		case "model":
Michael Yang's avatar
Michael Yang committed
315
316
317
318
319
320
321
322
323
			if strings.HasPrefix(c.Args, "@") {
				blobPath, err := GetBlobsPath(strings.TrimPrefix(c.Args, "@"))
				if err != nil {
					return err
				}

				c.Args = blobPath
			}

324
325
			pathName := realpath(modelFileDir, c.Args)

326
			ggufName, err := convertModel(name, pathName, fn)
327
			if err != nil {
328
				var pathErr *fs.PathError
329
330
331
				switch {
				case errors.Is(err, zip.ErrFormat):
					// it's not a safetensor archive
332
333
				case errors.As(err, &pathErr):
					// it's not a file on disk, could be a model reference
334
335
336
337
338
339
340
341
				default:
					return err
				}
			}

			if ggufName != "" {
				pathName = ggufName
				defer os.RemoveAll(ggufName)
Michael Yang's avatar
Michael Yang committed
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361

				if quantization != "" {
					quantization = strings.ToUpper(quantization)
					fn(api.ProgressResponse{Status: fmt.Sprintf("quantizing %s model to %s", "F16", quantization)})
					tempfile, err := os.CreateTemp(filepath.Dir(ggufName), quantization)
					if err != nil {
						return err
					}
					defer os.RemoveAll(tempfile.Name())

					if err := llm.Quantize(ggufName, tempfile.Name(), quantization); err != nil {
						return err
					}

					if err := tempfile.Close(); err != nil {
						return err
					}

					pathName = tempfile.Name()
				}
362
363
364
			}

			bin, err := os.Open(pathName)
365
			if err != nil {
Michael Yang's avatar
Michael Yang committed
366
367
368
369
370
371
				// not a file on disk so must be a model reference
				modelpath := ParseModelPath(c.Args)
				manifest, _, err := GetManifest(modelpath)
				switch {
				case errors.Is(err, os.ErrNotExist):
					fn(api.ProgressResponse{Status: "pulling model"})
Michael Yang's avatar
Michael Yang committed
372
					if err := PullModel(ctx, c.Args, &registryOptions{}, fn); err != nil {
373
374
375
						return err
					}

Michael Yang's avatar
Michael Yang committed
376
					manifest, _, err = GetManifest(modelpath)
377
378
379
					if err != nil {
						return err
					}
Michael Yang's avatar
Michael Yang committed
380
381
				case err != nil:
					return err
382
				}
383

384
				fn(api.ProgressResponse{Status: "reading model metadata"})
Michael Yang's avatar
Michael Yang committed
385
				fromConfigPath, err := GetBlobsPath(manifest.Config.Digest)
Michael Yang's avatar
Michael Yang committed
386
387
388
389
				if err != nil {
					return err
				}

Michael Yang's avatar
Michael Yang committed
390
				fromConfigFile, err := os.Open(fromConfigPath)
Michael Yang's avatar
Michael Yang committed
391
392
393
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
394
				defer fromConfigFile.Close()
Michael Yang's avatar
Michael Yang committed
395

Michael Yang's avatar
Michael Yang committed
396
397
				var fromConfig ConfigV2
				if err := json.NewDecoder(fromConfigFile).Decode(&fromConfig); err != nil {
Michael Yang's avatar
Michael Yang committed
398
399
400
					return err
				}

Bruce MacDonald's avatar
Bruce MacDonald committed
401
402
403
404
405
				// if the model is still not in gguf format, error out
				if fromConfig.ModelFormat != "gguf" {
					return fmt.Errorf("%s is not in gguf format, this base model is not compatible with this version of ollama", c.Args)
				}

Michael Yang's avatar
Michael Yang committed
406
407
408
409
				config.SetModelFormat(fromConfig.ModelFormat)
				config.SetModelFamily(append(fromConfig.ModelFamilies, fromConfig.ModelFamily)...)
				config.SetModelType(fromConfig.ModelType)
				config.SetFileType(fromConfig.FileType)
Michael Yang's avatar
Michael Yang committed
410

Michael Yang's avatar
Michael Yang committed
411
412
413
414
				for _, layer := range manifest.Layers {
					deleteMap[layer.Digest] = struct{}{}
					if layer.MediaType == "application/vnd.ollama.image.params" {
						fromParamsPath, err := GetBlobsPath(layer.Digest)
Michael Yang's avatar
Michael Yang committed
415
416
417
418
						if err != nil {
							return err
						}

Michael Yang's avatar
Michael Yang committed
419
						fromParamsFile, err := os.Open(fromParamsPath)
Michael Yang's avatar
Michael Yang committed
420
421
422
						if err != nil {
							return err
						}
Michael Yang's avatar
Michael Yang committed
423
						defer fromParamsFile.Close()
Michael Yang's avatar
Michael Yang committed
424

Michael Yang's avatar
Michael Yang committed
425
						if err := json.NewDecoder(fromParamsFile).Decode(&fromParams); err != nil {
Michael Yang's avatar
Michael Yang committed
426
427
428
429
							return err
						}
					}

Michael Yang's avatar
Michael Yang committed
430
					layer, err := NewLayerFromLayer(layer.Digest, layer.MediaType, modelpath.GetShortTagname())
431
432
433
					if err != nil {
						return err
					}
Michael Yang's avatar
Michael Yang committed
434

Michael Yang's avatar
Michael Yang committed
435
					layers.Add(layer)
436
				}
Michael Yang's avatar
Michael Yang committed
437
438
439

				deleteMap[manifest.Config.Digest] = struct{}{}
				continue
440
			}
Michael Yang's avatar
Michael Yang committed
441
			defer bin.Close()
442

443
444
445
			var offset int64
			for {
				fn(api.ProgressResponse{Status: "creating model layer"})
446
447
448
				if _, err := bin.Seek(offset, io.SeekStart); err != nil {
					return err
				}
449

Michael Yang's avatar
Michael Yang committed
450
451
452
453
454
455
456
				ggml, size, err := llm.DecodeGGML(bin)
				if errors.Is(err, io.EOF) {
					break
				} else if errors.Is(err, llm.ErrUnsupportedFormat) {
					return fmt.Errorf("model binary specified in FROM field is not a valid gguf format model, %w", err)
				} else if err != nil {
					return err
457
				}
Michael Yang's avatar
Michael Yang committed
458

Michael Yang's avatar
Michael Yang committed
459
				config.SetModelFormat(ggml.Name())
Michael Yang's avatar
Michael Yang committed
460
461
462
				config.SetModelFamily(ggml.KV().Architecture())
				config.SetModelType(format.HumanNumber(ggml.KV().ParameterCount()))
				config.SetFileType(ggml.KV().FileType())
463

464
				mediatype := mediatype
Michael Yang's avatar
Michael Yang committed
465
				if ggml.KV().Architecture() == "clip" {
466
467
					mediatype = "application/vnd.ollama.image.projector"
				}
468

Michael Yang's avatar
Michael Yang committed
469
				sr := io.NewSectionReader(bin, offset, size)
470
471
472
473
474
475
476
				layer, err := NewLayer(sr, mediatype)
				if err != nil {
					return err
				}

				layers.Add(layer)

Michael Yang's avatar
Michael Yang committed
477
				offset += size
478
			}
Michael Yang's avatar
Michael Yang committed
479
		case "adapter":
480
481
482
483
484
485
486
487
			if strings.HasPrefix(c.Args, "@") {
				blobPath, err := GetBlobsPath(strings.TrimPrefix(c.Args, "@"))
				if err != nil {
					return err
				}

				c.Args = blobPath
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
488

Michael Yang's avatar
Michael Yang committed
489
			fn(api.ProgressResponse{Status: "creating adapter layer"})
490
			bin, err := os.Open(realpath(modelFileDir, c.Args))
491
			if err != nil {
Michael Yang's avatar
Michael Yang committed
492
				return err
493
			}
Michael Yang's avatar
Michael Yang committed
494
			defer bin.Close()
495

Michael Yang's avatar
Michael Yang committed
496
			_, size, err := llm.DecodeGGML(bin)
Michael Yang's avatar
Michael Yang committed
497
498
499
500
			if err != nil {
				return err
			}

Michael Yang's avatar
Michael Yang committed
501
			sr := io.NewSectionReader(bin, 0, size)
Michael Yang's avatar
Michael Yang committed
502
			layer, err := NewLayer(sr, mediatype)
503
			if err != nil {
Michael Yang's avatar
Michael Yang committed
504
				return err
505
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
506

Michael Yang's avatar
Michael Yang committed
507
			layers.Add(layer)
Michael Yang's avatar
Michael Yang committed
508
509
		case "license":
			fn(api.ProgressResponse{Status: "creating license layer"})
Michael Yang's avatar
Michael Yang committed
510
511
512

			bin := strings.NewReader(c.Args)
			layer, err := NewLayer(bin, mediatype)
Bruce MacDonald's avatar
Bruce MacDonald committed
513
514
515
516
			if err != nil {
				return err
			}

Michael Yang's avatar
Michael Yang committed
517
			layers.Add(layer)
Michael Yang's avatar
Michael Yang committed
518
519
520
		case "template", "system":
			fn(api.ProgressResponse{Status: fmt.Sprintf("creating %s layer", c.Name)})

Michael Yang's avatar
Michael Yang committed
521
522
			bin := strings.NewReader(c.Args)
			layer, err := NewLayer(bin, mediatype)
523
			if err != nil {
524
				return err
525
			}
526

Michael Yang's avatar
Michael Yang committed
527
			layers.Replace(layer)
528
529
		case "message":
			messages = append(messages, c.Args)
530
		default:
531
			params[c.Name] = append(params[c.Name], c.Args)
532
533
534
		}
	}

535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
	if len(messages) > 0 {
		fn(api.ProgressResponse{Status: "creating parameters layer"})

		msgs := make([]api.Message, 0)

		for _, m := range messages {
			// todo: handle images
			msg := strings.SplitN(m, ": ", 2)
			msgs = append(msgs, api.Message{Role: msg[0], Content: msg[1]})
		}

		var b bytes.Buffer
		if err := json.NewEncoder(&b).Encode(msgs); err != nil {
			return err
		}

		layer, err := NewLayer(&b, "application/vnd.ollama.image.messages")
		if err != nil {
			return err
		}

		layers.Replace(layer)
	}

Michael Yang's avatar
Michael Yang committed
559
	if len(params) > 0 {
Michael Yang's avatar
Michael Yang committed
560
		fn(api.ProgressResponse{Status: "creating parameters layer"})
Michael Yang's avatar
Michael Yang committed
561

562
		formattedParams, err := api.FormatParams(params)
563
		if err != nil {
Michael Yang's avatar
Michael Yang committed
564
			return err
565
		}
566

Michael Yang's avatar
Michael Yang committed
567
		for k, v := range fromParams {
Michael Yang's avatar
Michael Yang committed
568
569
570
571
572
			if _, ok := formattedParams[k]; !ok {
				formattedParams[k] = v
			}
		}

Michael Yang's avatar
Michael Yang committed
573
574
		var b bytes.Buffer
		if err := json.NewEncoder(&b).Encode(formattedParams); err != nil {
575
576
577
			return err
		}

Michael Yang's avatar
Michael Yang committed
578
		fn(api.ProgressResponse{Status: "creating config layer"})
Michael Yang's avatar
Michael Yang committed
579
		layer, err := NewLayer(&b, "application/vnd.ollama.image.params")
580
		if err != nil {
Michael Yang's avatar
Michael Yang committed
581
			return err
582
		}
Michael Yang's avatar
Michael Yang committed
583

Michael Yang's avatar
Michael Yang committed
584
		layers.Replace(layer)
585
586
	}

Michael Yang's avatar
Michael Yang committed
587
588
589
	digests := make([]string, len(layers.items))
	for i, layer := range layers.items {
		digests[i] = layer.Digest
590
591
	}

Michael Yang's avatar
Michael Yang committed
592
	config.RootFS.DiffIDs = digests
Michael Yang's avatar
Michael Yang committed
593

Michael Yang's avatar
Michael Yang committed
594
595
	var b bytes.Buffer
	if err := json.NewEncoder(&b).Encode(config); err != nil {
596
597
598
		return err
	}

Michael Yang's avatar
Michael Yang committed
599
600
	configLayer, err := NewLayer(&b, "application/vnd.docker.container.image.v1+json")
	if err != nil {
601
602
603
		return err
	}

Michael Yang's avatar
Michael Yang committed
604
	delete(deleteMap, configLayer.Digest)
605

Michael Yang's avatar
Michael Yang committed
606
607
	for _, layer := range append(layers.items, configLayer) {
		committed, err := layer.Commit()
608
609
610
		if err != nil {
			return err
		}
611

Michael Yang's avatar
Michael Yang committed
612
613
614
		status := "writing layer"
		if !committed {
			status = "using already created layer"
615
616
		}

Michael Yang's avatar
Michael Yang committed
617
		fn(api.ProgressResponse{Status: fmt.Sprintf("%s %s", status, layer.Digest)})
618

Michael Yang's avatar
Michael Yang committed
619
		delete(deleteMap, layer.Digest)
620
621
	}

Michael Yang's avatar
Michael Yang committed
622
623
	fn(api.ProgressResponse{Status: "writing manifest"})
	if err := WriteManifest(name, configLayer, layers.items); err != nil {
624
625
		return err
	}
626

Michael Yang's avatar
Michael Yang committed
627
628
629
	if noprune := os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
		if err := deleteUnusedLayers(nil, deleteMap, false); err != nil {
			return err
630
631
632
		}
	}

Michael Yang's avatar
Michael Yang committed
633
634
	fn(api.ProgressResponse{Status: "success"})
	return nil
635
636
}

637
func convertModel(name, path string, fn func(resp api.ProgressResponse)) (string, error) {
638
	r, err := zip.OpenReader(path)
639
640
641
642
643
644
645
646
647
648
649
	if err != nil {
		return "", err
	}
	defer r.Close()

	tempDir, err := os.MkdirTemp("", "ollama-convert")
	if err != nil {
		return "", err
	}
	defer os.RemoveAll(tempDir)

650
	fn(api.ProgressResponse{Status: "unpacking model metadata"})
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
	for _, f := range r.File {
		fpath := filepath.Join(tempDir, f.Name)
		outFile, err := os.OpenFile(fpath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
		if err != nil {
			return "", err
		}

		rc, err := f.Open()
		if err != nil {
			return "", err
		}

		_, err = io.Copy(outFile, rc)
		if err != nil {
			return "", err
		}

		outFile.Close()
		rc.Close()
	}

672
	mf, err := convert.GetModelFormat(tempDir)
673
674
675
676
	if err != nil {
		return "", err
	}

677
	params, err := mf.GetParams(tempDir)
678
679
	if err != nil {
		return "", err
680
681
	}

682
683
684
685
686
687
	mArch, err := mf.GetModelArch(name, tempDir, params)
	if err != nil {
		return "", err
	}

	fn(api.ProgressResponse{Status: "processing tensors"})
688
	if err := mArch.GetTensors(); err != nil {
689
690
691
		return "", err
	}

692
	if err := mArch.LoadVocab(); err != nil {
693
694
695
		return "", err
	}

696
	fn(api.ProgressResponse{Status: "converting model"})
697
	path, err = mArch.WriteGGUF()
698
699
700
701
	if err != nil {
		return "", err
	}

702
	return path, nil
703
704
}

Michael Yang's avatar
Michael Yang committed
705
706
func CopyModel(src, dst model.Name) error {
	manifests, err := GetManifestPath()
707
708
709
710
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
711
712
	dstpath := filepath.Join(manifests, dst.FilepathNoBuild())
	if err := os.MkdirAll(filepath.Dir(dstpath), 0o755); err != nil {
713
714
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
715

Michael Yang's avatar
Michael Yang committed
716
717
	srcpath := filepath.Join(manifests, src.FilepathNoBuild())
	srcfile, err := os.Open(srcpath)
Patrick Devine's avatar
Patrick Devine committed
718
719
720
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
721
	defer srcfile.Close()
Patrick Devine's avatar
Patrick Devine committed
722

Michael Yang's avatar
Michael Yang committed
723
	dstfile, err := os.Create(dstpath)
Patrick Devine's avatar
Patrick Devine committed
724
725
726
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
727
	defer dstfile.Close()
Patrick Devine's avatar
Patrick Devine committed
728

Michael Yang's avatar
Michael Yang committed
729
730
	_, err = io.Copy(dstfile, srcfile)
	return err
Patrick Devine's avatar
Patrick Devine committed
731
732
}

Michael Yang's avatar
Michael Yang committed
733
func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{}, dryRun bool) error {
734
735
736
737
	fp, err := GetManifestPath()
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
738
739
740
741

	walkFunc := func(path string, info os.FileInfo, _ error) error {
		if info.IsDir() {
			return nil
742
743
		}

Michael Yang's avatar
Michael Yang committed
744
745
746
747
		dir, file := filepath.Split(path)
		dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
		tag := strings.Join([]string{dir, file}, ":")
		fmp := ParseModelPath(tag)
748

Michael Yang's avatar
Michael Yang committed
749
		// skip the manifest we're trying to delete
750
		if skipModelPath != nil && skipModelPath.GetFullTagname() == fmp.GetFullTagname() {
Michael Yang's avatar
Michael Yang committed
751
			return nil
752
		}
Michael Yang's avatar
Michael Yang committed
753
754
755
756

		// save (i.e. delete from the deleteMap) any files used in other manifests
		manifest, _, err := GetManifest(fmp)
		if err != nil {
Michael Yang's avatar
Michael Yang committed
757
			// nolint: nilerr
Michael Yang's avatar
Michael Yang committed
758
759
760
761
762
763
764
765
			return nil
		}

		for _, layer := range manifest.Layers {
			delete(deleteMap, layer.Digest)
		}

		delete(deleteMap, manifest.Config.Digest)
766
		return nil
Michael Yang's avatar
Michael Yang committed
767
768
769
	}

	if err := filepath.Walk(fp, walkFunc); err != nil {
Michael Yang's avatar
Michael Yang committed
770
771
		return err
	}
772
773

	// only delete the files which are still in the deleteMap
Michael Yang's avatar
Michael Yang committed
774
775
776
	for k := range deleteMap {
		fp, err := GetBlobsPath(k)
		if err != nil {
777
			slog.Info(fmt.Sprintf("couldn't get file path for '%s': %v", k, err))
Michael Yang's avatar
Michael Yang committed
778
779
780
781
			continue
		}
		if !dryRun {
			if err := os.Remove(fp); err != nil {
782
				slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err))
783
784
				continue
			}
Michael Yang's avatar
Michael Yang committed
785
		} else {
786
			slog.Info(fmt.Sprintf("wanted to remove: %s", fp))
787
788
789
		}
	}

790
791
792
793
	return nil
}

func PruneLayers() error {
Michael Yang's avatar
Michael Yang committed
794
	deleteMap := make(map[string]struct{})
795
796
797
798
799
800
801
	p, err := GetBlobsPath("")
	if err != nil {
		return err
	}

	blobs, err := os.ReadDir(p)
	if err != nil {
802
		slog.Info(fmt.Sprintf("couldn't read dir '%s': %v", p, err))
803
804
805
806
807
		return err
	}

	for _, blob := range blobs {
		name := blob.Name()
808
		name = strings.ReplaceAll(name, "-", ":")
Michael Yang's avatar
Michael Yang committed
809
810
811
		if strings.HasPrefix(name, "sha256:") {
			deleteMap[name] = struct{}{}
		}
812
813
	}

814
	slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap)))
815
816
817
818
819
820

	err = deleteUnusedLayers(nil, deleteMap, false)
	if err != nil {
		return err
	}

821
	slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
822
823
824
825

	return nil
}

Michael Yang's avatar
Michael Yang committed
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
func PruneDirectory(path string) error {
	info, err := os.Lstat(path)
	if err != nil {
		return err
	}

	if info.IsDir() && info.Mode()&os.ModeSymlink == 0 {
		entries, err := os.ReadDir(path)
		if err != nil {
			return err
		}

		for _, entry := range entries {
			if err := PruneDirectory(filepath.Join(path, entry.Name())); err != nil {
				return err
			}
		}

		entries, err = os.ReadDir(path)
		if err != nil {
			return err
		}

		if len(entries) > 0 {
			return nil
		}

		return os.Remove(path)
	}

	return nil
}

859
860
861
862
863
864
865
func DeleteModel(name string) error {
	mp := ParseModelPath(name)
	manifest, _, err := GetManifest(mp)
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
866
	deleteMap := make(map[string]struct{})
867
	for _, layer := range manifest.Layers {
Michael Yang's avatar
Michael Yang committed
868
		deleteMap[layer.Digest] = struct{}{}
869
	}
Michael Yang's avatar
Michael Yang committed
870
	deleteMap[manifest.Config.Digest] = struct{}{}
871
872
873
874
875
876

	err = deleteUnusedLayers(&mp, deleteMap, false)
	if err != nil {
		return err
	}

877
	fp, err := mp.GetManifestPath()
878
879
880
881
882
	if err != nil {
		return err
	}
	err = os.Remove(fp)
	if err != nil {
883
		slog.Info(fmt.Sprintf("couldn't remove manifest file '%s': %v", fp, err))
884
885
886
887
888
889
		return err
	}

	return nil
}

Patrick Devine's avatar
Patrick Devine committed
890
func ShowModelfile(model *Model) (string, error) {
Michael Yang's avatar
Michael Yang committed
891
	var mt struct {
Patrick Devine's avatar
Patrick Devine committed
892
		*Model
Michael Yang's avatar
Michael Yang committed
893
		From       string
Michael Yang's avatar
Michael Yang committed
894
		Parameters map[string][]any
Patrick Devine's avatar
Patrick Devine committed
895
896
	}

Michael Yang's avatar
Michael Yang committed
897
	mt.Parameters = make(map[string][]any)
Patrick Devine's avatar
Patrick Devine committed
898
	for k, v := range model.Options {
Michael Yang's avatar
Michael Yang committed
899
900
901
		if s, ok := v.([]any); ok {
			mt.Parameters[k] = s
			continue
Patrick Devine's avatar
Patrick Devine committed
902
903
		}

Michael Yang's avatar
Michael Yang committed
904
		mt.Parameters[k] = []any{v}
Patrick Devine's avatar
Patrick Devine committed
905
906
	}

Michael Yang's avatar
Michael Yang committed
907
908
	mt.Model = model
	mt.From = model.ModelPath
Patrick Devine's avatar
Patrick Devine committed
909

910
911
	if model.ParentModel != "" {
		mt.From = model.ParentModel
Patrick Devine's avatar
Patrick Devine committed
912
913
914
915
916
917
918
919
	}

	modelFile := `# Modelfile generated by "ollama show"
# To build a new Modelfile based on this one, replace the FROM line with:
# FROM {{ .ShortName }}

FROM {{ .From }}
TEMPLATE """{{ .Template }}"""
920
921

{{- if .System }}
Patrick Devine's avatar
Patrick Devine committed
922
SYSTEM """{{ .System }}"""
923
{{- end }}
924
925
926
927

{{- range $adapter := .AdapterPaths }}
ADAPTER {{ $adapter }}
{{- end }}
Michael Yang's avatar
Michael Yang committed
928

Michael Yang's avatar
Michael Yang committed
929
930
931
932
{{- range $k, $v := .Parameters }}
{{- range $parameter := $v }}
PARAMETER {{ $k }} {{ printf "%#v" $parameter }}
{{- end }}
Michael Yang's avatar
Michael Yang committed
933
{{- end }}`
Patrick Devine's avatar
Patrick Devine committed
934
935
936

	tmpl, err := template.New("").Parse(modelFile)
	if err != nil {
937
		slog.Info(fmt.Sprintf("error parsing template: %q", err))
Patrick Devine's avatar
Patrick Devine committed
938
939
940
941
942
943
		return "", err
	}

	var buf bytes.Buffer

	if err = tmpl.Execute(&buf, mt); err != nil {
944
		slog.Info(fmt.Sprintf("error executing template: %q", err))
Patrick Devine's avatar
Patrick Devine committed
945
946
947
948
949
950
		return "", err
	}

	return buf.String(), nil
}

Michael Yang's avatar
Michael Yang committed
951
func PushModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
952
	mp := ParseModelPath(name)
953
954
	fn(api.ProgressResponse{Status: "retrieving manifest"})

955
956
957
958
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
	}

Patrick Devine's avatar
Patrick Devine committed
959
	manifest, _, err := GetManifest(mp)
960
	if err != nil {
961
		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
962
963
964
965
		return err
	}

	var layers []*Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
966
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
967
	layers = append(layers, manifest.Config)
968
969

	for _, layer := range layers {
Michael Yang's avatar
Michael Yang committed
970
		if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
971
			slog.Info(fmt.Sprintf("error uploading blob: %v", err))
972
973
974
			if errors.Is(err, errUnauthorized) {
				return fmt.Errorf("unable to push %s, make sure this namespace exists and you are authorized to push to it", ParseModelPath(name).GetNamespaceRepository())
			}
975
976
			return err
		}
977
978
	}

979
	fn(api.ProgressResponse{Status: "pushing manifest"})
Michael Yang's avatar
Michael Yang committed
980
981
	requestURL := mp.BaseURL()
	requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
982
983
984
985
986
987

	manifestJSON, err := json.Marshal(manifest)
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
988
989
	headers := make(http.Header)
	headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
990
	resp, err := makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
991
992
993
994
995
	if err != nil {
		return err
	}
	defer resp.Body.Close()

996
	fn(api.ProgressResponse{Status: "success"})
997
998
999
1000

	return nil
}

Michael Yang's avatar
Michael Yang committed
1001
func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
1002
1003
	mp := ParseModelPath(name)

1004
1005
1006
1007
1008
	var manifest *ManifestV2
	var err error
	var noprune string

	// build deleteMap to prune unused layers
Michael Yang's avatar
Michael Yang committed
1009
	deleteMap := make(map[string]struct{})
1010
1011
1012
1013
1014
1015
1016
1017
1018

	if noprune = os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
		manifest, _, err = GetManifest(mp)
		if err != nil && !errors.Is(err, os.ErrNotExist) {
			return err
		}

		if manifest != nil {
			for _, l := range manifest.Layers {
Michael Yang's avatar
Michael Yang committed
1019
				deleteMap[l.Digest] = struct{}{}
1020
			}
Michael Yang's avatar
Michael Yang committed
1021
			deleteMap[manifest.Config.Digest] = struct{}{}
1022
1023
1024
		}
	}

1025
1026
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
1027
	}
1028

1029
	fn(api.ProgressResponse{Status: "pulling manifest"})
1030

1031
	manifest, err = pullModelManifest(ctx, mp, regOpts)
1032
	if err != nil {
1033
		return fmt.Errorf("pull model manifest: %s", err)
1034
1035
1036
	}

	var layers []*Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
1037
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
1038
	layers = append(layers, manifest.Config)
1039
1040

	for _, layer := range layers {
1041
1042
1043
1044
1045
1046
1047
1048
		if err := downloadBlob(
			ctx,
			downloadOpts{
				mp:      mp,
				digest:  layer.Digest,
				regOpts: regOpts,
				fn:      fn,
			}); err != nil {
1049
1050
			return err
		}
1051
		delete(deleteMap, layer.Digest)
1052
	}
1053
	delete(deleteMap, manifest.Config.Digest)
1054

Michael Yang's avatar
Michael Yang committed
1055
1056
1057
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
		if err := verifyBlob(layer.Digest); err != nil {
1058
1059
1060
1061
1062
1063
1064
1065
			if errors.Is(err, errDigestMismatch) {
				// something went wrong, delete the blob
				fp, err := GetBlobsPath(layer.Digest)
				if err != nil {
					return err
				}
				if err := os.Remove(fp); err != nil {
					// log this, but return the original error
1066
					slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err))
1067
1068
				}
			}
Michael Yang's avatar
Michael Yang committed
1069
1070
1071
1072
			return err
		}
	}

1073
	fn(api.ProgressResponse{Status: "writing manifest"})
1074

1075
	manifestJSON, err := json.Marshal(manifest)
1076
1077
1078
1079
	if err != nil {
		return err
	}

1080
	fp, err := mp.GetManifestPath()
1081
1082
1083
	if err != nil {
		return err
	}
1084
1085
1086
	if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
		return err
	}
1087

Bruce MacDonald's avatar
Bruce MacDonald committed
1088
	err = os.WriteFile(fp, manifestJSON, 0o644)
1089
	if err != nil {
1090
		slog.Info(fmt.Sprintf("couldn't write to %s", fp))
1091
1092
1093
		return err
	}

1094
1095
1096
1097
1098
1099
1100
1101
	if noprune == "" {
		fn(api.ProgressResponse{Status: "removing any unused layers"})
		err = deleteUnusedLayers(nil, deleteMap, false)
		if err != nil {
			return err
		}
	}

1102
	fn(api.ProgressResponse{Status: "success"})
1103
1104
1105
1106

	return nil
}

Michael Yang's avatar
Michael Yang committed
1107
func pullModelManifest(ctx context.Context, mp ModelPath, regOpts *registryOptions) (*ManifestV2, error) {
Michael Yang's avatar
Michael Yang committed
1108
	requestURL := mp.BaseURL().JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
1109

Michael Yang's avatar
Michael Yang committed
1110
1111
	headers := make(http.Header)
	headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
1112
	resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	var m *ManifestV2
	if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
		return nil, err
	}

	return m, err
}

// GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
Michael Yang's avatar
Michael Yang committed
1127
func GetSHA256Digest(r io.Reader) (string, int64) {
Michael Yang's avatar
Michael Yang committed
1128
1129
1130
1131
1132
1133
	h := sha256.New()
	n, err := io.Copy(h, r)
	if err != nil {
		log.Fatal(err)
	}

Michael Yang's avatar
Michael Yang committed
1134
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
1135
1136
}

1137
var errUnauthorized = errors.New("unauthorized")
1138

Michael Yang's avatar
Michael Yang committed
1139
func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *registryOptions) (*http.Response, error) {
Michael Yang's avatar
Michael Yang committed
1140
	for i := 0; i < 2; i++ {
Michael Yang's avatar
Michael Yang committed
1141
		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
Michael Yang's avatar
Michael Yang committed
1142
		if err != nil {
Michael Yang's avatar
Michael Yang committed
1143
			if !errors.Is(err, context.Canceled) {
1144
				slog.Info(fmt.Sprintf("request failed: %v", err))
Michael Yang's avatar
Michael Yang committed
1145
1146
			}

Michael Yang's avatar
Michael Yang committed
1147
1148
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
1149
1150
1151
1152

		switch {
		case resp.StatusCode == http.StatusUnauthorized:
			// Handle authentication error with one retry
Michael Yang's avatar
Michael Yang committed
1153
1154
			challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
			token, err := getAuthorizationToken(ctx, challenge)
Michael Yang's avatar
Michael Yang committed
1155
1156
1157
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
			regOpts.Token = token
			if body != nil {
				_, err = body.Seek(0, io.SeekStart)
				if err != nil {
					return nil, err
				}
			}
		case resp.StatusCode == http.StatusNotFound:
			return nil, os.ErrNotExist
		case resp.StatusCode >= http.StatusBadRequest:
			responseBody, err := io.ReadAll(resp.Body)
			if err != nil {
				return nil, fmt.Errorf("%d: %s", resp.StatusCode, err)
			}
			return nil, fmt.Errorf("%d: %s", resp.StatusCode, responseBody)
		default:
			return resp, nil
Michael Yang's avatar
Michael Yang committed
1175
1176
1177
		}
	}

Michael Yang's avatar
Michael Yang committed
1178
	return nil, errUnauthorized
Michael Yang's avatar
Michael Yang committed
1179
1180
}

Michael Yang's avatar
Michael Yang committed
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
func makeRequest(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.Reader, regOpts *registryOptions) (*http.Response, error) {
	if requestURL.Scheme != "http" && regOpts != nil && regOpts.Insecure {
		requestURL.Scheme = "http"
	}

	req, err := http.NewRequestWithContext(ctx, method, requestURL.String(), body)
	if err != nil {
		return nil, err
	}

	if headers != nil {
		req.Header = headers
	}

	if regOpts != nil {
		if regOpts.Token != "" {
			req.Header.Set("Authorization", "Bearer "+regOpts.Token)
		} else if regOpts.Username != "" && regOpts.Password != "" {
			req.SetBasicAuth(regOpts.Username, regOpts.Password)
		}
	}

	req.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))

	if s := req.Header.Get("Content-Length"); s != "" {
		contentLength, err := strconv.ParseInt(s, 10, 64)
		if err != nil {
			return nil, err
		}

		req.ContentLength = contentLength
	}

Michael Yang's avatar
Michael Yang committed
1214
	resp, err := http.DefaultClient.Do(req)
Michael Yang's avatar
Michael Yang committed
1215
1216
1217
1218
1219
1220
1221
	if err != nil {
		return nil, err
	}

	return resp, nil
}

Patrick Devine's avatar
Patrick Devine committed
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
func getValue(header, key string) string {
	startIdx := strings.Index(header, key+"=")
	if startIdx == -1 {
		return ""
	}

	// Move the index to the starting quote after the key.
	startIdx += len(key) + 2
	endIdx := startIdx

	for endIdx < len(header) {
		if header[endIdx] == '"' {
			if endIdx+1 < len(header) && header[endIdx+1] != ',' { // If the next character isn't a comma, continue
				endIdx++
				continue
			}
			break
		}
		endIdx++
	}
	return header[startIdx:endIdx]
}

Michael Yang's avatar
Michael Yang committed
1245
func parseRegistryChallenge(authStr string) registryChallenge {
Patrick Devine's avatar
Patrick Devine committed
1246
1247
	authStr = strings.TrimPrefix(authStr, "Bearer ")

Michael Yang's avatar
Michael Yang committed
1248
	return registryChallenge{
Patrick Devine's avatar
Patrick Devine committed
1249
1250
1251
1252
1253
1254
		Realm:   getValue(authStr, "realm"),
		Service: getValue(authStr, "service"),
		Scope:   getValue(authStr, "scope"),
	}
}

1255
var errDigestMismatch = errors.New("digest mismatch, file must be downloaded again")
1256

Michael Yang's avatar
Michael Yang committed
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
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 {
1271
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
1272
1273
1274
1275
	}

	return nil
}