Merge pull request #1064 from ecnahc515/expand_env2

Second attempt at expanding env variables in config

CHANGES.txt

@@ -12,6 +12,8 @@ Bug Handling
 Features
 --------
 
+* Support environment variables in config files (#1023).
+
 0.7.2 (2014-MM-DD)
 ==================
 

cmd/hekad/config.go

@@ -83,12 +83,20 @@ func LoadHekadConfig(configPath string) (config *HekadConfig, err error) {
 				continue
 			}
 			fPath := filepath.Join(configPath, fName)
-			if _, err = toml.DecodeFile(fPath, &configFile); err != nil {
+			contents, err := pipeline.ReplaceEnvsFile(fPath)
+			if err != nil {
+				return nil, err
+			}
+			if _, err = toml.Decode(contents, &configFile); err != nil {
 				return nil, fmt.Errorf("Error decoding config file: %s", err)
 			}
 		}
 	} else {
-		if _, err = toml.DecodeFile(configPath, &configFile); err != nil {
+		contents, err := pipeline.ReplaceEnvsFile(configPath)
+		if err != nil {
+			return nil, err
+		}
+		if _, err = toml.Decode(contents, &configFile); err != nil {
 			return nil, fmt.Errorf("Error decoding config file: %s", err)
 		}
 	}

docs/source/config/index.rst

@@ -218,6 +218,23 @@ Example hekad.toml file
 
 .. end-hekad-toml
 
+Using Environment Variables
+===========================
+
+If you wish to use environmental variables in your config files as a way to
+configure values, you can simply use ``%ENV[VARIABLE_NAME]`` and the text will
+be replaced with the value of the environmental variable ``VARIABLE_NAME``.
+
+Example:
+
+.. code-block:: ini
+
+    [AMQPInput]
+    url = "amqp://%ENV[USER]:%ENV[PASSWORD]@rabbitmq/"
+    exchange = "testout"
+    exchangeType = "fanout"
+
+
 .. start-restarting
 
 .. _configuring_restarting:

pipeline/config.go

@@ -17,9 +17,14 @@
 package pipeline
 
 import (
+	"bufio"
+	"bytes"
 	"code.google.com/p/go-uuid/uuid"
+	"errors"
 	"fmt"
 	"github.com/bbangert/toml"
+	"io"
+	"io/ioutil"
 	"log"
 	"os"
 	"reflect"
@@ -28,9 +33,17 @@ import (
 	"time"
 )
 
-const HEKA_DAEMON = "hekad"
+const (
+	HEKA_DAEMON     = "hekad"
+	invalidEnvChars = "\n\r\t "
+)
 
-var AvailablePlugins = make(map[string]func() interface{})
+var (
+	invalidEnvPrefix     = []byte("%ENV[")
+	AvailablePlugins     = make(map[string]func() interface{})
+	ErrMissingCloseDelim = errors.New("Missing closing delimiter")
+	ErrInvalidChars      = errors.New("Invalid characters in environmental variable")
+)
 
 // Adds a plugin to the set of usable Heka plugins that can be referenced from
 // a Heka config file.
@@ -814,7 +827,12 @@ const protobufEncoderToml = `
 func (self *PipelineConfig) LoadFromConfigFile(filename string) (err error) {
 	var configFile ConfigFile
 
-	if _, err = toml.DecodeFile(filename, &configFile); err != nil {
+	contents, err := ReplaceEnvsFile(filename)
+	if err != nil {
+		return err
+	}
+
+	if _, err = toml.Decode(contents, &configFile); err != nil {
 		return fmt.Errorf("Error decoding config file: %s", err)
 	}
 
@@ -960,3 +978,80 @@ func subsFromSection(section toml.Primitive) []string {
 	}
 	return subs
 }
+
+func ReplaceEnvsFile(path string) (string, error) {
+	file, err := os.Open(path)
+	if err != nil {
+		return "", err
+	}
+	r, err := EnvSub(file)
+	if err != nil {
+		return "", err
+	}
+	contents, err := ioutil.ReadAll(r)
+	if err != nil {
+		return "", err
+	}
+	return string(contents), nil
+}
+
+func EnvSub(r io.Reader) (io.Reader, error) {
+	bufIn := bufio.NewReader(r)
+	bufOut := new(bytes.Buffer)
+	for {
+		chunk, err := bufIn.ReadBytes(byte('%'))
+		if err != nil {
+			if err == io.EOF {
+				// We're done.
+				bufOut.Write(chunk)
+				break
+			}
+			return nil, err
+		}
+		bufOut.Write(chunk[:len(chunk)-1])
+
+		tmp := make([]byte, 4)
+		tmp, err = bufIn.Peek(4)
+		if err != nil {
+			if err == io.EOF {
+				// End of file, write the last few bytes out and exit.
+				bufOut.WriteRune('%')
+				bufOut.Write(tmp)
+				break
+			}
+			return nil, err
+		}
+
+		if string(tmp) == "ENV[" {
+			// Found opening delimiter, advance the read cursor and look for
+			// closing delimiter.
+			tmp, err = bufIn.ReadBytes(byte('['))
+			if err != nil {
+				// This shouldn't happen, since the Peek succeeded.
+				return nil, err
+			}
+			chunk, err = bufIn.ReadBytes(byte(']'))
+			if err != nil {
+				if err == io.EOF {
+					// No closing delimiter, return an error
+					return nil, ErrMissingCloseDelim
+				}
+				return nil, err
+			}
+			// `chunk` is now holding var name + closing delimiter.
+			// var name contains invalid characters, return an error
+			if bytes.IndexAny(chunk, invalidEnvChars) != -1 ||
+				bytes.Index(chunk, invalidEnvPrefix) != -1 {
+				return nil, ErrInvalidChars
+			}
+			varName := string(chunk[:len(chunk)-1])
+			varVal := os.Getenv(varName)
+			bufOut.WriteString(varVal)
+		} else {
+			// Just a random '%', not an opening delimiter, write it out and
+			// keep going.
+			bufOut.WriteRune('%')
+		}
+	}
+	return bufOut, nil
+}

plugins/config_test.go

@@ -26,6 +26,7 @@ import (
 	ts "github.com/mozilla-services/heka/plugins/testsupport"
 	_ "github.com/mozilla-services/heka/plugins/udp"
 	gs "github.com/rafrombrc/gospec/src/gospec"
+	"os"
 	"path/filepath"
 	"runtime"
 )
@@ -106,6 +107,30 @@ func LoadFromConfigSpec(c gs.Context) {
 			udp.Input().Stop()
 		})
 
+		c.Specify("works with env variables in config file", func() {
+			err := os.Setenv("LOG_ENCODER", "PayloadEncoder")
+			defer os.Setenv("LOG_ENCODER", "")
+			c.Assume(err, gs.IsNil)
+			err = pipeConfig.LoadFromConfigFile("./testsupport/config_env_test.toml")
+			c.Assume(err, gs.IsNil)
+
+			log, ok := pipeConfig.OutputRunners["LogOutput"]
+			c.Expect(ok, gs.IsTrue)
+			var encoder Encoder
+			encoder = log.Encoder()
+			_, ok = encoder.(*PayloadEncoder)
+			c.Expect(ok, gs.IsTrue)
+		})
+
+		c.Specify("returns an error with invalid env variables in config", func() {
+			err := pipeConfig.LoadFromConfigFile("./testsupport/bad_envs/config_1_test.toml")
+			c.Expect(err, gs.Equals, ErrMissingCloseDelim)
+
+			err = pipeConfig.LoadFromConfigFile("./testsupport/bad_envs/config_2_test.toml")
+			c.Expect(err, gs.Equals, ErrInvalidChars)
+
+		})
+
 		c.Specify("works w/ decoder defaults", func() {
 			err := pipeConfig.LoadFromConfigFile("./testsupport/config_test_defaults.toml")
 			c.Assume(err, gs.Not(gs.IsNil))

plugins/testsupport/bad_envs/config_1_test.toml

@@ -0,0 +1,5 @@
+[LogOutput]
+type = "LogOutput"
+message_matcher = "TRUE"
+encoder = "%ENV[this_has_no_ending"
+

plugins/testsupport/bad_envs/config_2_test.toml

@@ -0,0 +1,7 @@
+[TestOut]
+type = "LogOutput"
+message_matcher = "TRUE"
+encoder = "%ENV[ spaces arent allowed]"
+
+
+

plugins/testsupport/config_env_test.toml

@@ -0,0 +1,6 @@
+[%ENV[LOG_ENCODER]]
+
+[LogOutput]
+type = "LogOutput"
+message_matcher = "TRUE"
+encoder = "%ENV[LOG_ENCODER]"