Implement `hub-tool account rate-limiting`

Signed-off-by: Djordje Lukic <djordje.lukic@docker.com>

cmd/hub-tool/main.go

@@ -61,6 +61,8 @@ func main() {
 		hub.WithInStream(dockerCli.In()),
 		hub.WithOutStream(dockerCli.Out()),
 		hub.WithHubAccount(auth.Username),
+		hub.WithPassword(auth.Password),
+		hub.WithRefreshToken(auth.RefreshToken),
 		hub.WithHubToken(auth.Token))
 	if err != nil {
 		log.Fatal(err)

internal/ansi/ansi.go

@@ -37,8 +37,14 @@ var (
 	Error = utils.Red
 	// Emphasise color should be used with important content
 	Emphasise = utils.Green
+	// NoColor doesn't add any colors to the output
+	NoColor = noop
 )
 
+func noop(in string) string {
+	return in
+}
+
 // Link returns an ANSI terminal hyperlink
 func Link(url string, text string) string {
 	return fmt.Sprintf("\u001B]8;;%s\u0007%s\u001B]8;;\u0007", url, text)

internal/commands/account/cmd.go

@@ -42,6 +42,7 @@ func NewAccountCmd(streams command.Streams, hubClient *hub.Client) *cobra.Comman
 	}
 	cmd.AddCommand(
 		newInfoCmd(streams, hubClient, accountName),
+		newRateLimitingCmd(streams, hubClient, accountName),
 	)
 	return cmd
 }

internal/commands/account/ratelimiting.go

@@ -0,0 +1,94 @@
+/*
+   Copyright 2020 Docker Hub Tool authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package account
+
+import (
+	"fmt"
+	"io"
+	"time"
+
+	"github.com/docker/cli/cli"
+	"github.com/docker/cli/cli/command"
+	"github.com/docker/go-units"
+	"github.com/spf13/cobra"
+
+	"github.com/docker/hub-tool/internal/ansi"
+	"github.com/docker/hub-tool/internal/format"
+	"github.com/docker/hub-tool/internal/hub"
+	"github.com/docker/hub-tool/internal/metrics"
+)
+
+const (
+	rateLimitingName = "rate-limiting"
+)
+
+type rateLimitingOptions struct {
+	format.Option
+}
+
+func newRateLimitingCmd(streams command.Streams, hubClient *hub.Client, parent string) *cobra.Command {
+	var opts rateLimitingOptions
+	cmd := &cobra.Command{
+		Use:                   rateLimitingName,
+		Short:                 "Print the rate limiting information",
+		Args:                  cli.NoArgs,
+		DisableFlagsInUseLine: true,
+		Annotations: map[string]string{
+			"sudo": "true",
+		},
+		PreRun: func(cmd *cobra.Command, args []string) {
+			metrics.Send(parent, rateLimitingName)
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return runRateLimiting(streams, hubClient, opts)
+		},
+	}
+	opts.AddFormatFlag(cmd.Flags())
+
+	return cmd
+}
+
+func runRateLimiting(streams command.Streams, hubClient *hub.Client, opts rateLimitingOptions) error {
+	rl, err := hubClient.GetRateLimits()
+	if err != nil {
+		return err
+	}
+	value := &hub.RateLimits{}
+	if rl != nil {
+		value = rl
+	}
+	return opts.Print(streams.Out(), value, printRateLimit(rl))
+}
+
+func printRateLimit(rl *hub.RateLimits) func(io.Writer, interface{}) error {
+	return func(out io.Writer, _ interface{}) error {
+		if rl == nil {
+			fmt.Fprintln(out, ansi.Emphasise("Unlimited"))
+			return nil
+		}
+		color := ansi.NoColor
+		if *rl.Remaining <= 50 {
+			color = ansi.Warn
+		}
+		if *rl.Remaining <= 10 {
+			color = ansi.Error
+		}
+		fmt.Fprintf(out, color("Limit: %d, %s window\n"), rl.Limit, units.HumanDuration(time.Duration(*rl.LimitWindow)*time.Second))
+		fmt.Fprintf(out, color("Remaining: %d, %s window\n"), rl.Remaining, units.HumanDuration(time.Duration(*rl.RemainingWindow)*time.Second))
+		return nil
+	}
+}

internal/hub/client.go

@@ -50,6 +50,8 @@ type Client struct {
 
 	domain           string
 	token            string
+	refreshToken     string
+	password         string
 	account          string
 	fetchAllElements bool
 	in               io.Reader
@@ -153,6 +155,22 @@ func WithHubToken(token string) ClientOp {
 	}
 }
 
+// WithRefreshToken sets the refresh token to the client
+func WithRefreshToken(refreshToken string) ClientOp {
+	return func(c *Client) error {
+		c.refreshToken = refreshToken
+		return nil
+	}
+}
+
+// WithPassword sets the password to the client
+func WithPassword(password string) ClientOp {
+	return func(c *Client) error {
+		c.password = password
+		return nil
+	}
+}
+
 func withHubToken(token string) RequestOp {
 	return func(req *http.Request) error {
 		req.Header["Authorization"] = []string{fmt.Sprintf("Bearer %s", token)}

internal/hub/ratelimiting.go

@@ -0,0 +1,156 @@
+/*
+   Copyright 2020 Docker Hub Tool authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package hub
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"strconv"
+	"strings"
+)
+
+// RateLimits ...
+type RateLimits struct {
+	Limit           *int `json:",omitempty"`
+	LimitWindow     *int `json:",omitempty"`
+	Remaining       *int `json:",omitempty"`
+	RemainingWindow *int `json:",omitempty"`
+}
+
+const (
+	first  = "https://auth.docker.io/token?service=registry.docker.io&scope=repository:ratelimitpreview/test:pull"
+	second = "https://registry-1.docker.io/v2/ratelimitpreview/test/manifests/latest"
+)
+
+// GetRateLimits returns the rate limits for the authenticated user
+func (c *Client) GetRateLimits() (*RateLimits, error) {
+	token, err := tryGetToken(c)
+	if err != nil {
+		return nil, err
+	}
+
+	req, err := http.NewRequest("HEAD", second, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	resp, err := c.doRawRequest(req, withHubToken(token))
+	if err != nil {
+		return nil, err
+	}
+
+	limitHeader := resp.Header.Get("Ratelimit-Limit")
+	remainingHeader := resp.Header.Get("Ratelimit-Remaining")
+
+	if limitHeader == "" || remainingHeader == "" {
+		return nil, nil
+	}
+
+	limit, limitWindow, err := parseLimitHeader(limitHeader)
+	if err != nil {
+		return nil, err
+	}
+
+	remaining, remainingWindow, err := parseLimitHeader(remainingHeader)
+	if err != nil {
+		return nil, err
+	}
+
+	return &RateLimits{
+		Limit:           &limit,
+		LimitWindow:     &limitWindow,
+		Remaining:       &remaining,
+		RemainingWindow: &remainingWindow,
+	}, nil
+}
+
+func tryGetToken(c *Client) (string, error) {
+	token, err := c.getToken(c.password)
+	if err != nil {
+		token, err = c.getToken(c.refreshToken)
+		if err != nil {
+			token, err = c.getToken(c.token)
+			if err != nil {
+				return "", err
+			}
+		}
+	}
+	return token, nil
+}
+
+func (c *Client) getToken(password string) (string, error) {
+	req, err := http.NewRequest("GET", first, nil)
+	if err != nil {
+		return "", err
+	}
+
+	req.Header.Add("Authorization", "Basic "+basicAuth(c.account, password))
+	resp, err := c.doRawRequest(req)
+	if err != nil {
+		return "", err
+	}
+	if resp.Body != nil {
+		defer resp.Body.Close() //nolint:errcheck
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return "", errors.New("unable to get authorization token")
+	}
+
+	buf, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+	var t tokenResponse
+	if err := json.Unmarshal(buf, &t); err != nil {
+		return "", err
+	}
+
+	return t.Token, nil
+}
+
+func parseLimitHeader(value string) (int, int, error) {
+	parts := strings.Split(value, ";")
+	if len(parts) != 2 {
+		return 0, 0, fmt.Errorf("bad limit header %s", value)
+	}
+
+	v, err := strconv.Atoi(parts[0])
+	if err != nil {
+		return 0, 0, err
+	}
+
+	windowParts := strings.Split(parts[1], "=")
+	if len(windowParts) != 2 {
+		return 0, 0, fmt.Errorf("bad limit header %s", value)
+	}
+	w, err := strconv.Atoi(windowParts[1])
+	if err != nil {
+		return 0, 0, err
+	}
+
+	return v, w, nil
+}
+
+func basicAuth(username, password string) string {
+	auth := username + ":" + password
+	return base64.StdEncoding.EncodeToString([]byte(auth))
+}