images.go 29.1 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
func CopyModel(src, dst model.Name) error {
706
707
708
709
710
711
712
	if !dst.IsFullyQualified() {
		return model.Unqualified(dst)
	}
	if !src.IsFullyQualified() {
		return model.Unqualified(src)
	}

Michael Yang's avatar
Michael Yang committed
713
	manifests, err := GetManifestPath()
714
715
716
717
	if err != nil {
		return err
	}

718
	dstpath := filepath.Join(manifests, dst.Filepath())
Michael Yang's avatar
Michael Yang committed
719
	if err := os.MkdirAll(filepath.Dir(dstpath), 0o755); err != nil {
720
721
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
722

723
	srcpath := filepath.Join(manifests, src.Filepath())
Michael Yang's avatar
Michael Yang committed
724
	srcfile, err := os.Open(srcpath)
Patrick Devine's avatar
Patrick Devine committed
725
726
727
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
728
	defer srcfile.Close()
Patrick Devine's avatar
Patrick Devine committed
729

Michael Yang's avatar
Michael Yang committed
730
	dstfile, err := os.Create(dstpath)
Patrick Devine's avatar
Patrick Devine committed
731
732
733
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
734
	defer dstfile.Close()
Patrick Devine's avatar
Patrick Devine committed
735

Michael Yang's avatar
Michael Yang committed
736
737
	_, err = io.Copy(dstfile, srcfile)
	return err
Patrick Devine's avatar
Patrick Devine committed
738
739
}

Michael Yang's avatar
Michael Yang committed
740
func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{}, dryRun bool) error {
741
742
743
744
	fp, err := GetManifestPath()
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
745
746
747
748

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

Michael Yang's avatar
Michael Yang committed
751
752
753
754
		dir, file := filepath.Split(path)
		dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
		tag := strings.Join([]string{dir, file}, ":")
		fmp := ParseModelPath(tag)
755

Michael Yang's avatar
Michael Yang committed
756
		// skip the manifest we're trying to delete
757
		if skipModelPath != nil && skipModelPath.GetFullTagname() == fmp.GetFullTagname() {
Michael Yang's avatar
Michael Yang committed
758
			return nil
759
		}
Michael Yang's avatar
Michael Yang committed
760
761
762
763

		// 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
764
			// nolint: nilerr
Michael Yang's avatar
Michael Yang committed
765
766
767
768
769
770
771
772
			return nil
		}

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

		delete(deleteMap, manifest.Config.Digest)
773
		return nil
Michael Yang's avatar
Michael Yang committed
774
775
776
	}

	if err := filepath.Walk(fp, walkFunc); err != nil {
Michael Yang's avatar
Michael Yang committed
777
778
		return err
	}
779
780

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

797
798
799
800
	return nil
}

func PruneLayers() error {
Michael Yang's avatar
Michael Yang committed
801
	deleteMap := make(map[string]struct{})
802
803
804
805
806
807
808
	p, err := GetBlobsPath("")
	if err != nil {
		return err
	}

	blobs, err := os.ReadDir(p)
	if err != nil {
809
		slog.Info(fmt.Sprintf("couldn't read dir '%s': %v", p, err))
810
811
812
813
814
		return err
	}

	for _, blob := range blobs {
		name := blob.Name()
815
		name = strings.ReplaceAll(name, "-", ":")
Michael Yang's avatar
Michael Yang committed
816
817
818
		if strings.HasPrefix(name, "sha256:") {
			deleteMap[name] = struct{}{}
		}
819
820
	}

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

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

828
	slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
829
830
831
832

	return nil
}

Michael Yang's avatar
Michael Yang committed
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
859
860
861
862
863
864
865
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
}

866
867
868
869
870
871
872
func DeleteModel(name string) error {
	mp := ParseModelPath(name)
	manifest, _, err := GetManifest(mp)
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
873
	deleteMap := make(map[string]struct{})
874
	for _, layer := range manifest.Layers {
Michael Yang's avatar
Michael Yang committed
875
		deleteMap[layer.Digest] = struct{}{}
876
	}
Michael Yang's avatar
Michael Yang committed
877
	deleteMap[manifest.Config.Digest] = struct{}{}
878
879
880
881
882
883

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

884
	fp, err := mp.GetManifestPath()
885
886
887
888
889
	if err != nil {
		return err
	}
	err = os.Remove(fp)
	if err != nil {
890
		slog.Info(fmt.Sprintf("couldn't remove manifest file '%s': %v", fp, err))
891
892
893
894
895
896
		return err
	}

	return nil
}

Patrick Devine's avatar
Patrick Devine committed
897
func ShowModelfile(model *Model) (string, error) {
Michael Yang's avatar
Michael Yang committed
898
	var mt struct {
Patrick Devine's avatar
Patrick Devine committed
899
		*Model
Michael Yang's avatar
Michael Yang committed
900
		From       string
Michael Yang's avatar
Michael Yang committed
901
		Parameters map[string][]any
Patrick Devine's avatar
Patrick Devine committed
902
903
	}

Michael Yang's avatar
Michael Yang committed
904
	mt.Parameters = make(map[string][]any)
Patrick Devine's avatar
Patrick Devine committed
905
	for k, v := range model.Options {
Michael Yang's avatar
Michael Yang committed
906
907
908
		if s, ok := v.([]any); ok {
			mt.Parameters[k] = s
			continue
Patrick Devine's avatar
Patrick Devine committed
909
910
		}

Michael Yang's avatar
Michael Yang committed
911
		mt.Parameters[k] = []any{v}
Patrick Devine's avatar
Patrick Devine committed
912
913
	}

Michael Yang's avatar
Michael Yang committed
914
915
	mt.Model = model
	mt.From = model.ModelPath
Patrick Devine's avatar
Patrick Devine committed
916

917
918
	if model.ParentModel != "" {
		mt.From = model.ParentModel
Patrick Devine's avatar
Patrick Devine committed
919
920
921
922
923
924
925
926
	}

	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 }}"""
927
928

{{- if .System }}
Patrick Devine's avatar
Patrick Devine committed
929
SYSTEM """{{ .System }}"""
930
{{- end }}
931
932
933
934

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

Michael Yang's avatar
Michael Yang committed
936
937
938
939
{{- range $k, $v := .Parameters }}
{{- range $parameter := $v }}
PARAMETER {{ $k }} {{ printf "%#v" $parameter }}
{{- end }}
Michael Yang's avatar
Michael Yang committed
940
{{- end }}`
Patrick Devine's avatar
Patrick Devine committed
941
942
943

	tmpl, err := template.New("").Parse(modelFile)
	if err != nil {
944
		slog.Info(fmt.Sprintf("error parsing template: %q", err))
Patrick Devine's avatar
Patrick Devine committed
945
946
947
948
949
950
		return "", err
	}

	var buf bytes.Buffer

	if err = tmpl.Execute(&buf, mt); err != nil {
951
		slog.Info(fmt.Sprintf("error executing template: %q", err))
Patrick Devine's avatar
Patrick Devine committed
952
953
954
955
956
957
		return "", err
	}

	return buf.String(), nil
}

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

962
963
964
965
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
	}

Patrick Devine's avatar
Patrick Devine committed
966
	manifest, _, err := GetManifest(mp)
967
	if err != nil {
968
		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
969
970
971
972
		return err
	}

	var layers []*Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
973
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
974
	layers = append(layers, manifest.Config)
975
976

	for _, layer := range layers {
Michael Yang's avatar
Michael Yang committed
977
		if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
978
			slog.Info(fmt.Sprintf("error uploading blob: %v", err))
979
980
981
			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())
			}
982
983
			return err
		}
984
985
	}

986
	fn(api.ProgressResponse{Status: "pushing manifest"})
Michael Yang's avatar
Michael Yang committed
987
988
	requestURL := mp.BaseURL()
	requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
989
990
991
992
993
994

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

Michael Yang's avatar
Michael Yang committed
995
996
	headers := make(http.Header)
	headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
997
	resp, err := makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
998
999
1000
1001
1002
	if err != nil {
		return err
	}
	defer resp.Body.Close()

1003
	fn(api.ProgressResponse{Status: "success"})
1004
1005
1006
1007

	return nil
}

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

1011
1012
1013
1014
1015
	var manifest *ManifestV2
	var err error
	var noprune string

	// build deleteMap to prune unused layers
Michael Yang's avatar
Michael Yang committed
1016
	deleteMap := make(map[string]struct{})
1017
1018
1019
1020
1021
1022
1023
1024
1025

	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
1026
				deleteMap[l.Digest] = struct{}{}
1027
			}
Michael Yang's avatar
Michael Yang committed
1028
			deleteMap[manifest.Config.Digest] = struct{}{}
1029
1030
1031
		}
	}

1032
1033
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
1034
	}
1035

1036
	fn(api.ProgressResponse{Status: "pulling manifest"})
1037

1038
	manifest, err = pullModelManifest(ctx, mp, regOpts)
1039
	if err != nil {
1040
		return fmt.Errorf("pull model manifest: %s", err)
1041
1042
1043
	}

	var layers []*Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
1044
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
1045
	layers = append(layers, manifest.Config)
1046
1047

	for _, layer := range layers {
1048
1049
1050
1051
1052
1053
1054
1055
		if err := downloadBlob(
			ctx,
			downloadOpts{
				mp:      mp,
				digest:  layer.Digest,
				regOpts: regOpts,
				fn:      fn,
			}); err != nil {
1056
1057
			return err
		}
1058
		delete(deleteMap, layer.Digest)
1059
	}
1060
	delete(deleteMap, manifest.Config.Digest)
1061

Michael Yang's avatar
Michael Yang committed
1062
1063
1064
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
		if err := verifyBlob(layer.Digest); err != nil {
1065
1066
1067
1068
1069
1070
1071
1072
			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
1073
					slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err))
1074
1075
				}
			}
Michael Yang's avatar
Michael Yang committed
1076
1077
1078
1079
			return err
		}
	}

1080
	fn(api.ProgressResponse{Status: "writing manifest"})
1081

1082
	manifestJSON, err := json.Marshal(manifest)
1083
1084
1085
1086
	if err != nil {
		return err
	}

1087
	fp, err := mp.GetManifestPath()
1088
1089
1090
	if err != nil {
		return err
	}
1091
1092
1093
	if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
		return err
	}
1094

Bruce MacDonald's avatar
Bruce MacDonald committed
1095
	err = os.WriteFile(fp, manifestJSON, 0o644)
1096
	if err != nil {
1097
		slog.Info(fmt.Sprintf("couldn't write to %s", fp))
1098
1099
1100
		return err
	}

1101
1102
1103
1104
1105
1106
1107
1108
	if noprune == "" {
		fn(api.ProgressResponse{Status: "removing any unused layers"})
		err = deleteUnusedLayers(nil, deleteMap, false)
		if err != nil {
			return err
		}
	}

1109
	fn(api.ProgressResponse{Status: "success"})
1110
1111
1112
1113

	return nil
}

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

Michael Yang's avatar
Michael Yang committed
1117
1118
	headers := make(http.Header)
	headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
1119
	resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
	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
1134
func GetSHA256Digest(r io.Reader) (string, int64) {
Michael Yang's avatar
Michael Yang committed
1135
1136
1137
1138
1139
1140
	h := sha256.New()
	n, err := io.Copy(h, r)
	if err != nil {
		log.Fatal(err)
	}

Michael Yang's avatar
Michael Yang committed
1141
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
1142
1143
}

1144
var errUnauthorized = errors.New("unauthorized")
1145

Michael Yang's avatar
Michael Yang committed
1146
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
1147
	for i := 0; i < 2; i++ {
Michael Yang's avatar
Michael Yang committed
1148
		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
Michael Yang's avatar
Michael Yang committed
1149
		if err != nil {
Michael Yang's avatar
Michael Yang committed
1150
			if !errors.Is(err, context.Canceled) {
1151
				slog.Info(fmt.Sprintf("request failed: %v", err))
Michael Yang's avatar
Michael Yang committed
1152
1153
			}

Michael Yang's avatar
Michael Yang committed
1154
1155
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
1156
1157
1158
1159

		switch {
		case resp.StatusCode == http.StatusUnauthorized:
			// Handle authentication error with one retry
Michael Yang's avatar
Michael Yang committed
1160
1161
			challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
			token, err := getAuthorizationToken(ctx, challenge)
Michael Yang's avatar
Michael Yang committed
1162
1163
1164
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
			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
1182
1183
1184
		}
	}

Michael Yang's avatar
Michael Yang committed
1185
	return nil, errUnauthorized
Michael Yang's avatar
Michael Yang committed
1186
1187
}

Michael Yang's avatar
Michael Yang committed
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
1214
1215
1216
1217
1218
1219
1220
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
1221
	resp, err := http.DefaultClient.Do(req)
Michael Yang's avatar
Michael Yang committed
1222
1223
1224
1225
1226
1227
1228
	if err != nil {
		return nil, err
	}

	return resp, nil
}

Patrick Devine's avatar
Patrick Devine committed
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
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
1252
func parseRegistryChallenge(authStr string) registryChallenge {
Patrick Devine's avatar
Patrick Devine committed
1253
1254
	authStr = strings.TrimPrefix(authStr, "Bearer ")

Michael Yang's avatar
Michael Yang committed
1255
	return registryChallenge{
Patrick Devine's avatar
Patrick Devine committed
1256
1257
1258
1259
1260
1261
		Realm:   getValue(authStr, "realm"),
		Service: getValue(authStr, "service"),
		Scope:   getValue(authStr, "scope"),
	}
}

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

Michael Yang's avatar
Michael Yang committed
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
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 {
1278
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
1279
1280
1281
1282
	}

	return nil
}