ynab-portfolio-monitor/questrade/client.go

package questrade

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"time"
)

const loginServerURL = "https://login.questrade.com/oauth2/"

type LoginCredentials struct {
	AccessToken  string `json:"access_token"`
	TokenType    string `json:"token_type"`
	ExpiresIn    int    `json:"expires_in"`
	RefreshToken string `json:"refresh_token"`
	ApiServer    string `json:"api_server"`
}

// A client is the structure that will be used to consume the API
// endpoints. It holds the login credentials, http client/transport,
// rate limit information, and the login session timer.
type Client struct {
	Credentials        LoginCredentials
	SessionTimer       *time.Timer
	RateLimitRemaining int
	RateLimitReset     time.Time
	httpClient         *http.Client
	transport          *http.Transport
}

// authHeader is a shortcut that returns a string to be placed
// in the authorization header of API calls
func (l LoginCredentials) authHeader() string {
	return l.TokenType + " " + l.AccessToken
}

// Send an HTTP GET request, and return the processed response
func (c *Client) get(endpoint string, out interface{}, query url.Values) error {
	req, err := http.NewRequest("GET", c.Credentials.ApiServer+endpoint+query.Encode(), nil)
	if err != nil {
		return err
	}
	req.Header.Add("Authorization", c.Credentials.authHeader())

	res, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("http get request failed: %v", err)
	}

	err = c.processResponse(res, out)
	if err != nil {
		return err
	}
	return nil
}

// processResponse takes the body of an HTTP response, and either returns
// the error code, or unmarshalls the JSON response, extracts
// rate limit info, and places it into the object
// output parameter. This function closes the response body after reading it.
func (c *Client) processResponse(res *http.Response, out interface{}) error {
	body, err := io.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		return err
	}

	if res.StatusCode != 200 {
		return newQuestradeError(res, body)
	}

	err = json.Unmarshal(body, out)
	if err != nil {
		return err
	}

	reset, _ := strconv.Atoi(res.Header.Get("X-RateLimit-Reset"))
	c.RateLimitReset = time.Unix(int64(reset), 0)
	c.RateLimitRemaining, _ = strconv.Atoi(res.Header.Get("X-RateLimit-Remaining"))

	return nil
}

// Login takes the refresh token from the client login credentials
// and exchanges it for an access token. Returns a timer that
// expires when the login session is over.
// TODO - Return a proper error when login fails with HTTP 400 - Bad Request
func (c *Client) Login() error {
	login := loginServerURL

	vars := url.Values{"grant_type": {"refresh_token"}, "refresh_token": {c.Credentials.RefreshToken}}
	res, err := c.httpClient.PostForm(login+"token", vars)

	if err != nil {
		return err
	}

	err = c.processResponse(res, &c.Credentials)
	if err != nil {
		return err
	}

	c.SessionTimer = time.NewTimer(time.Duration(c.Credentials.ExpiresIn) * time.Second)

	return nil
}

// NewClient is the factory function for clients - takes a refresh token and logs in
func NewClient(refreshToken string) (*Client, error) {
	transport := &http.Transport{
		ResponseHeaderTimeout: 5 * time.Second,
	}

	client := &http.Client{
		Transport: transport,
	}

	// Create a new client
	c := &Client{
		Credentials: LoginCredentials{
			RefreshToken: refreshToken,
		},
		httpClient: client,
		transport:  transport,
	}

	err := c.Login()
	if err != nil {
		return nil, err
	}

	return c, nil
}