ente/cli/internal/api/client.go

package api

import (
	"context"
	"github.com/go-resty/resty/v2"
	"log"
	"time"
)

const (
	EnteAPIEndpoint = "https://api.ente.io"
	TokenHeader     = "X-Auth-Token"
	TokenQuery      = "token"
	ClientPkgHeader = "X-Client-Package"
)

var (
	RedactedHeaders = []string{TokenHeader, " X-Request-Id"}
)
var tokenMap map[string]string = make(map[string]string)

type Client struct {
	restClient *resty.Client
	// use separate client for downloading files
	downloadClient *resty.Client
}

type Params struct {
	Debug bool
	Trace bool
	Host  string
}

func readValueFromContext(ctx context.Context, key string) interface{} {
	value := ctx.Value(key)
	return value
}

func NewClient(p Params) *Client {
	enteAPI := resty.New()

	if p.Trace {
		enteAPI.EnableTrace()
	}
	enteAPI.OnBeforeRequest(func(c *resty.Client, req *resty.Request) error {
		app := readValueFromContext(req.Context(), "app")
		if app == nil {
			panic("app not set in context")
		}
		req.Header.Set(ClientPkgHeader, StringToApp(app.(string)).ClientPkg())
		attachToken(req)
		return nil
	})
	if p.Debug {
		enteAPI.OnBeforeRequest(func(c *resty.Client, req *resty.Request) error {
			logRequest(req)
			return nil
		})

		enteAPI.OnAfterResponse(func(c *resty.Client, resp *resty.Response) error {
			logResponse(resp)
			return nil
		})
	}
	if p.Host != "" {
		enteAPI.SetBaseURL(p.Host)
	} else {
		enteAPI.SetBaseURL(EnteAPIEndpoint)
	}
	return &Client{
		restClient: enteAPI,
		downloadClient: resty.New().
			SetRetryCount(3).
			SetRetryWaitTime(5 * time.Second).
			SetRetryMaxWaitTime(10 * time.Second).
			AddRetryCondition(func(r *resty.Response, err error) bool {
				shouldRetry := r.StatusCode() == 429 || r.StatusCode() > 500
				if shouldRetry {
					log.Printf("retrying download due to %d code", r.StatusCode())
				}
				return shouldRetry
			}),
	}
}

func attachToken(req *resty.Request) {
	accountKey := readValueFromContext(req.Context(), "account_key")
	if accountKey != nil && accountKey != "" {
		if token, ok := tokenMap[accountKey.(string)]; ok {
			req.SetHeader(TokenHeader, token)
		}
	}
}

func (c *Client) AddToken(id string, token string) {
	tokenMap[id] = token
}