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

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

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

26
	"github.com/jmorganca/ollama/api"
27
	"github.com/jmorganca/ollama/llm"
28
	"github.com/jmorganca/ollama/parser"
Michael Yang's avatar
Michael Yang committed
29
	"github.com/jmorganca/ollama/version"
30
31
)

32
33
34
35
type RegistryOptions struct {
	Insecure bool
	Username string
	Password string
Patrick Devine's avatar
Patrick Devine committed
36
	Token    string
37
38
}

39
type Model struct {
Michael Yang's avatar
Michael Yang committed
40
	Name           string `json:"name"`
41
	Config         ConfigV2
Michael Yang's avatar
Michael Yang committed
42
43
44
45
46
47
48
49
50
	ShortName      string
	ModelPath      string
	OriginalModel  string
	AdapterPaths   []string
	ProjectorPaths []string
	Template       string
	System         string
	License        []string
	Digest         string
Patrick Devine's avatar
Patrick Devine committed
51
	Size           int64
Michael Yang's avatar
Michael Yang committed
52
	Options        map[string]interface{}
53
54
}

Bruce MacDonald's avatar
Bruce MacDonald committed
55
56
57
58
59
60
type PromptVars struct {
	System   string
	Prompt   string
	Response string
	First    bool
}
61

62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
// extractParts extracts the parts of the template before and after the {{.Response}} node.
func extractParts(tmplStr string) (pre string, post string, err error) {
	tmpl, err := template.New("").Parse(tmplStr)
	if err != nil {
		return "", "", err
	}

	var foundResponse bool

	for _, node := range tmpl.Tree.Root.Nodes {
		if node.Type() == parse.NodeAction && node.String() == "{{.Response}}" {
			foundResponse = true
		}
		if !foundResponse {
			pre += node.String()
		} else {
			post += node.String()
		}
	}

	return pre, post, nil
}

func Prompt(promptTemplate string, p PromptVars) (string, error) {
Bruce MacDonald's avatar
Bruce MacDonald committed
86
	var prompt strings.Builder
87
	// Use the "missingkey=zero" option to handle missing variables without panicking
88
	tmpl, err := template.New("").Option("missingkey=zero").Parse(promptTemplate)
89
90
91
92
	if err != nil {
		return "", err
	}

93
94
95
96
97
	vars := map[string]any{
		"System":   p.System,
		"Prompt":   p.Prompt,
		"Response": p.Response,
		"First":    p.First,
98
99
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
100
	var sb strings.Builder
101
	if err := tmpl.Execute(&sb, vars); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
102
103
104
		return "", err
	}
	prompt.WriteString(sb.String())
105
106
107
108
109
110

	if !strings.Contains(prompt.String(), p.Response) {
		// if the response is not in the prompt template, append it to the end
		prompt.WriteString(p.Response)
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
111
112
	return prompt.String(), nil
}
113

114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
// PreResponsePrompt returns the prompt before the response tag
func (m *Model) PreResponsePrompt(p PromptVars) (string, error) {
	if p.System == "" {
		// use the default system prompt for this model if one is not specified
		p.System = m.System
	}
	pre, _, err := extractParts(m.Template)
	if err != nil {
		return "", err
	}

	return Prompt(pre, p)
}

// PostResponseTemplate returns the template after the response tag
func (m *Model) PostResponseTemplate(p PromptVars) (string, error) {
	if p.System == "" {
		// use the default system prompt for this model if one is not specified
		p.System = m.System
	}
	_, post, err := extractParts(m.Template)
	if err != nil {
		return "", err
	}

	if post == "" {
		// if there is no post-response template, return the provided response
		return p.Response, nil
	}

	return Prompt(post, p)
}

147
func (m *Model) ChatPrompt(msgs []api.Message) (string, []api.ImageData, error) {
Bruce MacDonald's avatar
Bruce MacDonald committed
148
149
	// build the prompt from the list of messages
	var prompt strings.Builder
150
	var currentImages []api.ImageData
Bruce MacDonald's avatar
Bruce MacDonald committed
151
	currentVars := PromptVars{
152
153
		First:  true,
		System: m.System,
Bruce MacDonald's avatar
Bruce MacDonald committed
154
155
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
156
	writePrompt := func() error {
157
		p, err := Prompt(m.Template, currentVars)
Bruce MacDonald's avatar
Bruce MacDonald committed
158
159
160
161
162
163
164
165
166
		if err != nil {
			return err
		}
		prompt.WriteString(p)
		currentVars = PromptVars{}
		return nil
	}

	for _, msg := range msgs {
167
		switch strings.ToLower(msg.Role) {
Bruce MacDonald's avatar
Bruce MacDonald committed
168
		case "system":
169
			if currentVars.System != "" {
Bruce MacDonald's avatar
Bruce MacDonald committed
170
				if err := writePrompt(); err != nil {
171
					return "", nil, err
Bruce MacDonald's avatar
Bruce MacDonald committed
172
173
174
175
				}
			}
			currentVars.System = msg.Content
		case "user":
176
			if currentVars.Prompt != "" {
Bruce MacDonald's avatar
Bruce MacDonald committed
177
				if err := writePrompt(); err != nil {
178
					return "", nil, err
Bruce MacDonald's avatar
Bruce MacDonald committed
179
180
181
				}
			}
			currentVars.Prompt = msg.Content
182
			currentImages = msg.Images
Bruce MacDonald's avatar
Bruce MacDonald committed
183
184
185
		case "assistant":
			currentVars.Response = msg.Content
			if err := writePrompt(); err != nil {
186
				return "", nil, err
Bruce MacDonald's avatar
Bruce MacDonald committed
187
188
			}
		default:
189
			return "", nil, fmt.Errorf("invalid role: %s, role must be one of [system, user, assistant]", msg.Role)
Bruce MacDonald's avatar
Bruce MacDonald committed
190
		}
191
192
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
193
194
	// Append the last set of vars if they are non-empty
	if currentVars.Prompt != "" || currentVars.System != "" {
195
196
197
		p, err := m.PreResponsePrompt(currentVars)
		if err != nil {
			return "", nil, fmt.Errorf("pre-response template: %w", err)
Bruce MacDonald's avatar
Bruce MacDonald committed
198
		}
199
		prompt.WriteString(p)
Bruce MacDonald's avatar
Bruce MacDonald committed
200
201
	}

202
	return prompt.String(), currentImages, nil
203
204
}

205
206
207
type ManifestV2 struct {
	SchemaVersion int      `json:"schemaVersion"`
	MediaType     string   `json:"mediaType"`
Michael Yang's avatar
Michael Yang committed
208
	Config        *Layer   `json:"config"`
209
210
211
212
	Layers        []*Layer `json:"layers"`
}

type ConfigV2 struct {
213
214
215
216
217
218
	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"`

219
	// required by spec
220
221
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
222
	RootFS       RootFS `json:"rootfs"`
223
224
}

Michael Yang's avatar
Michael Yang committed
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
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
	}
}

255
256
257
258
259
type RootFS struct {
	Type    string   `json:"type"`
	DiffIDs []string `json:"diff_ids"`
}

Michael Yang's avatar
Michael Yang committed
260
func (m *ManifestV2) GetTotalSize() (total int64) {
Patrick Devine's avatar
Patrick Devine committed
261
262
263
	for _, layer := range m.Layers {
		total += layer.Size
	}
Michael Yang's avatar
Michael Yang committed
264

Patrick Devine's avatar
Patrick Devine committed
265
266
267
268
	total += m.Config.Size
	return total
}

Patrick Devine's avatar
Patrick Devine committed
269
func GetManifest(mp ModelPath) (*ManifestV2, string, error) {
270
	fp, err := mp.GetManifestPath()
271
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
272
		return nil, "", err
273
	}
274

275
	if _, err = os.Stat(fp); err != nil {
Patrick Devine's avatar
Patrick Devine committed
276
		return nil, "", err
277
278
279
280
	}

	var manifest *ManifestV2

281
	bts, err := os.ReadFile(fp)
282
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
283
		return nil, "", fmt.Errorf("couldn't open file '%s'", fp)
284
285
	}

Patrick Devine's avatar
Patrick Devine committed
286
287
288
	shaSum := sha256.Sum256(bts)
	shaStr := hex.EncodeToString(shaSum[:])

289
	if err := json.Unmarshal(bts, &manifest); err != nil {
Patrick Devine's avatar
Patrick Devine committed
290
		return nil, "", err
291
292
	}

Patrick Devine's avatar
Patrick Devine committed
293
	return manifest, shaStr, nil
294
295
296
}

func GetModel(name string) (*Model, error) {
297
	mp := ParseModelPath(name)
Patrick Devine's avatar
Patrick Devine committed
298
	manifest, digest, err := GetManifest(mp)
299
300
301
302
303
	if err != nil {
		return nil, err
	}

	model := &Model{
304
305
306
307
308
		Name:      mp.GetFullTagname(),
		ShortName: mp.GetShortTagname(),
		Digest:    digest,
		Template:  "{{ .Prompt }}",
		License:   []string{},
Patrick Devine's avatar
Patrick Devine committed
309
		Size:      manifest.GetTotalSize(),
310
311
	}

312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
	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
	}

327
	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
328
		filename, err := GetBlobsPath(layer.Digest)
329
330
331
332
		if err != nil {
			return nil, err
		}

333
334
335
		switch layer.MediaType {
		case "application/vnd.ollama.image.model":
			model.ModelPath = filename
Patrick Devine's avatar
Patrick Devine committed
336
			model.OriginalModel = layer.From
337
		case "application/vnd.ollama.image.embed":
338
339
			// Deprecated in versions  > 0.1.2
			// TODO: remove this warning in a future version
340
			slog.Info("WARNING: model contains embeddings, but embeddings in modelfiles have been deprecated and will be ignored.")
341
342
		case "application/vnd.ollama.image.adapter":
			model.AdapterPaths = append(model.AdapterPaths, filename)
Michael Yang's avatar
Michael Yang committed
343
344
		case "application/vnd.ollama.image.projector":
			model.ProjectorPaths = append(model.ProjectorPaths, filename)
345
346
347
348
349
350
351
352
353
		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)
354
355
356
			if err != nil {
				return nil, err
			}
357
358

			model.System = string(bts)
359
360
361
362
363
364
365
		case "application/vnd.ollama.image.prompt":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

			model.Template = string(bts)
366
		case "application/vnd.ollama.image.params":
Michael Yang's avatar
Michael Yang committed
367
368
369
370
371
			params, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer params.Close()
372

373
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
374
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
375
376
				return nil, err
			}
Patrick Devine's avatar
Patrick Devine committed
377
378
379
380
381
382
		case "application/vnd.ollama.image.license":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}
			model.License = append(model.License, string(bts))
383
384
385
386
387
388
		}
	}

	return model, nil
}

389
390
func realpath(mfDir, from string) string {
	abspath, err := filepath.Abs(from)
Michael Yang's avatar
Michael Yang committed
391
	if err != nil {
392
		return from
393
394
	}

Michael Yang's avatar
Michael Yang committed
395
	home, err := os.UserHomeDir()
396
	if err != nil {
Michael Yang's avatar
Michael Yang committed
397
		return abspath
398
399
	}

400
	if from == "~" {
Michael Yang's avatar
Michael Yang committed
401
		return home
402
403
404
405
406
407
408
	} 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)
409
410
	}

Michael Yang's avatar
Michael Yang committed
411
412
413
	return abspath
}

414
func CreateModel(ctx context.Context, name, modelFileDir string, commands []parser.Command, fn func(resp api.ProgressResponse)) error {
415
416
	config := ConfigV2{
		OS:           "linux",
Michael Yang's avatar
Michael Yang committed
417
		Architecture: "amd64",
Michael Yang's avatar
Michael Yang committed
418
419
420
		RootFS: RootFS{
			Type: "layers",
		},
421
422
	}

Michael Yang's avatar
Michael Yang committed
423
424
	deleteMap := make(map[string]struct{})

Michael Yang's avatar
Michael Yang committed
425
	var layers Layers
Michael Yang's avatar
Michael Yang committed
426

427
	params := make(map[string][]string)
Michael Yang's avatar
Michael Yang committed
428
429
	fromParams := make(map[string]any)

430
	for _, c := range commands {
431
		slog.Info(fmt.Sprintf("[%s] - %s", c.Name, c.Args))
Michael Yang's avatar
Michael Yang committed
432
433
		mediatype := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)

434
435
		switch c.Name {
		case "model":
Michael Yang's avatar
Michael Yang committed
436
437
438
439
440
441
442
443
444
			if strings.HasPrefix(c.Args, "@") {
				blobPath, err := GetBlobsPath(strings.TrimPrefix(c.Args, "@"))
				if err != nil {
					return err
				}

				c.Args = blobPath
			}

445
			bin, err := os.Open(realpath(modelFileDir, c.Args))
446
			if err != nil {
Michael Yang's avatar
Michael Yang committed
447
448
449
450
451
452
453
				// 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"})
					if err := PullModel(ctx, c.Args, &RegistryOptions{}, fn); err != nil {
454
455
456
						return err
					}

Michael Yang's avatar
Michael Yang committed
457
					manifest, _, err = GetManifest(modelpath)
458
459
460
					if err != nil {
						return err
					}
Michael Yang's avatar
Michael Yang committed
461
462
				case err != nil:
					return err
463
				}
464

465
				fn(api.ProgressResponse{Status: "reading model metadata"})
Michael Yang's avatar
Michael Yang committed
466
				fromConfigPath, err := GetBlobsPath(manifest.Config.Digest)
Michael Yang's avatar
Michael Yang committed
467
468
469
470
				if err != nil {
					return err
				}

Michael Yang's avatar
Michael Yang committed
471
				fromConfigFile, err := os.Open(fromConfigPath)
Michael Yang's avatar
Michael Yang committed
472
473
474
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
475
				defer fromConfigFile.Close()
Michael Yang's avatar
Michael Yang committed
476

Michael Yang's avatar
Michael Yang committed
477
478
				var fromConfig ConfigV2
				if err := json.NewDecoder(fromConfigFile).Decode(&fromConfig); err != nil {
Michael Yang's avatar
Michael Yang committed
479
480
481
					return err
				}

Bruce MacDonald's avatar
Bruce MacDonald committed
482
483
484
485
486
				// 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
487
488
489
490
				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
491

Michael Yang's avatar
Michael Yang committed
492
493
494
495
				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
496
497
498
499
						if err != nil {
							return err
						}

Michael Yang's avatar
Michael Yang committed
500
						fromParamsFile, err := os.Open(fromParamsPath)
Michael Yang's avatar
Michael Yang committed
501
502
503
						if err != nil {
							return err
						}
Michael Yang's avatar
Michael Yang committed
504
						defer fromParamsFile.Close()
Michael Yang's avatar
Michael Yang committed
505

Michael Yang's avatar
Michael Yang committed
506
						if err := json.NewDecoder(fromParamsFile).Decode(&fromParams); err != nil {
Michael Yang's avatar
Michael Yang committed
507
508
509
510
							return err
						}
					}

Michael Yang's avatar
Michael Yang committed
511
					layer, err := NewLayerFromLayer(layer.Digest, layer.MediaType, modelpath.GetShortTagname())
512
513
514
					if err != nil {
						return err
					}
Michael Yang's avatar
Michael Yang committed
515

Michael Yang's avatar
Michael Yang committed
516
					layers.Add(layer)
517
				}
Michael Yang's avatar
Michael Yang committed
518
519
520

				deleteMap[manifest.Config.Digest] = struct{}{}
				continue
521
			}
Michael Yang's avatar
Michael Yang committed
522
			defer bin.Close()
523

524
			var offset int64
Bruce MacDonald's avatar
Bruce MacDonald committed
525
		CREATE:
526
527
			for {
				fn(api.ProgressResponse{Status: "creating model layer"})
528

529
530
				bin.Seek(offset, io.SeekStart)
				ggml, err := llm.DecodeGGML(bin)
Bruce MacDonald's avatar
Bruce MacDonald committed
531
532
533
534
535
536
537
538
539
				if err != nil {
					switch {
					case errors.Is(err, io.EOF):
						break CREATE
					case errors.Is(err, llm.ErrUnsupportedFormat):
						return fmt.Errorf("model binary specified in FROM field is not a valid gguf format model, %w", err)
					default:
						return err
					}
540
				}
Michael Yang's avatar
Michael Yang committed
541

Michael Yang's avatar
Michael Yang committed
542
543
544
545
				config.SetModelFormat(ggml.Name())
				config.SetModelFamily(ggml.ModelFamily())
				config.SetModelType(ggml.ModelType())
				config.SetFileType(ggml.FileType())
546

547
548
549
550
				mediatype := mediatype
				if ggml.ModelFamily() == "clip" {
					mediatype = "application/vnd.ollama.image.projector"
				}
551

552
553
554
555
556
557
558
559
560
561
				sr := io.NewSectionReader(bin, offset, ggml.Size)
				layer, err := NewLayer(sr, mediatype)
				if err != nil {
					return err
				}

				layers.Add(layer)

				offset += ggml.Size
			}
Michael Yang's avatar
Michael Yang committed
562
		case "adapter":
563
564
565
566
567
568
569
570
			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
571

Michael Yang's avatar
Michael Yang committed
572
			fn(api.ProgressResponse{Status: "creating adapter layer"})
573
			bin, err := os.Open(realpath(modelFileDir, c.Args))
574
			if err != nil {
Michael Yang's avatar
Michael Yang committed
575
				return err
576
			}
Michael Yang's avatar
Michael Yang committed
577
			defer bin.Close()
578

Michael Yang's avatar
Michael Yang committed
579
			layer, err := NewLayer(bin, mediatype)
580
			if err != nil {
Michael Yang's avatar
Michael Yang committed
581
				return err
582
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
583

Michael Yang's avatar
Michael Yang committed
584
			layers.Add(layer)
Michael Yang's avatar
Michael Yang committed
585
586
		case "license":
			fn(api.ProgressResponse{Status: "creating license layer"})
Michael Yang's avatar
Michael Yang committed
587
588
589

			bin := strings.NewReader(c.Args)
			layer, err := NewLayer(bin, mediatype)
Bruce MacDonald's avatar
Bruce MacDonald committed
590
591
592
593
			if err != nil {
				return err
			}

Michael Yang's avatar
Michael Yang committed
594
			layers.Add(layer)
Michael Yang's avatar
Michael Yang committed
595
596
597
		case "template", "system":
			fn(api.ProgressResponse{Status: fmt.Sprintf("creating %s layer", c.Name)})

Michael Yang's avatar
Michael Yang committed
598
599
			bin := strings.NewReader(c.Args)
			layer, err := NewLayer(bin, mediatype)
600
			if err != nil {
601
				return err
602
			}
603

Michael Yang's avatar
Michael Yang committed
604
			layers.Replace(layer)
605
		default:
606
			params[c.Name] = append(params[c.Name], c.Args)
607
608
609
		}
	}

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

613
		formattedParams, err := api.FormatParams(params)
614
		if err != nil {
Michael Yang's avatar
Michael Yang committed
615
			return err
616
		}
617

Michael Yang's avatar
Michael Yang committed
618
		for k, v := range fromParams {
Michael Yang's avatar
Michael Yang committed
619
620
621
622
623
			if _, ok := formattedParams[k]; !ok {
				formattedParams[k] = v
			}
		}

Patrick Devine's avatar
Patrick Devine committed
624
		// xxx - can this be removed?
Michael Yang's avatar
Michael Yang committed
625
		if config.ModelType == "65B" {
Michael Yang's avatar
Michael Yang committed
626
			if gqa, ok := formattedParams["gqa"].(int); ok && gqa == 8 {
Michael Yang's avatar
Michael Yang committed
627
628
629
630
				config.ModelType = "70B"
			}
		}

Michael Yang's avatar
Michael Yang committed
631
632
		var b bytes.Buffer
		if err := json.NewEncoder(&b).Encode(formattedParams); err != nil {
633
634
635
			return err
		}

Michael Yang's avatar
Michael Yang committed
636
		fn(api.ProgressResponse{Status: "creating config layer"})
Michael Yang's avatar
Michael Yang committed
637
		layer, err := NewLayer(&b, "application/vnd.ollama.image.params")
638
		if err != nil {
Michael Yang's avatar
Michael Yang committed
639
			return err
640
		}
Michael Yang's avatar
Michael Yang committed
641

Michael Yang's avatar
Michael Yang committed
642
		layers.Replace(layer)
643
644
	}

Michael Yang's avatar
Michael Yang committed
645
646
647
	digests := make([]string, len(layers.items))
	for i, layer := range layers.items {
		digests[i] = layer.Digest
648
649
	}

Michael Yang's avatar
Michael Yang committed
650
	config.RootFS.DiffIDs = digests
Michael Yang's avatar
Michael Yang committed
651

Michael Yang's avatar
Michael Yang committed
652
653
	var b bytes.Buffer
	if err := json.NewEncoder(&b).Encode(config); err != nil {
654
655
656
		return err
	}

Michael Yang's avatar
Michael Yang committed
657
658
	configLayer, err := NewLayer(&b, "application/vnd.docker.container.image.v1+json")
	if err != nil {
659
660
661
		return err
	}

Michael Yang's avatar
Michael Yang committed
662
	delete(deleteMap, configLayer.Digest)
663

Michael Yang's avatar
Michael Yang committed
664
665
	for _, layer := range append(layers.items, configLayer) {
		committed, err := layer.Commit()
666
667
668
		if err != nil {
			return err
		}
669

Michael Yang's avatar
Michael Yang committed
670
671
672
		status := "writing layer"
		if !committed {
			status = "using already created layer"
673
674
		}

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

Michael Yang's avatar
Michael Yang committed
677
		delete(deleteMap, layer.Digest)
678
679
	}

Michael Yang's avatar
Michael Yang committed
680
681
	fn(api.ProgressResponse{Status: "writing manifest"})
	if err := WriteManifest(name, configLayer, layers.items); err != nil {
682
683
		return err
	}
684

Michael Yang's avatar
Michael Yang committed
685
686
687
	if noprune := os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
		if err := deleteUnusedLayers(nil, deleteMap, false); err != nil {
			return err
688
689
690
		}
	}

Michael Yang's avatar
Michael Yang committed
691
692
	fn(api.ProgressResponse{Status: "success"})
	return nil
693
694
}

Patrick Devine's avatar
Patrick Devine committed
695
func CopyModel(src, dest string) error {
696
	srcModelPath := ParseModelPath(src)
697
	srcPath, err := srcModelPath.GetManifestPath()
698
699
700
701
	if err != nil {
		return err
	}

702
	destModelPath := ParseModelPath(dest)
703
	destPath, err := destModelPath.GetManifestPath()
Patrick Devine's avatar
Patrick Devine committed
704
705
706
	if err != nil {
		return err
	}
707
708
709
	if err := os.MkdirAll(filepath.Dir(destPath), 0o755); err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
710
711

	// copy the file
Michael Yang's avatar
Michael Yang committed
712
	input, err := os.ReadFile(srcPath)
Patrick Devine's avatar
Patrick Devine committed
713
714
715
716
717
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

Michael Yang's avatar
Michael Yang committed
718
	err = os.WriteFile(destPath, input, 0o644)
Patrick Devine's avatar
Patrick Devine committed
719
720
721
722
723
724
725
726
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

	return nil
}

Michael Yang's avatar
Michael Yang committed
727
func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{}, dryRun bool) error {
728
729
730
731
	fp, err := GetManifestPath()
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
732
733
734
735

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

Michael Yang's avatar
Michael Yang committed
738
739
740
741
		dir, file := filepath.Split(path)
		dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
		tag := strings.Join([]string{dir, file}, ":")
		fmp := ParseModelPath(tag)
742

Michael Yang's avatar
Michael Yang committed
743
		// skip the manifest we're trying to delete
744
		if skipModelPath != nil && skipModelPath.GetFullTagname() == fmp.GetFullTagname() {
Michael Yang's avatar
Michael Yang committed
745
			return nil
746
		}
Michael Yang's avatar
Michael Yang committed
747
748
749
750

		// 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
751
			// nolint: nilerr
Michael Yang's avatar
Michael Yang committed
752
753
754
755
756
757
758
759
			return nil
		}

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

		delete(deleteMap, manifest.Config.Digest)
760
		return nil
Michael Yang's avatar
Michael Yang committed
761
762
763
	}

	if err := filepath.Walk(fp, walkFunc); err != nil {
Michael Yang's avatar
Michael Yang committed
764
765
		return err
	}
766
767

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

784
785
786
787
	return nil
}

func PruneLayers() error {
Michael Yang's avatar
Michael Yang committed
788
	deleteMap := make(map[string]struct{})
789
790
791
792
793
794
795
	p, err := GetBlobsPath("")
	if err != nil {
		return err
	}

	blobs, err := os.ReadDir(p)
	if err != nil {
796
		slog.Info(fmt.Sprintf("couldn't read dir '%s': %v", p, err))
797
798
799
800
801
802
803
804
		return err
	}

	for _, blob := range blobs {
		name := blob.Name()
		if runtime.GOOS == "windows" {
			name = strings.ReplaceAll(name, "-", ":")
		}
Michael Yang's avatar
Michael Yang committed
805
806
807
		if strings.HasPrefix(name, "sha256:") {
			deleteMap[name] = struct{}{}
		}
808
809
	}

810
	slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap)))
811
812
813
814
815
816

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

817
	slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
818
819
820
821

	return nil
}

Michael Yang's avatar
Michael Yang committed
822
823
824
825
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
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
}

855
856
857
858
859
860
861
func DeleteModel(name string) error {
	mp := ParseModelPath(name)
	manifest, _, err := GetManifest(mp)
	if err != nil {
		return err
	}

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

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

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

	return nil
}

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

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

Michael Yang's avatar
Michael Yang committed
900
		mt.Parameters[k] = []any{v}
Patrick Devine's avatar
Patrick Devine committed
901
902
	}

Michael Yang's avatar
Michael Yang committed
903
904
	mt.Model = model
	mt.From = model.ModelPath
Patrick Devine's avatar
Patrick Devine committed
905

Michael Yang's avatar
Michael Yang committed
906
	if model.OriginalModel != "" {
Daniel Reis's avatar
Daniel Reis committed
907
		mt.From = model.OriginalModel
Patrick Devine's avatar
Patrick Devine committed
908
909
910
911
912
913
914
915
	}

	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 }}"""
916
917

{{- if .System }}
Patrick Devine's avatar
Patrick Devine committed
918
SYSTEM """{{ .System }}"""
919
{{- end }}
920
921
922
923

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

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

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

	var buf bytes.Buffer

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

	return buf.String(), nil
}

947
func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
948
	mp := ParseModelPath(name)
949
950
	fn(api.ProgressResponse{Status: "retrieving manifest"})

951
952
953
954
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
	}

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

	var layers []*Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
962
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
963
	layers = append(layers, manifest.Config)
964
965

	for _, layer := range layers {
Michael Yang's avatar
Michael Yang committed
966
		if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
967
			slog.Info(fmt.Sprintf("error uploading blob: %v", err))
968
969
970
			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())
			}
971
972
			return err
		}
973
974
	}

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

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

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

992
	fn(api.ProgressResponse{Status: "success"})
993
994
995
996

	return nil
}

997
func PullModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
998
999
	mp := ParseModelPath(name)

1000
1001
1002
1003
1004
	var manifest *ManifestV2
	var err error
	var noprune string

	// build deleteMap to prune unused layers
Michael Yang's avatar
Michael Yang committed
1005
	deleteMap := make(map[string]struct{})
1006
1007
1008
1009
1010
1011
1012
1013
1014

	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
1015
				deleteMap[l.Digest] = struct{}{}
1016
			}
Michael Yang's avatar
Michael Yang committed
1017
			deleteMap[manifest.Config.Digest] = struct{}{}
1018
1019
1020
		}
	}

1021
1022
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
1023
	}
1024

1025
	fn(api.ProgressResponse{Status: "pulling manifest"})
1026

1027
	manifest, err = pullModelManifest(ctx, mp, regOpts)
1028
	if err != nil {
1029
		return fmt.Errorf("pull model manifest: %s", err)
1030
1031
1032
	}

	var layers []*Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
1033
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
1034
	layers = append(layers, manifest.Config)
1035
1036

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

Michael Yang's avatar
Michael Yang committed
1051
1052
1053
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
		if err := verifyBlob(layer.Digest); err != nil {
1054
1055
1056
1057
1058
1059
1060
1061
			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
1062
					slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err))
1063
1064
				}
			}
Michael Yang's avatar
Michael Yang committed
1065
1066
1067
1068
			return err
		}
	}

1069
	fn(api.ProgressResponse{Status: "writing manifest"})
1070

1071
	manifestJSON, err := json.Marshal(manifest)
1072
1073
1074
1075
	if err != nil {
		return err
	}

1076
	fp, err := mp.GetManifestPath()
1077
1078
1079
	if err != nil {
		return err
	}
1080
1081
1082
	if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
		return err
	}
1083

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

1090
1091
1092
1093
1094
1095
1096
1097
	if noprune == "" {
		fn(api.ProgressResponse{Status: "removing any unused layers"})
		err = deleteUnusedLayers(nil, deleteMap, false)
		if err != nil {
			return err
		}
	}

1098
	fn(api.ProgressResponse{Status: "success"})
1099
1100
1101
1102

	return nil
}

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

Michael Yang's avatar
Michael Yang committed
1106
1107
	headers := make(http.Header)
	headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
1108
	resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
	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
1123
func GetSHA256Digest(r io.Reader) (string, int64) {
Michael Yang's avatar
Michael Yang committed
1124
1125
1126
1127
1128
1129
	h := sha256.New()
	n, err := io.Copy(h, r)
	if err != nil {
		log.Fatal(err)
	}

Michael Yang's avatar
Michael Yang committed
1130
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
1131
1132
}

1133
1134
var errUnauthorized = fmt.Errorf("unauthorized")

Michael Yang's avatar
Michael Yang committed
1135
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
1136
1137
	for i := 0; i < 2; i++ {
		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
Michael Yang's avatar
Michael Yang committed
1138
		if err != nil {
Michael Yang's avatar
Michael Yang committed
1139
			if !errors.Is(err, context.Canceled) {
1140
				slog.Info(fmt.Sprintf("request failed: %v", err))
Michael Yang's avatar
Michael Yang committed
1141
1142
			}

Michael Yang's avatar
Michael Yang committed
1143
1144
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
1145
1146
1147
1148
1149
1150
1151

		switch {
		case resp.StatusCode == http.StatusUnauthorized:
			// Handle authentication error with one retry
			auth := resp.Header.Get("www-authenticate")
			authRedir := ParseAuthRedirectString(auth)
			token, err := getAuthToken(ctx, authRedir)
Michael Yang's avatar
Michael Yang committed
1152
1153
1154
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
			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
1172
1173
1174
		}
	}

Michael Yang's avatar
Michael Yang committed
1175
	return nil, errUnauthorized
Michael Yang's avatar
Michael Yang committed
1176
1177
}

Michael Yang's avatar
Michael Yang committed
1178
func makeRequest(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) {
Michael Yang's avatar
Michael Yang committed
1179
	if requestURL.Scheme != "http" && regOpts != nil && regOpts.Insecure {
Michael Yang's avatar
Michael Yang committed
1180
		requestURL.Scheme = "http"
1181
1182
	}

Michael Yang's avatar
Michael Yang committed
1183
	req, err := http.NewRequestWithContext(ctx, method, requestURL.String(), body)
1184
1185
1186
1187
	if err != nil {
		return nil, err
	}

Michael Yang's avatar
Michael Yang committed
1188
1189
1190
1191
	if headers != nil {
		req.Header = headers
	}

Michael Yang's avatar
Michael Yang committed
1192
1193
1194
1195
1196
1197
	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)
		}
1198
1199
	}

Michael Yang's avatar
Michael Yang committed
1200
	req.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
1201

Michael Yang's avatar
Michael Yang committed
1202
1203
1204
1205
1206
1207
1208
1209
1210
	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
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
	proxyURL, err := http.ProxyFromEnvironment(req)
	if err != nil {
		return nil, err
	}

	client := http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyURL(proxyURL),
		},
	}

	resp, err := client.Do(req)
1223
1224
1225
1226
1227
1228
	if err != nil {
		return nil, err
	}

	return resp, nil
}
Michael Yang's avatar
Michael Yang committed
1229

Patrick Devine's avatar
Patrick Devine committed
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
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]
}

func ParseAuthRedirectString(authStr string) AuthRedirect {
	authStr = strings.TrimPrefix(authStr, "Bearer ")

	return AuthRedirect{
		Realm:   getValue(authStr, "realm"),
		Service: getValue(authStr, "service"),
		Scope:   getValue(authStr, "scope"),
	}
}

1263
1264
var errDigestMismatch = fmt.Errorf("digest mismatch, file must be downloaded again")

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

	return nil
}