add support for consul token auth

2022-08-22 12:26:35 -07:00
parent e285d3a1c5
commit e1825b4581
5 changed files with 68 additions and 18 deletions
--- a/controller/consul.go
+++ b/controller/consul.go
@@ -15,6 +15,7 @@ import (

 const (
 	consulNodeEnv        = "CONSUL_NODE"
+	consulToken          = "CONSUL_TOKEN"
 	allowStale           = "CONSUL_STALE"
 	matchTag             = "enable_gocast"
 	nodeURL              = "/catalog/node"
@@ -23,7 +24,7 @@ const (
 )

 type Clienter interface {
-	Get(url string) (*http.Response, error)
+	Do(req *http.Request) (*http.Response, error)
 }

 type Client struct {
@@ -32,6 +33,7 @@ type Client struct {

 type ConsulMon struct {
 	addr   string
+	token  string
 	node   string
 	client Clienter
 }
@@ -53,12 +55,26 @@ func contains(inp []string, elem string) bool {
 	return false
 }

-func NewConsulMon(addr string) (*ConsulMon, error) {
+func NewConsulMon(addr string, token string) (*ConsulMon, error) {
 	node := os.Getenv(consulNodeEnv)
 	if node == "" {
 		return nil, fmt.Errorf("%s env variable not set", consulNodeEnv)
 	}
-	return &ConsulMon{addr: addr, node: node, client: &http.Client{Timeout: 10 * time.Second}}, nil
+	return &ConsulMon{addr: addr, token: token, node: node, client: &http.Client{Timeout: 10 * time.Second}}, nil
+}
+
+func getHTTPReq(httpMethod string, addr string, tokenFrmCfg string) (*http.Request, error) {
+	req, err := http.NewRequest(httpMethod, addr, nil)
+	if err != nil {
+		return nil, err
+	}
+	tokenFrmEnv := os.Getenv(consulToken)
+	if tokenFrmEnv != "" {
+		req.Header.Set("X-Consul-Token", tokenFrmEnv)
+	} else if tokenFrmCfg != "" {
+		req.Header.Set("X-Consul-Token", tokenFrmCfg)
+	}
+	return req, nil
 }

 func (c *ConsulMon) queryServices() ([]*App, error) {
@@ -68,7 +84,8 @@ func (c *ConsulMon) queryServices() ([]*App, error) {
 		stale = "stale"
 	}
 	addr := c.addr + fmt.Sprintf("%s/%s?%s", nodeURL, c.node, stale)
-	resp, err := c.client.Get(addr)
+	req, err := getHTTPReq(http.MethodGet, addr, c.token)
+	resp, err := c.client.Do(req)
 	if err != nil {
 		return apps, err
 	}
@@ -125,7 +142,8 @@ func (c *ConsulMon) healthCheckLocal(service string) (bool, error) {
 	params := url.Values{}
 	params.Add("filter", "enable_gocast in ServiceTags")
 	addr := c.addr + fmt.Sprintf("%s?%s", localHealthCheckurl, params.Encode())
-	resp, err := c.client.Get(addr)
+	req, err := getHTTPReq(http.MethodGet, addr, c.token)
+	resp, err := c.client.Do(req)
 	if err != nil {
 		glog.V(2).Infof("Error getting %s with %s", addr, err)
 		return false, err
@@ -153,7 +171,8 @@ func (c *ConsulMon) healthCheckLocal(service string) (bool, error) {
 // This is the underlying api call: https://www.consul.io/api/health.html
 func (c *ConsulMon) healthCheckRemote(service string) (bool, error) {
 	addr := c.addr + fmt.Sprintf("%s/%s", remoteHealthCheckurl, service)
-	resp, err := c.client.Get(addr)
+	req, err := getHTTPReq(http.MethodGet, addr, c.token)
+	resp, err := c.client.Do(req)
 	if err != nil {
 		glog.V(2).Infof("Error getting %s with %s", addr, err)
 		return false, err
--- a/controller/consul_test.go
+++ b/controller/consul_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"io/ioutil"
 	"net/http"
+	"os"
 	"testing"

 	"github.com/mayuresh82/gocast/config"
@@ -77,16 +78,43 @@ var mockConsulCheckData = map[string]string{
 }

 type MockClient struct {
-	get func(url string) (*http.Response, error)
+	do func(*http.Request) (*http.Response, error)
 }

-func (c *MockClient) Get(url string) (*http.Response, error) {
-	if c.get != nil {
-		return c.get(url)
+func (c *MockClient) Do(*http.Request) (*http.Response, error) {
+	if c.do != nil {
+		return c.do(&http.Request{})
 	}
 	return nil, nil
 }

+func TestGetNewHTTPReq(t *testing.T) {
+	a := assert.New(t)
+
+	// test with consul token from config file
+	req, err := getHTTPReq("GET", "1.1.1.1", "3333-3333")
+	if err != nil {
+		t.Fatal(err)
+	}
+	a.Equal(req.Header.Get("X-Consul-Token"), "3333-3333")
+
+	// test with consul token from env variable
+	os.Setenv("CONSUL_TOKEN", "4444-4444")
+	req, err = getHTTPReq("GET", "1.1.1.1", "3333-3333")
+	os.Unsetenv("CONSUL_TOKEN")
+	if err != nil {
+		t.Fatal(err)
+	}
+	a.Equal(req.Header.Get("X-Consul-Token"), "4444-4444")
+
+	// test without consul token
+	req, err = getHTTPReq("GET", "1.1.1.1", "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	a.Equal(req.Header.Get("X-Consul-Token"), "")
+}
+
 func TestQueryServices(t *testing.T) {
 	a := assert.New(t)
 	client := &MockClient{}
@@ -95,7 +123,7 @@ func TestQueryServices(t *testing.T) {
 	}

 	// test valid app
-	client.get = func(url string) (*http.Response, error) {
+	client.do = func(*http.Request) (*http.Response, error) {
 		b := bytes.NewBuffer([]byte(mockConsulData["single-app"]))
 		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
 	}
@@ -110,7 +138,7 @@ func TestQueryServices(t *testing.T) {
 	a.True(app.Equal(apps[0]))

 	// test no match
-	client.get = func(url string) (*http.Response, error) {
+	client.do = func(*http.Request) (*http.Response, error) {
 		b := bytes.NewBuffer([]byte(mockConsulData["single-app-no-match"]))
 		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
 	}
@@ -121,7 +149,7 @@ func TestQueryServices(t *testing.T) {
 	a.Equal(0, len(apps))

 	// test missing vip
-	client.get = func(url string) (*http.Response, error) {
+	client.do = func(*http.Request) (*http.Response, error) {
 		b := bytes.NewBuffer([]byte(mockConsulData["single-app-no-vip"]))
 		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
 	}
@@ -136,7 +164,7 @@ func TestHealthCheck(t *testing.T) {

 	// test remote checks
 	cm.addr = "http://remote/check"
-	client.get = func(url string) (*http.Response, error) {
+	client.do = func(*http.Request) (*http.Response, error) {
 		b := bytes.NewBuffer([]byte(mockConsulCheckData["remote-pass"]))
 		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
 	}
@@ -145,7 +173,7 @@ func TestHealthCheck(t *testing.T) {
 		a.FailNow(err.Error())
 	}
 	a.True(check)
-	client.get = func(url string) (*http.Response, error) {
+	client.do = func(*http.Request) (*http.Response, error) {
 		b := bytes.NewBuffer([]byte(mockConsulCheckData["remote-fail"]))
 		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
 	}
@@ -154,7 +182,7 @@ func TestHealthCheck(t *testing.T) {

 	// test local checks
 	cm.addr = "http://localhost/check"
-	client.get = func(url string) (*http.Response, error) {
+	client.do = func(*http.Request) (*http.Response, error) {
 		b := bytes.NewBuffer([]byte(mockConsulCheckData["local-pass"]))
 		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
 	}
@@ -164,7 +192,7 @@ func TestHealthCheck(t *testing.T) {
 	}
 	a.True(check)
 	cm.addr = "http://127.0.0.1/check"
-	client.get = func(url string) (*http.Response, error) {
+	client.do = func(*http.Request) (*http.Response, error) {
 		b := bytes.NewBuffer([]byte(mockConsulCheckData["local-fail"]))
 		return &http.Response{Body: ioutil.NopCloser(b), StatusCode: http.StatusOK}, nil
 	}
--- a/controller/monitor.go
+++ b/controller/monitor.go
@@ -84,7 +84,7 @@ func NewMonitor(config *c.Config) *MonitorMgr {
 		cleanups: make(map[string]chan bool),
 	}
 	if config.Agent.ConsulAddr != "" {
-		cmon, err := NewConsulMon(config.Agent.ConsulAddr)
+		cmon, err := NewConsulMon(config.Agent.ConsulAddr, config.Agent.ConsulToken)
 		if err != nil {
 			glog.Errorf("Failed to start consul monitor: %v", err)
 		} else {