From b2581f5eb32df702401cd472fb729a341194359e Mon Sep 17 00:00:00 2001
From: Chris Hasenpflug <github@chris.hasenpflug.net>
Date: Fri, 2 Oct 2020 23:35:17 -0500
Subject: [PATCH] client: add content-type and csrf headers

required for UnifiOS controllers (UDM, UDM Pro)

Ref paultyng/terraform-provider-unifi#55
---
 unifi/unifi.go | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/unifi/unifi.go b/unifi/unifi.go
index c3130fb..2bfe71e 100644
--- a/unifi/unifi.go
+++ b/unifi/unifi.go
@@ -12,6 +12,7 @@ import (
 	"net/url"
 	"path"
 	"strings"
+	"sync"
 )
 
 const (
@@ -38,11 +39,20 @@ func (err *APIError) Error() string {
 }
 
 type Client struct {
+	// single thread client calls for CSRF, etc.
+	sync.Mutex
+
 	c       *http.Client
 	baseURL *url.URL
 
 	apiPath   string
 	loginPath string
+
+	csrf string
+}
+
+func (c *Client) CSRFToken() string {
+	return c.csrf
 }
 
 func (c *Client) SetBaseURL(base string) error {
@@ -133,6 +143,10 @@ func (c *Client) Login(ctx context.Context, user, pass string) error {
 }
 
 func (c *Client) do(ctx context.Context, method, relativeURL string, reqBody interface{}, respBody interface{}) error {
+	// single threading requests, this is mostly to assist in CSRF token propagation
+	c.Lock()
+	defer c.Unlock()
+
 	var (
 		reqReader io.Reader
 		err       error
@@ -161,6 +175,11 @@ func (c *Client) do(ctx context.Context, method, relativeURL string, reqBody int
 	}
 
 	req.Header.Set("User-Agent", "terraform-provider-unifi/0.1")
+	req.Header.Add("Content-Type", "application/json; charset=utf-8")
+
+	if c.csrf != "" {
+		req.Header.Set("X-CSRF-Token", c.csrf)
+	}
 
 	resp, err := c.c.Do(req)
 	if err != nil {
@@ -172,6 +191,10 @@ func (c *Client) do(ctx context.Context, method, relativeURL string, reqBody int
 		return &NotFoundError{}
 	}
 
+	if csrf := resp.Header.Get("x-csrf-token"); csrf != "" {
+		c.csrf = resp.Header.Get("x-csrf-token")
+	}
+
 	if resp.StatusCode != 200 {
 		fmt.Printf("Request Body:\n%s\n", string(reqBytes))
 		errBody := struct {