diff --git a/config.yaml b/config.yaml index f5972cc..f94aec0 100644 --- a/config.yaml +++ b/config.yaml @@ -9,6 +9,8 @@ agent: consul_addr: https://consul # interval to query consul for app discovery consul_query_interval: 5m + # token to authenticate client if consul requires it + consul_token: 00000000-0000-0000-0000-000000000000 bgp: local_as: 12345 diff --git a/config/config.go b/config/config.go index 35df0f7..8947c49 100644 --- a/config/config.go +++ b/config/config.go @@ -15,6 +15,7 @@ type AgentConfig struct { CleanupTimer time.Duration `yaml:"cleanup_timer"` ConsulAddr string `yaml:"consul_addr"` ConsulQueryInterval time.Duration `yaml:"consul_query_interval"` + ConsulToken string `yaml:"consul_token"` } type BgpConfig struct { diff --git a/controller/consul.go b/controller/consul.go index 5f53cc7..3444c35 100644 --- 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 diff --git a/controller/consul_test.go b/controller/consul_test.go index 8ddc115..2e1d9da 100644 --- 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 } diff --git a/controller/monitor.go b/controller/monitor.go index e629510..f5b9ddf 100644 --- 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 {