Repository: ollama/ollama
File: server/images.go (Go)

Code before:

var errUnauthorized = fmt.Errorf("unauthorized")

func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *RegistryOptions) (*http.Response, error) {
	for try := 0; try < maxRetries; try++ {
		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
		if err != nil {
			log.Printf("couldn't start upload: %v", err)
			return nil, err
		}

		switch {
		case resp.StatusCode == http.StatusUnauthorized:
			auth := resp.Header.Get("www-authenticate")
			authRedir := ParseAuthRedirectString(auth)
			token, err := getAuthToken(ctx, authRedir)
			if err != nil {
				return nil, err
			}

			regOpts.Token = token
			if body != nil {
				body.Seek(0, io.SeekStart)
			}

			if try == (maxRetries - 1) {
				// last retry reached
				return nil, errUnauthorized
			}
		case resp.StatusCode == http.StatusNotFound:
			return nil, os.ErrNotExist
		case resp.StatusCode >= http.StatusBadRequest:
			body, 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, body)
		default:
			return resp, nil
		}
	}

	return nil, errMaxRetriesExceeded
}

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


Review comment (suggestion):
could we set an outer `err` variable instead and then just return it at the end? feels strange to check boundary conditions

Code after:

var errUnauthorized = fmt.Errorf("unauthorized")

func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *RegistryOptions) (*http.Response, error) {
	lastErr := errMaxRetriesExceeded
	for try := 0; try < maxRetries; try++ {
		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
		if err != nil {
			log.Printf("couldn't start upload: %v", err)
			return nil, err
		}

		switch {
		case resp.StatusCode == http.StatusUnauthorized:
			auth := resp.Header.Get("www-authenticate")
			authRedir := ParseAuthRedirectString(auth)
			token, err := getAuthToken(ctx, authRedir)
			if err != nil {
				return nil, err
			}

			regOpts.Token = token
			if body != nil {
				body.Seek(0, io.SeekStart)
			}
			lastErr = errUnauthorized
		case resp.StatusCode == http.StatusNotFound:
			return nil, os.ErrNotExist
		case resp.StatusCode >= http.StatusBadRequest:
			body, 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, body)
		default:
			return resp, nil
		}
	}

	return nil, lastErr
}

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