Repository: jumptrading/influx-spout
Pull request #126 (2018-09-18 03:38:35)

Support InfluxDB HTTP basic authentication
The username and password are read from two environment variables (INFLUXDB_USER and INFLUXDB_PASS).


Patch:
diff --git a/README.md b/README.md
index de0ba37..e294750 100644
--- a/README.md
+++ b/README.md
@@ -491,6 +491,11 @@ A writer will batch up messages until one of the limits defined by the
 `batch_max_count`, `batch_max_size` or `batch_max_age` options is
 reached.
 
+The writer also supports optional HTTP basic authentication to the
+InfluxDB backend. The username and password are set using the
+`INFLUXDB_USER` and `INFLUXDB_PASS` environment variables. These
+options can't be set via the configuration file.
+
 ### Monitor
 
 The monitor is responsible for collecting metrics from the other
diff --git a/config/config.go b/config/config.go
index b11ddac..6a4d0ef 100644
--- a/config/config.go
+++ b/config/config.go
@@ -44,6 +44,8 @@ type Config struct {
 	NATSSubjectJunkyard string            `toml:"nats_subject_junkyard"`
 	InfluxDBAddress     string            `toml:"influxdb_address"`
 	InfluxDBPort        int               `toml:"influxdb_port"`
+	InfluxDBUser        string            `toml:"-"`
+	InfluxDBPass        string            `toml:"-"`
 	DBName              string            `toml:"influxdb_dbname"`
 	BatchMaxCount       int               `toml:"batch_max_count"`
 	BatchMaxSize        datasize.ByteSize `toml:"batch_max_size"`
@@ -189,6 +191,12 @@ func (c *Config) validateWriter() error {
 	if c.InfluxDBPort < 1 || c.InfluxDBPort > 65535 {
 		return errors.New("influxdb_port out of range")
 	}
+	if c.InfluxDBUser != "" && c.InfluxDBPass == "" {
+		return errors.New("$INFLUXDB_USER without $INFLUXDB_PASS")
+	}
+	if c.InfluxDBUser == "" && c.InfluxDBPass != "" {
+		return errors.New("$INFLUXDB_PASS without $INFLUXDB_USER")
+	}
 
 	for _, rule := range c.Rule {
 		if rule.Rtype == "" {
@@ -248,6 +256,11 @@ func NewConfigFromFile(fileName string) (*Config, error) {
 		conf.DownsamplePeriod.Duration = time.Minute
 	}
 
+	if conf.Mode == "writer" {
+		conf.InfluxDBUser = os.Getenv("INFLUXDB_USER")
+		conf.InfluxDBPass = os.Getenv("INFLUXDB_PASS")
+	}
+
 	if err := conf.Validate(); err != nil {
 		return nil, err
 	}
diff --git a/writer/writer.go b/writer/writer.go
index afb56f7..c08cc3d 100644
--- a/writer/writer.go
+++ b/writer/writer.go
@@ -217,8 +217,15 @@ func (w *Writer) shouldSend(batch *batch.Batch) bool {
 
 // sendBatch sends the accumulated batch via HTTP to InfluxDB.
 func (w *Writer) sendBatch(batch *batch.Batch, client *http.Client) error {
-	reader := bytes.NewReader(batch.Bytes())
-	resp, err := client.Post(w.url, "application/json; charset=UTF-8", reader)
+	req, err := http.NewRequest("POST", w.url, bytes.NewReader(batch.Bytes()))
+	if err != nil {
+		return err
+	}
+	req.Header.Add("Content-Type", "application/json; charset=UTF-8")
+	if w.c.InfluxDBUser != "" {
+		req.SetBasicAuth(w.c.InfluxDBUser, w.c.InfluxDBPass)
+	}
+	resp, err := client.Do(req)
 	if err != nil {
 		return fmt.Errorf("failed to send HTTP request: %v\n", err)
 	}


Test patch:
diff --git a/config/config_small_test.go b/config/config_small_test.go
index 096bc1c..5b500f4 100644
--- a/config/config_small_test.go
+++ b/config/config_small_test.go
@@ -18,6 +18,7 @@ package config
 
 import (
 	"errors"
+	"os"
 	"path/filepath"
 	"testing"
 	"time"
@@ -32,6 +33,12 @@ const testConfigBaseName = "something"
 
 var testConfigFileName = filepath.Join("some", "dir", testConfigBaseName+".toml")
 
+func init() {
+	// Clear these to ensure consistent test environment.
+	os.Setenv("INFLUXDB_USER", "")
+	os.Setenv("INFLUXDB_PASS", "")
+}
+
 func TestCorrectConfigFile(t *testing.T) {
 	const validConfigSample = `
 name = "thor"
@@ -101,6 +108,8 @@ func TestAllDefaults(t *testing.T) {
 	assert.Equal(t, "influx-spout-junk", conf.NATSSubjectJunkyard)
 	assert.Equal(t, "localhost", conf.InfluxDBAddress)
 	assert.Equal(t, 8086, conf.InfluxDBPort)
+	assert.Equal(t, "", conf.InfluxDBUser)
+	assert.Equal(t, "", conf.InfluxDBPass)
 	assert.Equal(t, "influx-spout-junk", conf.DBName)
 	assert.Equal(t, 10, conf.BatchMaxCount)
 	assert.Equal(t, 10*datasize.MB, conf.BatchMaxSize)
@@ -157,6 +166,18 @@ func TestDefaultDownsamplePeriod(t *testing.T) {
 	assert.Equal(t, time.Minute, conf.DownsamplePeriod.Duration)
 }
 
+func TestInfluxAuth(t *testing.T) {
+	os.Setenv("INFLUXDB_USER", "user")
+	defer os.Setenv("INFLUXDB_USER", "")
+	os.Setenv("INFLUXDB_PASS", "secret")
+	defer os.Setenv("INFLUXDB_PASS", "")
+
+	conf, err := parseConfig(`mode = "writer"`)
+	require.NoError(t, err)
+	assert.Equal(t, "user", conf.InfluxDBUser)
+	assert.Equal(t, "secret", conf.InfluxDBPass)
+}
+
 func TestNoMode(t *testing.T) {
 	_, err := parseConfig("")
 	assert.EqualError(t, err, "mode not specified in config")
@@ -533,3 +554,15 @@ downsample_suffix = ""
 		assert.EqualError(t, err, test.err)
 	}
 }
+
+func TestTextValidateAuth(t *testing.T) {
+	config0 := newDefaultConfig()
+	config0.Mode = "writer"
+	config0.InfluxDBUser = "hello"
+	assert.EqualError(t, config0.Validate(), "$INFLUXDB_USER without $INFLUXDB_PASS")
+
+	config1 := newDefaultConfig()
+	config1.Mode = "writer"
+	config1.InfluxDBPass = "hello"
+	assert.EqualError(t, config1.Validate(), "$INFLUXDB_PASS without $INFLUXDB_USER")
+}
diff --git a/writer/writer_medium_test.go b/writer/writer_medium_test.go
index 76594f9..535aebc 100644
--- a/writer/writer_medium_test.go
+++ b/writer/writer_medium_test.go
@@ -116,6 +116,24 @@ func TestBasicWriter(t *testing.T) {
 	})
 }
 
+func TestInfluxDBAuth(t *testing.T) {
+	nc, closeNATS := runGnatsd(t)
+	defer closeNATS()
+
+	influxd := runTestInfluxd()
+	defer influxd.Stop()
+
+	conf := testConfig()
+	conf.InfluxDBUser = "user"
+	conf.InfluxDBPass = "pass"
+	w := startWriter(t, conf)
+	defer w.Stop()
+
+	line := "What's in a name?"
+	publish(t, nc, conf.NATSSubject[0], line)
+	influxd.AssertWriteWithAuth(t, line, "user", "pass")
+}
+
 func TestBatchMBLimit(t *testing.T) {
 	nc, closeNATS := runGnatsd(t)
 	defer closeNATS()
@@ -146,7 +164,7 @@ func TestBatchMBLimit(t *testing.T) {
 	// BatchMaxMB is exceed
 	select {
 	case msg := <-influxd.Writes:
-		assert.Len(t, msg, totalSize)
+		assert.Len(t, msg.Body, totalSize)
 	case <-time.After(spouttest.LongWait):
 		t.Fatal("timed out waiting for messages")
 	}
@@ -328,7 +346,7 @@ func publish(t require.TestingT, nc *nats.Conn, subject, msg string) {
 type testInfluxd struct {
 	server *http.Server
 	wg     sync.WaitGroup
-	Writes chan string
+	Writes chan write
 	ready  chan struct{}
 }
 
@@ -337,7 +355,7 @@ func runTestInfluxd() *testInfluxd {
 		server: &http.Server{
 			Addr: fmt.Sprintf(":%d", influxPort),
 		},
-		Writes: make(chan string, 99),
+		Writes: make(chan write, 99),
 		ready:  make(chan struct{}),
 	}
 
@@ -359,8 +377,21 @@ func (s *testInfluxd) Stop() {
 
 func (s *testInfluxd) AssertWrite(t *testing.T, expected string) {
 	select {
-	case msg := <-s.Writes:
-		assert.Equal(t, msg, expected)
+	case write := <-s.Writes:
+		assert.Equal(t, expected, write.Body)
+		assert.Equal(t, "", write.Username)
+		assert.Equal(t, "", write.Password)
+	case <-time.After(spouttest.LongWait):
+		t.Fatal("timed out waiting for message")
+	}
+}
+
+func (s *testInfluxd) AssertWriteWithAuth(t *testing.T, expected, username, password string) {
+	select {
+	case write := <-s.Writes:
+		assert.Equal(t, expected, write.Body)
+		assert.Equal(t, username, write.Username)
+		assert.Equal(t, password, write.Password)
 	case <-time.After(spouttest.LongWait):
 		t.Fatal("timed out waiting for message")
 	}
@@ -391,5 +422,14 @@ func (s *testInfluxd) handleWrite(w http.ResponseWriter, r *http.Request) {
 	}
 	w.WriteHeader(http.StatusNoContent)
 
-	s.Writes <- string(body)
+	username, password, _ := r.BasicAuth()
+	s.Writes <- write{
+		Body:     string(body),
+		Username: username,
+		Password: password,
+	}
+}
+
+type write struct {
+	Body, Username, Password string
 }