portainer/api/http/handler/endpoints/endpoint_dockerhub_status.go

package endpoints

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"

	portainer "github.com/portainer/portainer/api"
	"github.com/portainer/portainer/api/http/client"
	"github.com/portainer/portainer/api/internal/endpointutils"
	httperror "github.com/portainer/portainer/pkg/libhttp/error"
	"github.com/portainer/portainer/pkg/libhttp/request"
	"github.com/portainer/portainer/pkg/libhttp/response"

	"github.com/segmentio/encoding/json"
)

type dockerhubStatusResponse struct {
	// Remaiming images to pull
	Remaining int `json:"remaining"`
	// Daily limit
	Limit int `json:"limit"`
}

// @id endpointDockerhubStatus
// @summary fetch docker pull limits
// @description get docker pull limits for a docker hub registry in the environment
// @description **Access policy**:
// @tags endpoints
// @security ApiKeyAuth
// @security jwt
// @produce json
// @param id path int true "endpoint ID"
// @param registryId path int true "registry ID"
// @success 200 {object} dockerhubStatusResponse "Success"
// @failure 400 "Invalid request"
// @failure 403 "Permission denied"
// @failure 404 "registry or endpoint not found"
// @failure 500 "Server error"
// @router /endpoints/{id}/dockerhub/{registryId} [get]
func (handler *Handler) endpointDockerhubStatus(w http.ResponseWriter, r *http.Request) *httperror.HandlerError {
	endpointID, err := request.RetrieveNumericRouteVariableValue(r, "id")
	if err != nil {
		return httperror.BadRequest("Invalid environment identifier route variable", err)
	}

	endpoint, err := handler.DataStore.Endpoint().Endpoint(portainer.EndpointID(endpointID))
	if handler.DataStore.IsErrObjectNotFound(err) {
		return httperror.NotFound("Unable to find an environment with the specified identifier inside the database", err)
	} else if err != nil {
		return httperror.InternalServerError("Unable to find an environment with the specified identifier inside the database", err)
	}

	if !endpointutils.IsLocalEndpoint(endpoint) {
		return httperror.BadRequest("Invalid environment type", errors.New("Invalid environment type"))
	}

	registryID, err := request.RetrieveNumericRouteVariableValue(r, "registryId")
	if err != nil {
		return httperror.BadRequest("Invalid registry identifier route variable", err)
	}

	var registry *portainer.Registry

	if registryID == 0 {
		registry = &portainer.Registry{}
	} else {
		registry, err = handler.DataStore.Registry().Read(portainer.RegistryID(registryID))
		if handler.DataStore.IsErrObjectNotFound(err) {
			return httperror.NotFound("Unable to find a registry with the specified identifier inside the database", err)
		} else if err != nil {
			return httperror.InternalServerError("Unable to find a registry with the specified identifier inside the database", err)
		}

		if registry.Type != portainer.DockerHubRegistry {
			return httperror.BadRequest("Invalid registry type", errors.New("Invalid registry type"))
		}
	}

	httpClient := client.NewHTTPClient()
	token, err := getDockerHubToken(httpClient, registry)
	if err != nil {
		return httperror.InternalServerError("Unable to retrieve DockerHub token from DockerHub", err)
	}

	resp, err := getDockerHubLimits(httpClient, token)
	if err != nil {
		return httperror.InternalServerError("Unable to retrieve DockerHub rate limits from DockerHub", err)
	}

	return response.JSON(w, resp)
}

func getDockerHubToken(httpClient *client.HTTPClient, registry *portainer.Registry) (string, error) {
	type dockerhubTokenResponse struct {
		Token string `json:"token"`
	}

	requestURL := "https://auth.docker.io/token?service=registry.docker.io&scope=repository:ratelimitpreview/test:pull"

	req, err := http.NewRequest(http.MethodGet, requestURL, nil)
	if err != nil {
		return "", err
	}

	if registry.Authentication {
		req.SetBasicAuth(registry.Username, registry.Password)
	}

	resp, err := httpClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", errors.New("failed fetching dockerhub token")
	}

	var data dockerhubTokenResponse
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return "", err
	}

	return data.Token, nil
}

func getDockerHubLimits(httpClient *client.HTTPClient, token string) (*dockerhubStatusResponse, error) {
	requestURL := "https://registry-1.docker.io/v2/ratelimitpreview/test/manifests/latest"

	req, err := http.NewRequest(http.MethodHead, requestURL, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", token))

	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	io.Copy(io.Discard, resp.Body)
	resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, errors.New("failed fetching dockerhub limits")
	}

	// An error with rateLimit-Limit or RateLimit-Remaining is likely for dockerhub pro accounts where there is no rate limit.
	// In that specific case the headers will not be present.  Don't bubble up the error as its normal
	// See: https://docs.docker.com/docker-hub/download-rate-limit/
	rateLimit, err := parseRateLimitHeader(resp.Header, "RateLimit-Limit")
	if err != nil {
		return nil, nil
	}

	rateLimitRemaining, err := parseRateLimitHeader(resp.Header, "RateLimit-Remaining")
	if err != nil {
		return nil, nil
	}

	return &dockerhubStatusResponse{
		Limit:     rateLimit,
		Remaining: rateLimitRemaining,
	}, nil
}

func parseRateLimitHeader(headers http.Header, headerKey string) (int, error) {
	headerValue := headers.Get(headerKey)
	if headerValue == "" {
		return 0, fmt.Errorf("Missing %s header", headerKey)
	}

	matches := strings.Split(headerValue, ";")
	value, err := strconv.Atoi(matches[0])
	if err != nil {
		return 0, err
	}

	return value, nil
}