Merge pull request #44 from bcwaldon/ip-sub

feat($ip): Substitute $[public|private]_ipv4 in whole user-data

coreos-cloudinit.go

@@ -1,8 +1,6 @@
 package main
 
 import (
-	"bufio"
-	"bytes"
 	"flag"
 	"fmt"
 	"log"
@@ -58,7 +56,7 @@ func main() {
 	}
 
 	log.Printf("Fetching user-data from datasource of type %q", ds.Type())
-	userdata, err := ds.Fetch()
+	userdataBytes, err := ds.Fetch()
 	if err != nil {
 		log.Printf("Failed fetching user-data from datasource: %v", err)
 		if ignoreFailure {
@@ -68,11 +66,16 @@ func main() {
 		}
 	}
 
-	if len(userdata) == 0 {
+	if len(userdataBytes) == 0 {
 		log.Printf("No user data to handle, exiting.")
 		os.Exit(0)
 	}
 
+	env := initialize.NewEnvironment("/", workspace)
+
+	userdata := string(userdataBytes)
+	userdata = env.Apply(userdata)
+
 	parsed, err := ParseUserData(userdata)
 	if err != nil {
 		log.Printf("Failed parsing user-data: %v", err)
@@ -83,7 +86,6 @@ func main() {
 		}
 	}
 
-	env := initialize.NewEnvironment("/", workspace)
 	err = initialize.PrepWorkspace(env.Workspace())
 	if err != nil {
 		log.Fatalf("Failed preparing workspace: %v", err)
@@ -107,16 +109,14 @@ func main() {
 	}
 }
 
-func ParseUserData(contents []byte) (interface{}, error) {
-	bytereader := bytes.NewReader(contents)
-	bufreader := bufio.NewReader(bytereader)
-	header, _ := bufreader.ReadString('\n')
+func ParseUserData(contents string) (interface{}, error) {
+	header := strings.SplitN(contents, "\n", 2)[0]
 
 	if strings.HasPrefix(header, "#!") {
 		log.Printf("Parsing user-data as script")
 		return system.Script(contents), nil
 
-	} else if header == "#cloud-config\n" {
+	} else if header == "#cloud-config" {
 		log.Printf("Parsing user-data as cloud-config")
 		cfg, err := initialize.NewCloudConfig(contents)
 		if err != nil {

initialize/config.go

@@ -21,9 +21,9 @@ type CloudConfig struct {
 	Users      []system.User
 }
 
-func NewCloudConfig(contents []byte) (*CloudConfig, error) {
+func NewCloudConfig(contents string) (*CloudConfig, error) {
 	var cfg CloudConfig
-	err := goyaml.Unmarshal(contents, &cfg)
+	err := goyaml.Unmarshal([]byte(contents), &cfg)
 	return &cfg, err
 }
 

initialize/config_test.go

@@ -7,7 +7,7 @@ import (
 
 // Assert that the parsing of a cloud config file "generally works"
 func TestCloudConfigEmpty(t *testing.T) {
-	cfg, err := NewCloudConfig([]byte{})
+	cfg, err := NewCloudConfig("")
 	if err != nil {
 		t.Fatalf("Encountered unexpected error :%v", err)
 	}
@@ -28,7 +28,7 @@ func TestCloudConfigEmpty(t *testing.T) {
 
 // Assert that the parsing of a cloud config file "generally works"
 func TestCloudConfig(t *testing.T) {
-	contents := []byte(`
+	contents := `
 coreos: 
   etcd:
     discovery: "https://discovery.etcd.io/827c73219eeb2fa5530027c37bf18877"
@@ -56,7 +56,7 @@ write_files:
     permissions: '0644'
     owner: root:dogepack
 hostname: trontastic
-`)
+`
 	cfg, err := NewCloudConfig(contents)
 	if err != nil {
 		t.Fatalf("Encountered unexpected error :%v", err)
@@ -123,10 +123,10 @@ Address=10.209.171.177/19
 
 // Assert that our interface conversion doesn't panic
 func TestCloudConfigKeysNotList(t *testing.T) {
-	contents := []byte(`
+	contents := `
 ssh_authorized_keys:
   - foo: bar
-`)
+`
 	cfg, err := NewCloudConfig(contents)
 	if err != nil {
 		t.Fatalf("Encountered unexpected error :%v", err)
@@ -139,7 +139,7 @@ ssh_authorized_keys:
 }
 
 func TestCloudConfigSerializationHeader(t *testing.T) {
-	cfg, _ := NewCloudConfig([]byte{})
+	cfg, _ := NewCloudConfig("")
 	contents := cfg.String()
 	header := strings.SplitN(contents, "\n", 2)[0]
 	if header != "#cloud-config" {
@@ -148,7 +148,7 @@ func TestCloudConfigSerializationHeader(t *testing.T) {
 }
 
 func TestCloudConfigUsers(t *testing.T) {
-	contents := []byte(`
+	contents := `
 users:
   - name: elroy
     passwd: somehash
@@ -164,7 +164,7 @@ users:
     no-user-group: true
     system: y
     no-log-init: True
-`)
+`
 	cfg, err := NewCloudConfig(contents)
 	if err != nil {
 		t.Fatalf("Encountered unexpected error: %v", err)

initialize/env.go

@@ -1,19 +1,26 @@
 package initialize
 
 import (
+	"os"
 	"path"
+	"strings"
 )
 
 const DefaultSSHKeyName = "coreos-cloudinit"
 
 type Environment struct {
-	root      string
-	workspace string
-	sshKeyName string
+	root          string
+	workspace     string
+	sshKeyName    string
+	substitutions map[string]string
 }
 
 func NewEnvironment(root, workspace string) *Environment {
-	return &Environment{root, workspace, DefaultSSHKeyName}
+	substitutions := map[string]string{
+		"$public_ipv4":  os.Getenv("COREOS_PUBLIC_IPV4"),
+		"$private_ipv4": os.Getenv("COREOS_PRIVATE_IPV4"),
+	}
+	return &Environment{root, workspace, DefaultSSHKeyName, substitutions}
 }
 
 func (self *Environment) Workspace() string {
@@ -31,3 +38,10 @@ func (self *Environment) SSHKeyName() string {
 func (self *Environment) SetSSHKeyName(name string) {
 	self.sshKeyName = name
 }
+
+func (self *Environment) Apply(data string) string {
+	for key, val := range self.substitutions {
+		data = strings.Replace(data, key, val, -1)
+	}
+	return data
+}

initialize/env_test.go

@@ -0,0 +1,27 @@
+package initialize
+
+import (
+	"os"
+	"testing"
+)
+
+func TestEnvironmentApply(t *testing.T) {
+	os.Setenv("COREOS_PUBLIC_IPV4", "192.0.2.3")
+	os.Setenv("COREOS_PRIVATE_IPV4", "192.0.2.203")
+	env := NewEnvironment("./", "./")
+	input := `[Service]
+ExecStart=/usr/bin/echo "$public_ipv4"
+ExecStop=/usr/bin/echo $private_ipv4
+ExecStop=/usr/bin/echo $unknown
+`
+	expected := `[Service]
+ExecStart=/usr/bin/echo "192.0.2.3"
+ExecStop=/usr/bin/echo 192.0.2.203
+ExecStop=/usr/bin/echo $unknown
+`
+
+	output := env.Apply(input)
+	if output != expected {
+		t.Fatalf("Environment incorrectly applied.\nOutput:\n%s\nExpected:\n%s", output, expected)
+	}
+}

initialize/etcd.go

@@ -3,7 +3,6 @@ package initialize
 import (
 	"errors"
 	"fmt"
-	"os"
 	"path"
 	"strings"
 
@@ -32,20 +31,9 @@ func (ec EtcdEnvironment) String() (out string) {
 		}
 	}
 
-	public := os.Getenv("COREOS_PUBLIC_IPV4")
-	private := os.Getenv("COREOS_PRIVATE_IPV4")
-
 	out += "[Service]\n"
 
 	for key, val := range norm {
-		if public != "" {
-			val = strings.Replace(val, "$public_ipv4", public, -1)
-		}
-
-		if private != "" {
-			val = strings.Replace(val, "$private_ipv4", private, -1)
-		}
-
 		out += fmt.Sprintf("Environment=\"ETCD_%s=%s\"\n", key, val)
 	}
 

initialize/etcd_test.go

@@ -58,25 +58,6 @@ Environment="ETCD_PEER_BIND_ADDR=127.0.0.1:7002"
 	}
 }
 
-func TestEtcdEnvironmentReplacement(t *testing.T) {
-	os.Clearenv()
-	os.Setenv("COREOS_PUBLIC_IPV4", "203.0.113.29")
-	os.Setenv("COREOS_PRIVATE_IPV4", "192.0.2.13")
-
-	cfg := make(EtcdEnvironment, 0)
-	cfg["bind-addr"] = "$public_ipv4:4001"
-	cfg["peer-bind-addr"] = "$private_ipv4:7001"
-
-	env := cfg.String()
-	expect := `[Service]
-Environment="ETCD_BIND_ADDR=203.0.113.29:4001"
-Environment="ETCD_PEER_BIND_ADDR=192.0.2.13:7001"
-`
-	if env != expect {
-		t.Errorf("Generated environment:\n%s\nExpected environment:\n%s", env, expect)
-	}
-}
-
 func TestEtcdEnvironmentWrittenToDisk(t *testing.T) {
 	ec := EtcdEnvironment{
 		"name": "node001",

initialize/github_test.go

@@ -45,11 +45,11 @@ func TestCloudConfigUsersGithubMarshal(t *testing.T) {
 }
 func TestCloudConfigUsersGithubUser(t *testing.T) {
 
-	contents := []byte(`
+	contents := `
 users:
   - name: elroy
     coreos-ssh-import-github: bcwaldon
-`)
+`
 	cfg, err := NewCloudConfig(contents)
 	if err != nil {
 		t.Fatalf("Encountered unexpected error: %v", err)