updater.go 7.58 KB
Newer Older
1
2
3
//go:build windows || darwin

package updater
4
5
6

import (
	"context"
Michael Yang's avatar
Michael Yang committed
7
	"crypto/rand"
8
9
10
11
12
13
14
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"mime"
	"net/http"
Michael Yang's avatar
Michael Yang committed
15
	"net/url"
16
17
18
19
	"os"
	"path"
	"path/filepath"
	"runtime"
Michael Yang's avatar
lint  
Michael Yang committed
20
	"strconv"
21
22
23
	"strings"
	"time"

24
25
	"github.com/ollama/ollama/app/store"
	"github.com/ollama/ollama/app/version"
26
	"github.com/ollama/ollama/auth"
27
28
29
)

var (
30
31
32
33
34
35
36
37
38
39
40
41
	UpdateCheckURLBase      = "https://ollama.com/api/update"
	UpdateDownloaded        = false
	UpdateCheckInterval     = 60 * 60 * time.Second
	UpdateCheckInitialDelay = 3 * time.Second // 30 * time.Second

	UpdateStageDir    string
	UpgradeLogFile    string
	UpgradeMarkerFile string
	Installer         string
	UserAgentOS       string

	VerifyDownload func() error
42
43
44
45
46
47
48
49
)

// TODO - maybe move up to the API package?
type UpdateResponse struct {
	UpdateURL     string `json:"url"`
	UpdateVersion string `json:"version"`
}

50
func (u *Updater) checkForUpdate(ctx context.Context) (bool, UpdateResponse) {
51
	var updateResp UpdateResponse
Michael Yang's avatar
Michael Yang committed
52
53

	requestURL, err := url.Parse(UpdateCheckURLBase)
54
	if err != nil {
Michael Yang's avatar
Michael Yang committed
55
56
57
58
59
60
61
		return false, updateResp
	}

	query := requestURL.Query()
	query.Add("os", runtime.GOOS)
	query.Add("arch", runtime.GOARCH)
	query.Add("version", version.Version)
Michael Yang's avatar
lint  
Michael Yang committed
62
	query.Add("ts", strconv.FormatInt(time.Now().Unix(), 10))
Michael Yang's avatar
Michael Yang committed
63

64
65
66
67
68
69
	// The original macOS app used to use the device ID
	// to check for updates so include it if present
	if runtime.GOOS == "darwin" {
		if id, err := u.Store.ID(); err == nil && id != "" {
			query.Add("id", id)
		}
70
	}
Michael Yang's avatar
Michael Yang committed
71

72
	var signature string
Michael Yang's avatar
Michael Yang committed
73

74
	nonce, err := auth.NewNonce(rand.Reader, 16)
Michael Yang's avatar
Michael Yang committed
75
	if err != nil {
76
77
78
79
80
81
82
83
84
85
86
		// Don't sign if we haven't yet generated a key pair for the server
		slog.Debug("unable to generate nonce for update check request", "error", err)
	} else {
		query.Add("nonce", nonce)
		requestURL.RawQuery = query.Encode()

		data := []byte(fmt.Sprintf("%s,%s", http.MethodGet, requestURL.RequestURI()))
		signature, err = auth.Sign(ctx, data)
		if err != nil {
			slog.Debug("unable to generate signature for update check request", "error", err)
		}
Michael Yang's avatar
Michael Yang committed
87
88
	}

jmorganca's avatar
jmorganca committed
89
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL.String(), nil)
90
91
92
93
	if err != nil {
		slog.Warn(fmt.Sprintf("failed to check for update: %s", err))
		return false, updateResp
	}
94
95
96
97
98
	if signature != "" {
		req.Header.Set("Authorization", signature)
	}
	ua := fmt.Sprintf("ollama/%s %s Go/%s %s", version.Version, runtime.GOARCH, runtime.Version(), UserAgentOS)
	req.Header.Set("User-Agent", ua)
99

100
	slog.Debug("checking for available update", "requestURL", requestURL, "User-Agent", ua)
Michael Yang's avatar
Michael Yang committed
101
	resp, err := http.DefaultClient.Do(req)
102
103
104
105
106
107
	if err != nil {
		slog.Warn(fmt.Sprintf("failed to check for update: %s", err))
		return false, updateResp
	}
	defer resp.Body.Close()

Michael Yang's avatar
lint  
Michael Yang committed
108
	if resp.StatusCode == http.StatusNoContent {
109
110
111
112
113
114
115
		slog.Debug("check update response 204 (current version is up to date)")
		return false, updateResp
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		slog.Warn(fmt.Sprintf("failed to read body response: %s", err))
	}
116

Michael Yang's avatar
lint  
Michael Yang committed
117
	if resp.StatusCode != http.StatusOK {
118
119
120
		slog.Info(fmt.Sprintf("check update error %d - %.96s", resp.StatusCode, string(body)))
		return false, updateResp
	}
121
122
123
124
125
126
127
128
129
130
131
132
	err = json.Unmarshal(body, &updateResp)
	if err != nil {
		slog.Warn(fmt.Sprintf("malformed response checking for update: %s", err))
		return false, updateResp
	}
	// Extract the version string from the URL in the github release artifact path
	updateResp.UpdateVersion = path.Base(path.Dir(updateResp.UpdateURL))

	slog.Info("New update available at " + updateResp.UpdateURL)
	return true, updateResp
}

133
func (u *Updater) DownloadNewRelease(ctx context.Context, updateResp UpdateResponse) error {
134
135
136
137
138
	// Do a head first to check etag info
	req, err := http.NewRequestWithContext(ctx, http.MethodHead, updateResp.UpdateURL, nil)
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
139

140
141
142
143
144
145
146
147
148
149
150
151
152
153
	// In case of slow downloads, continue the update check in the background
	bgctx, cancel := context.WithCancel(ctx)
	defer cancel()
	go func() {
		for {
			select {
			case <-bgctx.Done():
				return
			case <-time.After(UpdateCheckInterval):
				u.checkForUpdate(bgctx)
			}
		}
	}()

Michael Yang's avatar
Michael Yang committed
154
	resp, err := http.DefaultClient.Do(req)
155
156
157
	if err != nil {
		return fmt.Errorf("error checking update: %w", err)
	}
Michael Yang's avatar
lint  
Michael Yang committed
158
	if resp.StatusCode != http.StatusOK {
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
		return fmt.Errorf("unexpected status attempting to download update %d", resp.StatusCode)
	}
	resp.Body.Close()
	etag := strings.Trim(resp.Header.Get("etag"), "\"")
	if etag == "" {
		slog.Debug("no etag detected, falling back to filename based dedup")
		etag = "_"
	}
	filename := Installer
	_, params, err := mime.ParseMediaType(resp.Header.Get("content-disposition"))
	if err == nil {
		filename = params["filename"]
	}

	stageFilename := filepath.Join(UpdateStageDir, etag, filename)

	// Check to see if we already have it downloaded
	_, err = os.Stat(stageFilename)
	if err == nil {
178
		slog.Info("update already downloaded", "bundle", stageFilename)
179
180
181
		return nil
	}

182
	cleanupOldDownloads(UpdateStageDir)
183
184

	req.Method = http.MethodGet
Michael Yang's avatar
Michael Yang committed
185
	resp, err = http.DefaultClient.Do(req)
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
	if err != nil {
		return fmt.Errorf("error checking update: %w", err)
	}
	defer resp.Body.Close()
	etag = strings.Trim(resp.Header.Get("etag"), "\"")
	if etag == "" {
		slog.Debug("no etag detected, falling back to filename based dedup") // TODO probably can get rid of this redundant log
		etag = "_"
	}

	stageFilename = filepath.Join(UpdateStageDir, etag, filename)

	_, err = os.Stat(filepath.Dir(stageFilename))
	if errors.Is(err, os.ErrNotExist) {
		if err := os.MkdirAll(filepath.Dir(stageFilename), 0o755); err != nil {
			return fmt.Errorf("create ollama dir %s: %v", filepath.Dir(stageFilename), err)
		}
	}

	payload, err := io.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("failed to read body response: %w", err)
	}
	fp, err := os.OpenFile(stageFilename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o755)
	if err != nil {
		return fmt.Errorf("write payload %s: %w", stageFilename, err)
	}
	defer fp.Close()
	if n, err := fp.Write(payload); err != nil || n != len(payload) {
		return fmt.Errorf("write payload %s: %d vs %d -- %w", stageFilename, n, len(payload), err)
	}
	slog.Info("new update downloaded " + stageFilename)

219
220
221
222
	if err := VerifyDownload(); err != nil {
		_ = os.Remove(stageFilename)
		return fmt.Errorf("%s - %s", resp.Request.URL.String(), err)
	}
223
224
225
226
	UpdateDownloaded = true
	return nil
}

227
228
func cleanupOldDownloads(stageDir string) {
	files, err := os.ReadDir(stageDir)
229
230
231
232
233
234
235
236
	if err != nil && errors.Is(err, os.ErrNotExist) {
		// Expected behavior on first run
		return
	} else if err != nil {
		slog.Warn(fmt.Sprintf("failed to list stage dir: %s", err))
		return
	}
	for _, file := range files {
237
		fullname := filepath.Join(stageDir, file.Name())
238
239
240
241
242
243
244
245
		slog.Debug("cleaning up old download: " + fullname)
		err = os.RemoveAll(fullname)
		if err != nil {
			slog.Warn(fmt.Sprintf("failed to cleanup stale update download %s", err))
		}
	}
}

246
247
248
249
250
type Updater struct {
	Store *store.Store
}

func (u *Updater) StartBackgroundUpdaterChecker(ctx context.Context, cb func(string) error) {
251
252
	go func() {
		// Don't blast an update message immediately after startup
253
254
		time.Sleep(UpdateCheckInitialDelay)
		slog.Info("beginning update checker", "interval", UpdateCheckInterval)
255
		for {
256
			available, resp := u.checkForUpdate(ctx)
257
			if available {
258
				err := u.DownloadNewRelease(ctx, resp)
259
260
				if err != nil {
					slog.Error(fmt.Sprintf("failed to download new release: %s", err))
261
262
263
264
265
				} else {
					err = cb(resp.UpdateVersion)
					if err != nil {
						slog.Warn(fmt.Sprintf("failed to register update available with tray: %s", err))
					}
266
267
268
269
270
271
272
				}
			}
			select {
			case <-ctx.Done():
				slog.Debug("stopping background update checker")
				return
			default:
273
				time.Sleep(UpdateCheckInterval)
274
275
276
277
			}
		}
	}()
}