images.go 28.6 KB
Newer Older
1
2
3
4
package server

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

26
	"github.com/ollama/ollama/api"
27
	"github.com/ollama/ollama/auth"
Michael Yang's avatar
Michael Yang committed
28
	"github.com/ollama/ollama/envconfig"
Michael Yang's avatar
Michael Yang committed
29
	"github.com/ollama/ollama/format"
30
	"github.com/ollama/ollama/llama"
31
	"github.com/ollama/ollama/llm"
32
	"github.com/ollama/ollama/parser"
Michael Yang's avatar
Michael Yang committed
33
	"github.com/ollama/ollama/template"
34
	"github.com/ollama/ollama/types/errtypes"
Michael Yang's avatar
Michael Yang committed
35
	"github.com/ollama/ollama/types/model"
36
	"github.com/ollama/ollama/version"
37
38
)

39
40
41
42
43
44
var (
	errCapabilities         = errors.New("does not support")
	errCapabilityCompletion = errors.New("completion")
	errCapabilityTools      = errors.New("tools")
	errCapabilityInsert     = errors.New("insert")
)
Michael Yang's avatar
Michael Yang committed
45

Michael Yang's avatar
Michael Yang committed
46
47
type Capability string

Michael Yang's avatar
tools  
Michael Yang committed
48
49
50
const (
	CapabilityCompletion = Capability("completion")
	CapabilityTools      = Capability("tools")
51
	CapabilityInsert     = Capability("insert")
Michael Yang's avatar
tools  
Michael Yang committed
52
)
Michael Yang's avatar
Michael Yang committed
53

Michael Yang's avatar
Michael Yang committed
54
55
56
57
58
type registryOptions struct {
	Insecure bool
	Username string
	Password string
	Token    string
59
60

	CheckRedirect func(req *http.Request, via []*http.Request) error
Michael Yang's avatar
Michael Yang committed
61
62
}

63
type Model struct {
Michael Yang's avatar
Michael Yang committed
64
	Name           string `json:"name"`
65
	Config         ConfigV2
Michael Yang's avatar
Michael Yang committed
66
67
	ShortName      string
	ModelPath      string
68
	ParentModel    string
Michael Yang's avatar
Michael Yang committed
69
70
71
72
73
74
	AdapterPaths   []string
	ProjectorPaths []string
	System         string
	License        []string
	Digest         string
	Options        map[string]interface{}
Michael Yang's avatar
Michael Yang committed
75
	Messages       []api.Message
Michael Yang's avatar
Michael Yang committed
76
77

	Template *template.Template
78
79
}

Michael Yang's avatar
Michael Yang committed
80
81
82
83
// CheckCapabilities checks if the model has the specified capabilities returning an error describing
// any missing or unknown capabilities
func (m *Model) CheckCapabilities(caps ...Capability) error {
	var errs []error
Michael Yang's avatar
Michael Yang committed
84
85
86
	for _, cap := range caps {
		switch cap {
		case CapabilityCompletion:
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
			f, err := os.Open(m.ModelPath)
			if err != nil {
				slog.Error("couldn't open model file", "error", err)
				continue
			}
			defer f.Close()

			// TODO(mxyng): decode the GGML into model to avoid doing this multiple times
			ggml, _, err := llm.DecodeGGML(f, 0)
			if err != nil {
				slog.Error("couldn't decode ggml", "error", err)
				continue
			}

			if _, ok := ggml.KV()[fmt.Sprintf("%s.pooling_type", ggml.KV().Architecture())]; ok {
Michael Yang's avatar
Michael Yang committed
102
				errs = append(errs, errCapabilityCompletion)
Michael Yang's avatar
Michael Yang committed
103
			}
Michael Yang's avatar
tools  
Michael Yang committed
104
105
		case CapabilityTools:
			if !slices.Contains(m.Template.Vars(), "tools") {
106
107
108
109
110
111
				errs = append(errs, errCapabilityTools)
			}
		case CapabilityInsert:
			vars := m.Template.Vars()
			if !slices.Contains(vars, "suffix") {
				errs = append(errs, errCapabilityInsert)
Michael Yang's avatar
tools  
Michael Yang committed
112
			}
Michael Yang's avatar
Michael Yang committed
113
114
		default:
			slog.Error("unknown capability", "capability", cap)
Michael Yang's avatar
Michael Yang committed
115
			return fmt.Errorf("unknown capability: %s", cap)
Michael Yang's avatar
Michael Yang committed
116
117
118
		}
	}

Michael Yang's avatar
Michael Yang committed
119
	if err := errors.Join(errs...); err != nil {
120
		return fmt.Errorf("%w %w", errCapabilities, errors.Join(errs...))
Michael Yang's avatar
Michael Yang committed
121
122
123
	}

	return nil
124
125
}

Michael Yang's avatar
Michael Yang committed
126
func (m *Model) String() string {
127
	var modelfile parser.File
Michael Yang's avatar
Michael Yang committed
128

129
	modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
130
131
132
		Name: "model",
		Args: m.ModelPath,
	})
133

Michael Yang's avatar
Michael Yang committed
134
	for _, adapter := range m.AdapterPaths {
135
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
136
137
			Name: "adapter",
			Args: adapter,
Michael Yang's avatar
Michael Yang committed
138
		})
139
140
	}

Michael Yang's avatar
Michael Yang committed
141
	for _, projector := range m.ProjectorPaths {
142
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
143
144
			Name: "model",
			Args: projector,
Michael Yang's avatar
Michael Yang committed
145
		})
146
147
	}

Michael Yang's avatar
Michael Yang committed
148
	if m.Template != nil {
149
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
150
			Name: "template",
Michael Yang's avatar
Michael Yang committed
151
			Args: m.Template.String(),
Michael Yang's avatar
Michael Yang committed
152
		})
153
154
	}

Michael Yang's avatar
Michael Yang committed
155
	if m.System != "" {
156
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
157
158
			Name: "system",
			Args: m.System,
Michael Yang's avatar
Michael Yang committed
159
		})
160
161
162
163
164
165
	}

	for k, v := range m.Options {
		switch v := v.(type) {
		case []any:
			for _, s := range v {
166
				modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
167
168
169
					Name: k,
					Args: fmt.Sprintf("%v", s),
				})
170
171
			}
		default:
172
			modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
173
174
175
				Name: k,
				Args: fmt.Sprintf("%v", v),
			})
176
177
178
179
		}
	}

	for _, license := range m.License {
180
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
181
182
183
			Name: "license",
			Args: license,
		})
184
185
186
	}

	for _, msg := range m.Messages {
187
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
188
			Name: "message",
Michael Yang's avatar
Michael Yang committed
189
			Args: fmt.Sprintf("%s: %s", msg.Role, msg.Content),
Michael Yang's avatar
Michael Yang committed
190
		})
191
192
	}

Michael Yang's avatar
Michael Yang committed
193
	return modelfile.String()
194
195
}

196
type ConfigV2 struct {
197
198
199
200
201
202
	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"`

203
	// required by spec
204
205
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
206
	RootFS       RootFS `json:"rootfs"`
207
208
209
210
211
212
213
}

type RootFS struct {
	Type    string   `json:"type"`
	DiffIDs []string `json:"diff_ids"`
}

Michael Yang's avatar
Michael Yang committed
214
func GetManifest(mp ModelPath) (*Manifest, string, error) {
215
	fp, err := mp.GetManifestPath()
216
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
217
		return nil, "", err
218
	}
219

Michael Yang's avatar
Michael Yang committed
220
	f, err := os.Open(fp)
221
	if err != nil {
Michael Yang's avatar
Michael Yang committed
222
		return nil, "", err
223
	}
Michael Yang's avatar
Michael Yang committed
224
	defer f.Close()
225

Michael Yang's avatar
Michael Yang committed
226
	sha256sum := sha256.New()
Patrick Devine's avatar
Patrick Devine committed
227

Michael Yang's avatar
Michael Yang committed
228
229
	var manifest Manifest
	if err := json.NewDecoder(io.TeeReader(f, sha256sum)).Decode(&manifest); err != nil {
Patrick Devine's avatar
Patrick Devine committed
230
		return nil, "", err
231
232
	}

Michael Yang's avatar
Michael Yang committed
233
	return &manifest, hex.EncodeToString(sha256sum.Sum(nil)), nil
234
235
236
}

func GetModel(name string) (*Model, error) {
237
	mp := ParseModelPath(name)
Patrick Devine's avatar
Patrick Devine committed
238
	manifest, digest, err := GetManifest(mp)
239
240
241
242
243
	if err != nil {
		return nil, err
	}

	model := &Model{
244
245
246
		Name:      mp.GetFullTagname(),
		ShortName: mp.GetShortTagname(),
		Digest:    digest,
Michael Yang's avatar
Michael Yang committed
247
		Template:  template.DefaultTemplate,
248
249
	}

250
251
252
253
254
	if manifest.Config.Digest != "" {
		filename, err := GetBlobsPath(manifest.Config.Digest)
		if err != nil {
			return nil, err
		}
255

256
257
258
259
260
		configFile, err := os.Open(filename)
		if err != nil {
			return nil, err
		}
		defer configFile.Close()
261

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

267
	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
268
		filename, err := GetBlobsPath(layer.Digest)
269
270
271
272
		if err != nil {
			return nil, err
		}

273
274
275
		switch layer.MediaType {
		case "application/vnd.ollama.image.model":
			model.ModelPath = filename
276
			model.ParentModel = layer.From
277
		case "application/vnd.ollama.image.embed":
278
279
			// Deprecated in versions  > 0.1.2
			// TODO: remove this warning in a future version
280
			slog.Info("WARNING: model contains embeddings, but embeddings in modelfiles have been deprecated and will be ignored.")
281
282
		case "application/vnd.ollama.image.adapter":
			model.AdapterPaths = append(model.AdapterPaths, filename)
Michael Yang's avatar
Michael Yang committed
283
284
		case "application/vnd.ollama.image.projector":
			model.ProjectorPaths = append(model.ProjectorPaths, filename)
Michael Yang's avatar
Michael Yang committed
285
286
		case "application/vnd.ollama.image.prompt",
			"application/vnd.ollama.image.template":
287
288
289
290
291
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

Michael Yang's avatar
Michael Yang committed
292
			model.Template, err = template.Parse(string(bts))
293
294
295
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
296
		case "application/vnd.ollama.image.system":
297
298
299
300
301
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

Michael Yang's avatar
Michael Yang committed
302
			model.System = string(bts)
303
		case "application/vnd.ollama.image.params":
Michael Yang's avatar
Michael Yang committed
304
305
306
307
308
			params, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer params.Close()
309

310
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
311
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
312
313
				return nil, err
			}
314
315
316
317
318
319
320
321
322
323
		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
324
325
326
327
328
329
		case "application/vnd.ollama.image.license":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}
			model.License = append(model.License, string(bts))
330
331
332
333
334
335
		}
	}

	return model, nil
}

Michael Yang's avatar
Michael Yang committed
336
func realpath(rel, from string) string {
337
	abspath, err := filepath.Abs(from)
Michael Yang's avatar
Michael Yang committed
338
	if err != nil {
339
		return from
340
341
	}

Michael Yang's avatar
Michael Yang committed
342
	home, err := os.UserHomeDir()
343
	if err != nil {
Michael Yang's avatar
Michael Yang committed
344
		return abspath
345
346
	}

347
	if from == "~" {
Michael Yang's avatar
Michael Yang committed
348
		return home
349
350
351
352
	} else if strings.HasPrefix(from, "~/") {
		return filepath.Join(home, from[2:])
	}

Michael Yang's avatar
Michael Yang committed
353
	if _, err := os.Stat(filepath.Join(rel, from)); err == nil {
354
		// this is a file relative to the Modelfile
Michael Yang's avatar
Michael Yang committed
355
		return filepath.Join(rel, from)
356
357
	}

Michael Yang's avatar
Michael Yang committed
358
359
360
	return abspath
}

361
func CreateModel(ctx context.Context, name model.Name, modelFileDir, quantization string, modelfile *parser.File, fn func(resp api.ProgressResponse)) (err error) {
362
363
	config := ConfigV2{
		OS:           "linux",
Michael Yang's avatar
Michael Yang committed
364
		Architecture: "amd64",
Michael Yang's avatar
Michael Yang committed
365
366
367
		RootFS: RootFS{
			Type: "layers",
		},
368
369
	}

Michael Yang's avatar
Michael Yang committed
370
371
	var messages []*api.Message
	parameters := make(map[string]any)
Michael Yang's avatar
Michael Yang committed
372

373
	var layers []Layer
374
	var baseLayers []*layerGGML
Michael Yang's avatar
Michael Yang committed
375
	for _, c := range modelfile.Commands {
Michael Yang's avatar
Michael Yang committed
376
		mediatype := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)
377
		command := c.Name
Michael Yang's avatar
Michael Yang committed
378

379
		switch command {
Michael Yang's avatar
Michael Yang committed
380
		case "model", "adapter":
381
			if name := model.ParseName(c.Args); name.IsValid() && command == "model" {
Michael Yang's avatar
Michael Yang committed
382
				baseLayers, err = parseFromModel(ctx, name, fn)
Michael Yang's avatar
Michael Yang committed
383
384
385
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
386
			} else if strings.HasPrefix(c.Args, "@") {
387
				digest := strings.TrimPrefix(c.Args, "@")
Michael Yang's avatar
Michael Yang committed
388
389
				if ib, ok := intermediateBlobs[digest]; ok {
					p, err := GetBlobsPath(ib)
390
391
392
393
394
395
396
397
398
					if err != nil {
						return err
					}

					if _, err := os.Stat(p); errors.Is(err, os.ErrNotExist) {
						// pass
					} else if err != nil {
						return err
					} else {
Michael Yang's avatar
Michael Yang committed
399
400
						fn(api.ProgressResponse{Status: fmt.Sprintf("using cached layer %s", ib)})
						digest = ib
401
402
403
404
					}
				}

				blobpath, err := GetBlobsPath(digest)
Michael Yang's avatar
Michael Yang committed
405
				if err != nil {
Michael Yang's avatar
Michael Yang committed
406
					return err
407
				}
408

Michael Yang's avatar
Michael Yang committed
409
				blob, err := os.Open(blobpath)
Michael Yang's avatar
Michael Yang committed
410
411
412
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
413
				defer blob.Close()
Michael Yang's avatar
Michael Yang committed
414

415
				baseLayers, err = parseFromFile(ctx, command, baseLayers, blob, digest, fn)
Michael Yang's avatar
Michael Yang committed
416
417
418
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
419
420
			} else if file, err := os.Open(realpath(modelFileDir, c.Args)); err == nil {
				defer file.Close()
Michael Yang's avatar
Michael Yang committed
421

422
				baseLayers, err = parseFromFile(ctx, command, baseLayers, file, "", fn)
Michael Yang's avatar
Michael Yang committed
423
				if err != nil {
Michael Yang's avatar
Michael Yang committed
424
425
					return err
				}
Michael Yang's avatar
Michael Yang committed
426
427
428
			} else {
				return fmt.Errorf("invalid model reference: %s", c.Args)
			}
Michael Yang's avatar
Michael Yang committed
429

Michael Yang's avatar
Michael Yang committed
430
			for _, baseLayer := range baseLayers {
Michael Yang's avatar
Michael Yang committed
431
432
433
434
				if quantization != "" &&
					baseLayer.MediaType == "application/vnd.ollama.image.model" &&
					baseLayer.GGML != nil &&
					baseLayer.GGML.Name() == "gguf" {
Michael Yang's avatar
Michael Yang committed
435
					want, err := llm.ParseFileType(quantization)
436
					if err != nil {
Michael Yang's avatar
Michael Yang committed
437
						return err
438
					}
Michael Yang's avatar
Michael Yang committed
439

Michael Yang's avatar
Michael Yang committed
440
441
					ft := baseLayer.GGML.KV().FileType()
					if !slices.Contains([]string{"F16", "F32"}, ft.String()) {
Michael Yang's avatar
Michael Yang committed
442
						return errors.New("quantization is only supported for F16 and F32 models")
Michael Yang's avatar
Michael Yang committed
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
					} else if want != ft {
						fn(api.ProgressResponse{Status: fmt.Sprintf("quantizing %s model to %s", ft, quantization)})

						blob, err := GetBlobsPath(baseLayer.Digest)
						if err != nil {
							return err
						}

						temp, err := os.CreateTemp(filepath.Dir(blob), quantization)
						if err != nil {
							return err
						}
						defer temp.Close()
						defer os.Remove(temp.Name())

458
						if err := llama.Quantize(blob, temp.Name(), uint32(want)); err != nil {
Michael Yang's avatar
Michael Yang committed
459
460
461
							return err
						}

Michael Yang's avatar
Michael Yang committed
462
						layer, err := NewLayer(temp, baseLayer.MediaType)
Michael Yang's avatar
Michael Yang committed
463
464
465
						if err != nil {
							return err
						}
Michael Yang's avatar
Michael Yang committed
466

Michael Yang's avatar
Michael Yang committed
467
468
469
470
						if _, err := temp.Seek(0, io.SeekStart); err != nil {
							return err
						}

471
						ggml, _, err := llm.DecodeGGML(temp, 0)
Michael Yang's avatar
Michael Yang committed
472
473
						if err != nil {
							return err
Michael Yang's avatar
Michael Yang committed
474
475
						}

Michael Yang's avatar
Michael Yang committed
476
477
						baseLayer.Layer = layer
						baseLayer.GGML = ggml
Michael Yang's avatar
Michael Yang committed
478
					}
479
				}
480

Michael Yang's avatar
Michael Yang committed
481
482
483
484
				if baseLayer.GGML != nil {
					config.ModelFormat = cmp.Or(config.ModelFormat, baseLayer.GGML.Name())
					config.ModelFamily = cmp.Or(config.ModelFamily, baseLayer.GGML.KV().Architecture())
					config.ModelType = cmp.Or(config.ModelType, format.HumanNumber(baseLayer.GGML.KV().ParameterCount()))
Michael Yang's avatar
Michael Yang committed
485
					config.FileType = cmp.Or(config.FileType, baseLayer.GGML.KV().FileType().String())
Michael Yang's avatar
Michael Yang committed
486
					config.ModelFamilies = append(config.ModelFamilies, baseLayer.GGML.KV().Architecture())
487
488
				}

Michael Yang's avatar
Michael Yang committed
489
				layers = append(layers, baseLayer.Layer)
490
			}
Michael Yang's avatar
Michael Yang committed
491
		case "license", "template", "system":
Josh's avatar
Josh committed
492
493
494
495
496
497
			if c.Name == "template" {
				if _, err := template.Parse(c.Args); err != nil {
					return fmt.Errorf("%w: %s", errBadTemplate, err)
				}
			}

Michael Yang's avatar
Michael Yang committed
498
499
			if c.Name != "license" {
				// replace
500
				layers = slices.DeleteFunc(layers, func(layer Layer) bool {
501
502
503
504
505
506
507
508
509
					if layer.MediaType != mediatype {
						return false
					}

					if err := layer.Remove(); err != nil {
						return false
					}

					return true
Michael Yang's avatar
Michael Yang committed
510
				})
Michael Yang's avatar
Michael Yang committed
511
512
			}

513
514
515
516
517
518
			blob := strings.NewReader(c.Args)
			layer, err := NewLayer(blob, mediatype)
			if err != nil {
				return err
			}

Michael Yang's avatar
Michael Yang committed
519
520
521
522
523
524
525
526
527
528
			layers = append(layers, layer)
		case "message":
			role, content, ok := strings.Cut(c.Args, ": ")
			if !ok {
				return fmt.Errorf("invalid message: %s", c.Args)
			}

			messages = append(messages, &api.Message{Role: role, Content: content})
		default:
			ps, err := api.FormatParams(map[string][]string{c.Name: {c.Args}})
529
			if err != nil {
Michael Yang's avatar
Michael Yang committed
530
				return err
531
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
532

Michael Yang's avatar
Michael Yang committed
533
534
535
536
537
538
539
540
541
542
543
			for k, v := range ps {
				if ks, ok := parameters[k].([]string); ok {
					parameters[k] = append(ks, v.([]string)...)
				} else if vs, ok := v.([]string); ok {
					parameters[k] = vs
				} else {
					parameters[k] = v
				}
			}
		}
	}
Michael Yang's avatar
Michael Yang committed
544

Michael Yang's avatar
Michael Yang committed
545
	var err2 error
546
	layers = slices.DeleteFunc(layers, func(layer Layer) bool {
Michael Yang's avatar
Michael Yang committed
547
548
549
550
551
552
553
554
555
556
557
		switch layer.MediaType {
		case "application/vnd.ollama.image.message":
			// if there are new messages, remove the inherited ones
			if len(messages) > 0 {
				return true
			}

			return false
		case "application/vnd.ollama.image.params":
			// merge inherited parameters with new ones
			r, err := layer.Open()
Bruce MacDonald's avatar
Bruce MacDonald committed
558
			if err != nil {
Michael Yang's avatar
Michael Yang committed
559
560
				err2 = err
				return false
Bruce MacDonald's avatar
Bruce MacDonald committed
561
			}
Michael Yang's avatar
Michael Yang committed
562
			defer r.Close()
Bruce MacDonald's avatar
Bruce MacDonald committed
563

Michael Yang's avatar
Michael Yang committed
564
565
566
567
568
			var ps map[string]any
			if err := json.NewDecoder(r).Decode(&ps); err != nil {
				err2 = err
				return false
			}
Michael Yang's avatar
Michael Yang committed
569

Michael Yang's avatar
Michael Yang committed
570
571
572
573
			for k, v := range ps {
				if _, ok := parameters[k]; !ok {
					parameters[k] = v
				}
574
			}
575

Michael Yang's avatar
Michael Yang committed
576
			return true
577
		default:
Michael Yang's avatar
Michael Yang committed
578
			return false
579
		}
Michael Yang's avatar
Michael Yang committed
580
581
582
583
	})

	if err2 != nil {
		return err2
584
585
	}

586
587
	if len(messages) > 0 {
		var b bytes.Buffer
Michael Yang's avatar
Michael Yang committed
588
		if err := json.NewEncoder(&b).Encode(messages); err != nil {
589
590
591
592
593
594
595
596
			return err
		}

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

Michael Yang's avatar
Michael Yang committed
597
		layers = append(layers, layer)
598
599
	}

Michael Yang's avatar
Michael Yang committed
600
	if len(parameters) > 0 {
Michael Yang's avatar
Michael Yang committed
601
		var b bytes.Buffer
Michael Yang's avatar
Michael Yang committed
602
		if err := json.NewEncoder(&b).Encode(parameters); err != nil {
603
604
605
			return err
		}

Michael Yang's avatar
Michael Yang committed
606
		layer, err := NewLayer(&b, "application/vnd.ollama.image.params")
607
		if err != nil {
Michael Yang's avatar
Michael Yang committed
608
			return err
609
		}
Michael Yang's avatar
Michael Yang committed
610

Michael Yang's avatar
Michael Yang committed
611
		layers = append(layers, layer)
612
613
	}

Michael Yang's avatar
Michael Yang committed
614
615
	digests := make([]string, len(layers))
	for i, layer := range layers {
Michael Yang's avatar
Michael Yang committed
616
		digests[i] = layer.Digest
617
618
	}

Michael Yang's avatar
Michael Yang committed
619
	config.RootFS.DiffIDs = digests
Michael Yang's avatar
Michael Yang committed
620

Michael Yang's avatar
Michael Yang committed
621
622
	var b bytes.Buffer
	if err := json.NewEncoder(&b).Encode(config); err != nil {
623
624
625
		return err
	}

626
	configLayer, err := NewLayer(&b, "application/vnd.docker.container.image.v1+json")
Michael Yang's avatar
Michael Yang committed
627
	if err != nil {
628
629
630
		return err
	}

631
	for _, layer := range append(layers, configLayer) {
Michael Yang's avatar
Michael Yang committed
632
633
		if layer.status != "" {
			fn(api.ProgressResponse{Status: layer.status})
634
		}
Michael Yang's avatar
Michael Yang committed
635
	}
636

637
	old, _ := ParseNamedManifest(name)
638

Michael Yang's avatar
Michael Yang committed
639
	fn(api.ProgressResponse{Status: "writing manifest"})
640
	if err := WriteManifest(name, configLayer, layers); err != nil {
641
642
		return err
	}
643

Michael Yang's avatar
bool  
Michael Yang committed
644
	if !envconfig.NoPrune() && old != nil {
645
		if err := old.RemoveLayers(); err != nil {
Michael Yang's avatar
Michael Yang committed
646
			return err
647
648
649
		}
	}

Michael Yang's avatar
Michael Yang committed
650
651
	fn(api.ProgressResponse{Status: "success"})
	return nil
652
653
}

Michael Yang's avatar
Michael Yang committed
654
func CopyModel(src, dst model.Name) error {
655
656
657
658
659
660
661
	if !dst.IsFullyQualified() {
		return model.Unqualified(dst)
	}
	if !src.IsFullyQualified() {
		return model.Unqualified(src)
	}

662
663
664
665
	if src.Filepath() == dst.Filepath() {
		return nil
	}

Michael Yang's avatar
Michael Yang committed
666
	manifests, err := GetManifestPath()
667
668
669
670
	if err != nil {
		return err
	}

671
	dstpath := filepath.Join(manifests, dst.Filepath())
Michael Yang's avatar
Michael Yang committed
672
	if err := os.MkdirAll(filepath.Dir(dstpath), 0o755); err != nil {
673
674
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
675

676
	srcpath := filepath.Join(manifests, src.Filepath())
Michael Yang's avatar
Michael Yang committed
677
	srcfile, err := os.Open(srcpath)
Patrick Devine's avatar
Patrick Devine committed
678
679
680
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
681
	defer srcfile.Close()
Patrick Devine's avatar
Patrick Devine committed
682

Michael Yang's avatar
Michael Yang committed
683
	dstfile, err := os.Create(dstpath)
Patrick Devine's avatar
Patrick Devine committed
684
685
686
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
687
	defer dstfile.Close()
Patrick Devine's avatar
Patrick Devine committed
688

Michael Yang's avatar
Michael Yang committed
689
690
	_, err = io.Copy(dstfile, srcfile)
	return err
Patrick Devine's avatar
Patrick Devine committed
691
692
}

Michael Yang's avatar
Michael Yang committed
693
func deleteUnusedLayers(deleteMap map[string]struct{}) error {
694
695
	// Ignore corrupt manifests to avoid blocking deletion of layers that are freshly orphaned
	manifests, err := Manifests(true)
696
697
698
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
699

Michael Yang's avatar
Michael Yang committed
700
	for _, manifest := range manifests {
Michael Yang's avatar
Michael Yang committed
701
702
703
704
705
		for _, layer := range manifest.Layers {
			delete(deleteMap, layer.Digest)
		}

		delete(deleteMap, manifest.Config.Digest)
Michael Yang's avatar
Michael Yang committed
706
	}
707
708

	// only delete the files which are still in the deleteMap
Michael Yang's avatar
Michael Yang committed
709
710
711
	for k := range deleteMap {
		fp, err := GetBlobsPath(k)
		if err != nil {
712
			slog.Info(fmt.Sprintf("couldn't get file path for '%s': %v", k, err))
Michael Yang's avatar
Michael Yang committed
713
714
			continue
		}
715
716
717
		if err := os.Remove(fp); err != nil {
			slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err))
			continue
718
719
720
		}
	}

721
722
723
724
	return nil
}

func PruneLayers() error {
Michael Yang's avatar
Michael Yang committed
725
	deleteMap := make(map[string]struct{})
726
727
728
729
730
731
732
	p, err := GetBlobsPath("")
	if err != nil {
		return err
	}

	blobs, err := os.ReadDir(p)
	if err != nil {
733
		slog.Info(fmt.Sprintf("couldn't read dir '%s': %v", p, err))
734
735
736
737
738
		return err
	}

	for _, blob := range blobs {
		name := blob.Name()
739
		name = strings.ReplaceAll(name, "-", ":")
740
741
742
743
744
745
746
747
748
749
750

		_, err := GetBlobsPath(name)
		if err != nil {
			if errors.Is(err, ErrInvalidDigestFormat) {
				// remove invalid blobs (e.g. partial downloads)
				if err := os.Remove(filepath.Join(p, blob.Name())); err != nil {
					slog.Error("couldn't remove blob", "blob", blob.Name(), "error", err)
				}
			}

			continue
Michael Yang's avatar
Michael Yang committed
751
		}
752
753

		deleteMap[name] = struct{}{}
754
755
	}

756
	slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap)))
757

Michael Yang's avatar
Michael Yang committed
758
	if err := deleteUnusedLayers(deleteMap); err != nil {
759
		slog.Error(fmt.Sprintf("couldn't remove unused layers: %v", err))
760
		return nil
761
762
	}

763
	slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
764
765
766
767

	return nil
}

Michael Yang's avatar
Michael Yang committed
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
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
}

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

805
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
Michael Yang's avatar
lint  
Michael Yang committed
806
		return errors.New("insecure protocol http")
807
808
	}

Patrick Devine's avatar
Patrick Devine committed
809
	manifest, _, err := GetManifest(mp)
810
	if err != nil {
811
		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
812
813
814
		return err
	}

815
	var layers []Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
816
	layers = append(layers, manifest.Layers...)
817
	if manifest.Config.Digest != "" {
818
		layers = append(layers, manifest.Config)
819
	}
820
821

	for _, layer := range layers {
Michael Yang's avatar
Michael Yang committed
822
		if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
823
			slog.Info(fmt.Sprintf("error uploading blob: %v", err))
824
825
			return err
		}
826
827
	}

828
	fn(api.ProgressResponse{Status: "pushing manifest"})
Michael Yang's avatar
Michael Yang committed
829
830
	requestURL := mp.BaseURL()
	requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
831
832
833
834
835
836

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

Michael Yang's avatar
Michael Yang committed
837
838
	headers := make(http.Header)
	headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
839
	resp, err := makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
840
841
842
843
844
	if err != nil {
		return err
	}
	defer resp.Body.Close()

845
	fn(api.ProgressResponse{Status: "success"})
846
847
848
849

	return nil
}

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

853
	// build deleteMap to prune unused layers
Michael Yang's avatar
Michael Yang committed
854
	deleteMap := make(map[string]struct{})
Michael Yang's avatar
Michael Yang committed
855
856
857
	manifest, _, err := GetManifest(mp)
	if errors.Is(err, os.ErrNotExist) {
		// noop
858
859
	} else if err != nil {
		slog.Warn("pulling model with bad existing manifest", "name", name, "error", err)
Michael Yang's avatar
Michael Yang committed
860
861
862
	} else {
		for _, l := range manifest.Layers {
			deleteMap[l.Digest] = struct{}{}
863
		}
Michael Yang's avatar
Michael Yang committed
864
865
		if manifest.Config.Digest != "" {
			deleteMap[manifest.Config.Digest] = struct{}{}
866
867
868
		}
	}

869
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
Michael Yang's avatar
lint  
Michael Yang committed
870
		return errors.New("insecure protocol http")
871
	}
872

873
	fn(api.ProgressResponse{Status: "pulling manifest"})
874

875
	manifest, err = pullModelManifest(ctx, mp, regOpts)
876
	if err != nil {
877
		return fmt.Errorf("pull model manifest: %s", err)
878
879
	}

880
	var layers []Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
881
	layers = append(layers, manifest.Layers...)
882
	if manifest.Config.Digest != "" {
883
		layers = append(layers, manifest.Config)
884
	}
885

886
	skipVerify := make(map[string]bool)
887
	for _, layer := range layers {
888
889
890
891
892
893
894
		cacheHit, err := downloadBlob(ctx, downloadOpts{
			mp:      mp,
			digest:  layer.Digest,
			regOpts: regOpts,
			fn:      fn,
		})
		if err != nil {
895
896
			return err
		}
897
		skipVerify[layer.Digest] = cacheHit
898
		delete(deleteMap, layer.Digest)
899
	}
900
	delete(deleteMap, manifest.Config.Digest)
901

Michael Yang's avatar
Michael Yang committed
902
903
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
904
905
906
		if skipVerify[layer.Digest] {
			continue
		}
Michael Yang's avatar
Michael Yang committed
907
		if err := verifyBlob(layer.Digest); err != nil {
908
909
910
911
912
913
914
915
			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
916
					slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err))
917
918
				}
			}
Michael Yang's avatar
Michael Yang committed
919
920
921
922
			return err
		}
	}

923
	fn(api.ProgressResponse{Status: "writing manifest"})
924

925
	manifestJSON, err := json.Marshal(manifest)
926
927
928
929
	if err != nil {
		return err
	}

930
	fp, err := mp.GetManifestPath()
931
932
933
	if err != nil {
		return err
	}
934
935
936
	if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
		return err
	}
937

Bruce MacDonald's avatar
Bruce MacDonald committed
938
	err = os.WriteFile(fp, manifestJSON, 0o644)
939
	if err != nil {
940
		slog.Info(fmt.Sprintf("couldn't write to %s", fp))
941
942
943
		return err
	}

Michael Yang's avatar
Michael Yang committed
944
945
	if !envconfig.NoPrune() && len(deleteMap) > 0 {
		fn(api.ProgressResponse{Status: "removing unused layers"})
Michael Yang's avatar
Michael Yang committed
946
		if err := deleteUnusedLayers(deleteMap); err != nil {
947
			fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't remove unused layers: %v", err)})
948
949
950
		}
	}

951
	fn(api.ProgressResponse{Status: "success"})
952
953
954
955

	return nil
}

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

Michael Yang's avatar
Michael Yang committed
959
960
	headers := make(http.Header)
	headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
961
	resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
962
963
964
965
966
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

Michael Yang's avatar
Michael Yang committed
967
	var m Manifest
968
969
970
971
	if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
		return nil, err
	}

Michael Yang's avatar
Michael Yang committed
972
	return &m, err
973
974
975
}

// GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
Michael Yang's avatar
Michael Yang committed
976
func GetSHA256Digest(r io.Reader) (string, int64) {
Michael Yang's avatar
Michael Yang committed
977
978
979
980
981
982
	h := sha256.New()
	n, err := io.Copy(h, r)
	if err != nil {
		log.Fatal(err)
	}

Michael Yang's avatar
Michael Yang committed
983
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
984
985
}

Michael Yang's avatar
lint  
Michael Yang committed
986
var errUnauthorized = errors.New("unauthorized: access denied")
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015

// getTokenSubject returns the subject of a JWT token, it does not validate the token
func getTokenSubject(token string) string {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return ""
	}

	payload := parts[1]
	payloadBytes, err := base64.RawURLEncoding.DecodeString(payload)
	if err != nil {
		slog.Error(fmt.Sprintf("failed to decode jwt payload: %v", err))
		return ""
	}

	var payloadMap map[string]interface{}
	if err := json.Unmarshal(payloadBytes, &payloadMap); err != nil {
		slog.Error(fmt.Sprintf("failed to unmarshal payload JSON: %v", err))
		return ""
	}

	sub, ok := payloadMap["sub"]
	if !ok {
		slog.Error("jwt does not contain 'sub' field")
		return ""
	}

	return fmt.Sprintf("%s", sub)
}
1016

Michael Yang's avatar
Michael Yang committed
1017
func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *registryOptions) (*http.Response, error) {
1018
	anonymous := true // access will default to anonymous if no user is found associated with the public key
Michael Yang's avatar
lint  
Michael Yang committed
1019
	for range 2 {
Michael Yang's avatar
Michael Yang committed
1020
		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
Michael Yang's avatar
Michael Yang committed
1021
		if err != nil {
Michael Yang's avatar
Michael Yang committed
1022
			if !errors.Is(err, context.Canceled) {
1023
				slog.Info(fmt.Sprintf("request failed: %v", err))
Michael Yang's avatar
Michael Yang committed
1024
1025
			}

Michael Yang's avatar
Michael Yang committed
1026
1027
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
1028
1029
1030

		switch {
		case resp.StatusCode == http.StatusUnauthorized:
1031
1032
			resp.Body.Close()

Michael Yang's avatar
Michael Yang committed
1033
			// Handle authentication error with one retry
Michael Yang's avatar
Michael Yang committed
1034
1035
			challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
			token, err := getAuthorizationToken(ctx, challenge)
Michael Yang's avatar
Michael Yang committed
1036
1037
1038
			if err != nil {
				return nil, err
			}
1039
			anonymous = getTokenSubject(token) == "anonymous"
Michael Yang's avatar
Michael Yang committed
1040
1041
1042
1043
1044
1045
1046
1047
			regOpts.Token = token
			if body != nil {
				_, err = body.Seek(0, io.SeekStart)
				if err != nil {
					return nil, err
				}
			}
		case resp.StatusCode == http.StatusNotFound:
1048
			resp.Body.Close()
Michael Yang's avatar
Michael Yang committed
1049
1050
			return nil, os.ErrNotExist
		case resp.StatusCode >= http.StatusBadRequest:
1051
			defer resp.Body.Close()
Michael Yang's avatar
Michael Yang committed
1052
1053
1054
1055
1056
1057
1058
			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
1059
1060
1061
		}
	}

1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
	if anonymous {
		// no user is associated with the public key, and the request requires non-anonymous access
		pubKey, nestedErr := auth.GetPublicKey()
		if nestedErr != nil {
			slog.Error(fmt.Sprintf("couldn't get public key: %v", nestedErr))
			return nil, errUnauthorized
		}
		return nil, &errtypes.UnknownOllamaKey{Key: pubKey}
	}
	// user is associated with the public key, but is not authorized to make the request
Michael Yang's avatar
Michael Yang committed
1072
	return nil, errUnauthorized
Michael Yang's avatar
Michael Yang committed
1073
1074
}

1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
// testMakeRequestDialContext specifies the dial function for the http client in
// makeRequest. It can be used to resolve hosts in model names to local
// addresses for testing. For example, the model name ("example.com/my/model")
// can be directed to push/pull from "127.0.0.1:1234".
//
// This is not safe to set across goroutines. It should be set in
// the main test goroutine, and not by tests marked to run in parallel with
// t.Parallel().
//
// It should be cleared after use, otherwise it will affect other tests.
//
// Ideally we would have some set this up the stack, but the code is not
// structured in a way that makes this easy, so this will have to do for now.
var testMakeRequestDialContext func(ctx context.Context, network, addr string) (net.Conn, error)

Michael Yang's avatar
Michael Yang committed
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
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)
		}
	}

1112
	req.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
Michael Yang's avatar
Michael Yang committed
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122

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

		req.ContentLength = contentLength
	}

1123
	resp, err := (&http.Client{
1124
1125
1126
		Transport: &http.Transport{
			DialContext: testMakeRequestDialContext,
		},
1127
1128
		CheckRedirect: regOpts.CheckRedirect,
	}).Do(req)
Michael Yang's avatar
Michael Yang committed
1129
1130
1131
1132
1133
1134
1135
	if err != nil {
		return nil, err
	}

	return resp, nil
}

Patrick Devine's avatar
Patrick Devine committed
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
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
1159
func parseRegistryChallenge(authStr string) registryChallenge {
Patrick Devine's avatar
Patrick Devine committed
1160
1161
	authStr = strings.TrimPrefix(authStr, "Bearer ")

Michael Yang's avatar
Michael Yang committed
1162
	return registryChallenge{
Patrick Devine's avatar
Patrick Devine committed
1163
1164
1165
1166
1167
1168
		Realm:   getValue(authStr, "realm"),
		Service: getValue(authStr, "service"),
		Scope:   getValue(authStr, "scope"),
	}
}

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

Michael Yang's avatar
Michael Yang committed
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
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 {
1185
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
1186
1187
1188
1189
	}

	return nil
}