ente/cli/internal/api/client.go
2024-03-01 12:39:01 +05:30

98 lines
2.1 KiB
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
}