Merge pull request #6893 from justinsb/fake_atomic_write

Avoid concurrent write corruption to /etc/hosts

protokube/pkg/gossip/dns/hosts/BUILD.bazel

@@ -1,4 +1,4 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "go_default_library",
@@ -7,3 +7,10 @@ go_library(
     visibility = ["//visibility:public"],
     deps = ["//vendor/k8s.io/klog:go_default_library"],
 )
+
+go_test(
+    name = "go_default_test",
+    srcs = ["hosts_test.go"],
+    embed = [":go_default_library"],
+    deps = ["//pkg/diff:go_default_library"],
+)

protokube/pkg/gossip/dns/hosts/hosts.go

@@ -17,12 +17,16 @@ limitations under the License.
 package hosts
 
 import (
+	"bytes"
 	"fmt"
 	"io/ioutil"
+	math_rand "math/rand"
 	"os"
 	"path/filepath"
 	"sort"
 	"strings"
+	"sync"
+	"time"
 
 	"k8s.io/klog"
 )
@@ -32,7 +36,13 @@ const (
 	GUARD_END   = "# End host entries managed by kops"
 )
 
+var hostsFileMutex sync.Mutex
+
 func UpdateHostsFileWithRecords(p string, addrToHosts map[string][]string) error {
+	// For safety / sanity, we avoid concurrent updates from one process
+	hostsFileMutex.Lock()
+	defer hostsFileMutex.Unlock()
+
 	stat, err := os.Stat(p)
 	if err != nil {
 		return fmt.Errorf("error getting file status of %q: %v", p, err)
@@ -44,19 +54,28 @@ func UpdateHostsFileWithRecords(p string, addrToHosts map[string][]string) error
 	}
 
 	var out []string
-	depth := 0
+	inGuardBlock := false
 	for _, line := range strings.Split(string(data), "\n") {
 		k := strings.TrimSpace(line)
 		if k == GUARD_BEGIN {
-			depth++
+			if inGuardBlock {
+				klog.Warningf("/etc/hosts guard-block begin seen while in guard block; will ignore")
+			}
+			inGuardBlock = true
 		}
 
-		if depth <= 0 {
+		if !inGuardBlock {
 			out = append(out, line)
 		}
 
 		if k == GUARD_END {
-			depth--
+			if !inGuardBlock {
+				klog.Warningf("/etc/hosts guard-block end seen before guard-block start; will ignore end")
+				// Don't output the line
+				out = out[:len(out)-1]
+			}
+
+			inGuardBlock = false
 		}
 	}
 
@@ -74,25 +93,72 @@ func UpdateHostsFileWithRecords(p string, addrToHosts map[string][]string) error
 	}
 	out = append(out, "")
 
-	out = append(out, GUARD_BEGIN)
+	var block []string
 	for addr, hosts := range addrToHosts {
 		sort.Strings(hosts)
-		out = append(out, addr+"\t"+strings.Join(hosts, " "))
+		block = append(block, addr+"\t"+strings.Join(hosts, " "))
 	}
+	// Sort into a consistent order to minimize updates
+	sort.Strings(block)
+
+	out = append(out, GUARD_BEGIN)
+	out = append(out, block...)
 	out = append(out, GUARD_END)
 	out = append(out, "")
 
+	updated := []byte(strings.Join(out, "\n"))
+
+	if bytes.Equal(updated, data) {
+		klog.V(2).Infof("skipping update of unchanged /etc/hosts")
+		return nil
+	}
+
 	// Note that because we are bind mounting /etc/hosts, we can't do a normal atomic file write
 	// (where we write a temp file and rename it)
-	// TODO: We should just hold the file open while we read & write it
-	err = ioutil.WriteFile(p, []byte(strings.Join(out, "\n")), stat.Mode().Perm())
-	if err != nil {
+	if err := pseudoAtomicWrite(p, updated, stat.Mode()); err != nil {
 		return fmt.Errorf("error writing file %q: %v", p, err)
 	}
 
 	return nil
 }
 
+// Because we are bind-mounting /etc/hosts, we can't do a normal
+// atomic file write (where we write a temp file and rename it);
+// instead we write the file, pause, re-read and see if anyone else
+// wrote in the meantime; if so we rewrite again.  By pausing for a
+// random amount of time, eventually we'll win the write race and
+// exit.  This doesn't guarantee fairness, but it should mean that the
+// end-result is not malformed (i.e. partial writes).
+func pseudoAtomicWrite(p string, b []byte, mode os.FileMode) error {
+	attempt := 0
+	for {
+		attempt++
+		if attempt > 10 {
+			return fmt.Errorf("failed to consistently write file %q - too many retries", p)
+		}
+
+		if err := ioutil.WriteFile(p, b, mode); err != nil {
+			klog.Warningf("error writing file %q: %v", p, err)
+			continue
+		}
+
+		n := 1 + math_rand.Intn(20)
+		time.Sleep(time.Duration(n) * time.Millisecond)
+
+		contents, err := ioutil.ReadFile(p)
+		if err != nil {
+			klog.Warningf("error re-reading file %q: %v", p, err)
+			continue
+		}
+
+		if bytes.Equal(contents, b) {
+			return nil
+		}
+
+		klog.Warningf("detected concurrent write to file %q, will retry", p)
+	}
+}
+
 func atomicWriteFile(filename string, data []byte, perm os.FileMode) error {
 	dir := filepath.Dir(filename)
 

protokube/pkg/gossip/dns/hosts/hosts_test.go

@@ -0,0 +1,159 @@
+/*
+Copyright 2019 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package hosts
+
+import (
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"k8s.io/kops/pkg/diff"
+)
+
+func TestRemovesDuplicateGuardedBlocks(t *testing.T) {
+	in := `
+foo 10.2.3.4
+
+# Begin host entries managed by etcd-manager[etcd] - do not edit
+# End host entries managed by etcd-manager[etcd]
+# Begin host entries managed by etcd-manager[etcd] - do not edit
+# End host entries managed by etcd-manager[etcd]
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+`
+
+	expected := `
+foo 10.2.3.4
+
+# Begin host entries managed by etcd-manager[etcd] - do not edit
+# End host entries managed by etcd-manager[etcd]
+# Begin host entries managed by etcd-manager[etcd] - do not edit
+# End host entries managed by etcd-manager[etcd]
+
+# Begin host entries managed by kops - do not edit
+a\t10.0.1.1 10.0.1.2
+b\t10.0.2.1
+c\t
+# End host entries managed by kops
+`
+
+	runTest(t, in, expected)
+}
+
+func TestRecoversFromBadNesting(t *testing.T) {
+	in := `
+foo 10.2.3.4
+
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+# End host entries managed by kops
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# Begin host entries managed by kops - do not edit
+# Begin host entries managed by kops - do not edit
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+# Begin host entries managed by kops - do not edit
+# End host entries managed by kops
+
+bar 10.1.2.3
+`
+
+	expected := `
+foo 10.2.3.4
+
+
+bar 10.1.2.3
+
+# Begin host entries managed by kops - do not edit
+a\t10.0.1.1 10.0.1.2
+b\t10.0.2.1
+c\t
+# End host entries managed by kops
+`
+
+	runTest(t, in, expected)
+}
+
+func runTest(t *testing.T, in string, expected string) {
+	expected = strings.Replace(expected, "\\t", "\t", -1)
+
+	dir, err := ioutil.TempDir("", "")
+	if err != nil {
+		t.Fatalf("error creating temp dir: %v", err)
+	}
+	defer func() {
+		err := os.RemoveAll(dir)
+		if err != nil {
+			t.Errorf("failed to remove temp dir %q: %v", dir, err)
+		}
+	}()
+
+	p := filepath.Join(dir, "hosts")
+	addrToHosts := map[string][]string{
+		"a": {"10.0.1.2", "10.0.1.1"},
+		"b": {"10.0.2.1"},
+		"c": {},
+	}
+
+	if err := ioutil.WriteFile(p, []byte(in), 0755); err != nil {
+		t.Fatalf("error writing hosts file: %v", err)
+	}
+
+	// We run it repeatedly to make sure we don't change it accidentally
+	for i := 0; i < 100; i++ {
+		if err := UpdateHostsFileWithRecords(p, addrToHosts); err != nil {
+			t.Fatalf("error updating hosts file: %v", err)
+		}
+
+		b, err := ioutil.ReadFile(p)
+		if err != nil {
+			t.Fatalf("error reading output file: %v", err)
+		}
+
+		actual := string(b)
+		if actual != expected {
+			diffString := diff.FormatDiff(expected, actual)
+			t.Logf("diff:\n%s\n", diffString)
+			t.Errorf("unexpected output.  expected=%q, actual=%q", expected, actual)
+		}
+	}
+}