feat(update): add more configuration options for update.conf

initialize/update.go

@@ -2,6 +2,8 @@ package initialize
 
 import (
 	"bufio"
+	"errors"
+	"fmt"
 	"os"
 	"path"
 	"strings"
@@ -9,27 +11,78 @@ import (
 	"github.com/coreos/coreos-cloudinit/system"
 )
 
-const locksmithUnit = "locksmithd.service"
+const (
+	locksmithUnit = "locksmithd.service"
+)
+
+// updateOption represents a configurable update option, which, if set, will be
+// written into update.conf, replacing any existing value for the option
+type updateOption struct {
+	key    string   // key used to configure this option in cloud-config
+	valid  []string // valid values for the option
+	prefix string   // prefix for the option in the update.conf file
+	value  string   // used to store the new value in update.conf (including prefix)
+	seen   bool     // whether the option has been seen in any existing update.conf
+}
+
+// updateOptions defines the update options understood by cloud-config.
+// The keys represent the string used in cloud-config to configure the option.
+var updateOptions = []*updateOption{
+	&updateOption{
+		key:    "reboot-strategy",
+		prefix: "REBOOT_STRATEGY=",
+		valid:  []string{"best-effort", "etcd-lock", "reboot", "off"},
+	},
+	&updateOption{
+		key:    "group",
+		prefix: "GROUP=",
+		valid:  []string{"master", "beta", "alpha", "stable"},
+	},
+	&updateOption{
+		key:    "server",
+		prefix: "SERVER=",
+	},
+}
+
+// isValid checks whether a supplied value is valid for this option
+func (uo updateOption) isValid(val string) bool {
+	if len(uo.valid) == 0 {
+		return true
+	}
+	for _, v := range uo.valid {
+		if val == v {
+			return true
+		}
+	}
+	return false
+}
 
 type UpdateConfig map[string]string
 
-func (uc UpdateConfig) strategy() string {
+// File generates an `/etc/coreos/update.conf` file (if any update
-	s, _ := uc["reboot-strategy"]
+// configuration options are set in cloud-config) by either rewriting the
-	return s
+// existing file on disk, or starting from `/usr/share/coreos/update.conf`
-}
-
-// File creates an `/etc/coreos/update.conf` file with the requested
-// strategy, by either rewriting the existing file on disk, or starting
-// from `/usr/share/coreos/update.conf`
 func (uc UpdateConfig) File(root string) (*system.File, error) {
-
+	if len(uc) < 1 {
-	// If no reboot-strategy is set, we don't need to generate a new config
-	if _, ok := uc["reboot-strategy"]; !ok {
 		return nil, nil
 	}
 
 	var out string
 
+	// Generate the list of possible substitutions to be performed based on the options that are configured
+	subs := make([]*updateOption, 0)
+	for _, uo := range updateOptions {
+		val, ok := uc[uo.key]
+		if !ok {
+			continue
+		}
+		if !uo.isValid(val) {
+			return nil, errors.New(fmt.Sprintf("invalid value %v for option %v (valid options: %v)", val, uo.key, uo.valid))
+		}
+		uo.value = uo.prefix + val
+		subs = append(subs, uo)
+	}
+
 	etcUpdate := path.Join(root, "etc", "coreos", "update.conf")
 	usrUpdate := path.Join(root, "usr", "share", "coreos", "update.conf")
 
@@ -43,13 +96,14 @@ func (uc UpdateConfig) File(root string) (*system.File, error) {
 
 	scanner := bufio.NewScanner(conf)
 
-	sawStrat := false
-	stratLine := "REBOOT_STRATEGY=" + uc.strategy()
 	for scanner.Scan() {
 		line := scanner.Text()
-		if strings.HasPrefix(line, "REBOOT_STRATEGY=") {
+		for _, s := range subs {
-			line = stratLine
+			if strings.HasPrefix(line, s.prefix) {
-			sawStrat = true
+				line = s.value
+				s.seen = true
+				break
+			}
 		}
 		out += line
 		out += "\n"
@@ -58,10 +112,13 @@ func (uc UpdateConfig) File(root string) (*system.File, error) {
 		}
 	}
 
-	if !sawStrat {
+	for _, s := range subs {
-		out += stratLine
+		if !s.seen {
-		out += "\n"
+			out += s.value
+			out += "\n"
+		}
 	}
+
 	return &system.File{
 		Path:               path.Join("etc", "coreos", "update.conf"),
 		RawFilePermissions: "0644",
@@ -69,9 +126,14 @@ func (uc UpdateConfig) File(root string) (*system.File, error) {
 	}, nil
 }
 
-// Unit generates a locksmith system.Unit for the cloud-init initializer to
+// GetUnit generates a locksmith system.Unit, if reboot-strategy was set in
-// act on appropriately
+// cloud-config, for the cloud-init initializer to act on appropriately
 func (uc UpdateConfig) Unit(root string) (*system.Unit, error) {
+	strategy, ok := uc["reboot-strategy"]
+	if !ok {
+		return nil, nil
+	}
+
 	u := &system.Unit{
 		Name:    locksmithUnit,
 		Enable:  true,
@@ -79,7 +141,7 @@ func (uc UpdateConfig) Unit(root string) (*system.Unit, error) {
 		Mask:    false,
 	}
 
-	if uc.strategy() == "off" {
+	if strategy == "off" {
 		u.Enable = false
 		u.Command = "stop"
 		u.Mask = true

initialize/update_test.go

@@ -4,6 +4,8 @@ import (
 	"io/ioutil"
 	"os"
 	"path"
+	"sort"
+	"strings"
 	"testing"
 
 	"github.com/coreos/coreos-cloudinit/system"
@@ -12,11 +14,9 @@ import (
 const (
 	base = `SERVER=https://example.com
 GROUP=thegroupc`
-
 	configured = base + `
 REBOOT_STRATEGY=awesome
 `
-
 	expected = base + `
 REBOOT_STRATEGY=etcd-lock
 `
@@ -29,7 +29,143 @@ func setupFixtures(dir string) {
 	ioutil.WriteFile(path.Join(dir, "usr", "share", "coreos", "update.conf"), []byte(base), 0644)
 }
 
-func TestLocksmithEnvironmentWrittenToDisk(t *testing.T) {
+func TestEmptyUpdateConfig(t *testing.T) {
+	uc := &UpdateConfig{}
+	f, err := uc.File("")
+	if err != nil {
+		t.Error("unexpected error getting file from empty UpdateConfig")
+	}
+	if f != nil {
+		t.Errorf("getting file from empty UpdateConfig should have returned nil, got %v", f)
+	}
+	u, err := uc.Unit("")
+	if err != nil {
+		t.Error("unexpected error getting unit from empty UpdateConfig")
+	}
+	if u != nil {
+		t.Errorf("getting unit from empty UpdateConfig should have returned nil, got %v", u)
+	}
+}
+
+func TestInvalidUpdateOptions(t *testing.T) {
+	uon := &updateOption{
+		key:    "numbers",
+		prefix: "numero_",
+		valid:  []string{"one", "two"},
+	}
+	uoa := &updateOption{
+		key:    "any_will_do",
+		prefix: "any_",
+	}
+
+	if !uon.isValid("one") {
+		t.Error("update option did not accept valid option \"one\"")
+	}
+	if uon.isValid("three") {
+		t.Error("update option accepted invalid option \"three\"")
+	}
+	for _, s := range []string{"one", "asdf", "foobarbaz"} {
+		if !uoa.isValid(s) {
+			t.Errorf("update option with no \"valid\" field did not accept %q", s)
+		}
+	}
+
+	uc := &UpdateConfig{"reboot-strategy": "wizzlewazzle"}
+	f, err := uc.File("")
+	if err == nil {
+		t.Errorf("File did not give an error on invalid UpdateOption")
+	}
+	if f != nil {
+		t.Errorf("File did not return a nil file on invalid UpdateOption")
+	}
+}
+
+func TestServerGroupOptions(t *testing.T) {
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+	setupFixtures(dir)
+	u := &UpdateConfig{"group": "master", "server": "http://foo.com"}
+
+	want := `
+GROUP=master
+SERVER=http://foo.com`
+
+	f, err := u.File(dir)
+	if err != nil {
+		t.Errorf("unexpected error getting file from UpdateConfig: %v", err)
+	} else if f == nil {
+		t.Error("unexpectedly got empty file from UpdateConfig")
+	} else {
+		out := strings.Split(f.Content, "\n")
+		sort.Strings(out)
+		got := strings.Join(out, "\n")
+		if got != want {
+			t.Errorf("File has incorrect contents, got %v, want %v", got, want)
+		}
+	}
+}
+
+func TestRebootStrategies(t *testing.T) {
+	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
+	if err != nil {
+		t.Fatalf("Unable to create tempdir: %v", err)
+	}
+	defer os.RemoveAll(dir)
+	setupFixtures(dir)
+	strategies := []struct {
+		name     string
+		line     string
+		uMask    bool
+		uCommand string
+	}{
+		{"best-effort", "REBOOT_STRATEGY=best-effort", false, "restart"},
+		{"etcd-lock", "REBOOT_STRATEGY=etcd-lock", false, "restart"},
+		{"reboot", "REBOOT_STRATEGY=reboot", false, "restart"},
+		{"off", "REBOOT_STRATEGY=off", true, "stop"},
+	}
+	for _, s := range strategies {
+		uc := &UpdateConfig{"reboot-strategy": s.name}
+		f, err := uc.File(dir)
+		if err != nil {
+			t.Errorf("update failed to generate file for reboot-strategy=%v: %v", s.name, err)
+		} else if f == nil {
+			t.Errorf("generated empty file for reboot-strategy=%v", s.name)
+		} else {
+			seen := false
+			for _, line := range strings.Split(f.Content, "\n") {
+				if line == s.line {
+					seen = true
+					break
+				}
+			}
+			if !seen {
+				t.Errorf("couldn't find expected line %v for reboot-strategy=%v", s.line)
+			}
+		}
+		u, err := uc.Unit(dir)
+		if err != nil {
+			t.Errorf("failed to generate unit for reboot-strategy=%v!", s.name)
+		} else if u == nil {
+			t.Errorf("generated empty unit for reboot-strategy=%v", s.name)
+		} else {
+			if u.Name != locksmithUnit {
+				t.Errorf("unit generated for reboot strategy=%v had bad name: %v", s.name, u.Name)
+			}
+			if u.Mask != s.uMask {
+				t.Errorf("unit generated for reboot strategy=%v had bad mask: %t", s.name, u.Mask)
+			}
+			if u.Command != s.uCommand {
+				t.Errorf("unit generated for reboot strategy=%v had bad command: %v", s.name, u.Command)
+			}
+		}
+	}
+
+}
+
+func TestUpdateConfWrittenToDisk(t *testing.T) {
 	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
 	if err != nil {
 		t.Fatalf("Unable to create tempdir: %v", err)
@@ -49,9 +185,8 @@ func TestLocksmithEnvironmentWrittenToDisk(t *testing.T) {
 		f, err := uc.File(dir)
 		if err != nil {
 			t.Fatalf("Processing UpdateConfig failed: %v", err)
-		}
+		} else if f == nil {
-		if f == nil {
+			t.Fatal("Unexpectedly got nil updateconfig file")
-			t.Fatalf("UpdateConfig generated nil file unexpectedly")
 		}
 
 		f.Path = path.Join(dir, f.Path)
@@ -80,32 +215,3 @@ func TestLocksmithEnvironmentWrittenToDisk(t *testing.T) {
 		}
 	}
 }
-func TestLocksmithEnvironmentMasked(t *testing.T) {
-	dir, err := ioutil.TempDir(os.TempDir(), "coreos-cloudinit-")
-	if err != nil {
-		t.Fatalf("Unable to create tempdir: %v", err)
-	}
-	defer os.RemoveAll(dir)
-	setupFixtures(dir)
-
-	uc := &UpdateConfig{"reboot-strategy": "off"}
-
-	u, err := uc.Unit(dir)
-	if err != nil {
-		t.Fatalf("Processing UpdateConfig failed: %v", err)
-	}
-	if u == nil {
-		t.Fatalf("UpdateConfig generated nil unit unexpectedly")
-	}
-
-	system.MaskUnit(u.Name, dir)
-
-	fullPath := path.Join(dir, "etc", "systemd", "system", "locksmithd.service")
-	target, err := os.Readlink(fullPath)
-	if err != nil {
-		t.Fatalf("Unable to read link %v", err)
-	}
-	if target != "/dev/null" {
-		t.Fatalf("Locksmith not masked, unit target %v", target)
-	}
-}