From 96f24ca7bb57206cf0da7370e508b91dbb582d71 Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Fri, 3 Aug 2018 14:01:54 -0700
Subject: [PATCH] executor: improve hosts cleanup

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 executor/containerdexecutor/executor.go |  7 +++----
 executor/oci/hosts.go                   | 26 +++++++++++++------------
 executor/runcexecutor/executor.go       |  6 +++---
 3 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/executor/containerdexecutor/executor.go b/executor/containerdexecutor/executor.go
index f271f036..cb076686 100644
--- a/executor/containerdexecutor/executor.go
+++ b/executor/containerdexecutor/executor.go
@@ -3,7 +3,6 @@ package containerdexecutor
 import (
 	"context"
 	"io"
-	"os"
 	"syscall"
 	"time"
 
@@ -40,12 +39,12 @@ func (w containerdExecutor) Exec(ctx context.Context, meta executor.Meta, root c
 		return err
 	}
 
-	hostsFile, err := oci.GetHostsFile(ctx, w.root, meta.ExtraHosts)
+	hostsFile, clean, err := oci.GetHostsFile(ctx, w.root, meta.ExtraHosts)
 	if err != nil {
 		return err
 	}
-	if len(meta.ExtraHosts) > 0 {
-		defer os.RemoveAll(hostsFile)
+	if clean != nil {
+		defer clean()
 	}
 
 	mountable, err := root.Mount(ctx, false)
diff --git a/executor/oci/hosts.go b/executor/oci/hosts.go
index d01b9e6e..405d758d 100644
--- a/executor/oci/hosts.go
+++ b/executor/oci/hosts.go
@@ -17,51 +17,53 @@ const hostsContent = `
 ::1	localhost ip6-localhost ip6-loopback
 `
 
-func GetHostsFile(ctx context.Context, stateDir string, extraHosts []executor.HostIP) (string, error) {
+func GetHostsFile(ctx context.Context, stateDir string, extraHosts []executor.HostIP) (string, func(), error) {
 	if len(extraHosts) == 0 {
 		_, err := g.Do(ctx, stateDir, func(ctx context.Context) (interface{}, error) {
-			_, err := makeHostsFile(stateDir, nil)
+			_, _, err := makeHostsFile(stateDir, nil)
 			return nil, err
 		})
 		if err != nil {
-			return "", err
+			return "", nil, err
 		}
-		return filepath.Join(stateDir, "hosts"), nil
+		return filepath.Join(stateDir, "hosts"), func() {}, nil
 	}
 	return makeHostsFile(stateDir, extraHosts)
 }
 
-func makeHostsFile(stateDir string, extraHosts []executor.HostIP) (string, error) {
+func makeHostsFile(stateDir string, extraHosts []executor.HostIP) (string, func(), error) {
 	p := filepath.Join(stateDir, "hosts")
 	if len(extraHosts) != 0 {
 		p += "." + identity.NewID()
 	}
 	_, err := os.Stat(p)
 	if err == nil {
-		return "", nil
+		return "", func() {}, nil
 	}
 	if !os.IsNotExist(err) {
-		return "", err
+		return "", nil, err
 	}
 
 	b := &bytes.Buffer{}
 
 	if _, err := b.Write([]byte(hostsContent)); err != nil {
-		return "", err
+		return "", nil, err
 	}
 
 	for _, h := range extraHosts {
 		if _, err := b.Write([]byte(fmt.Sprintf("%s\t%s\n", h.IP.String(), h.Host))); err != nil {
-			return "", err
+			return "", nil, err
 		}
 	}
 
 	if err := ioutil.WriteFile(p+".tmp", b.Bytes(), 0644); err != nil {
-		return "", err
+		return "", nil, err
 	}
 
 	if err := os.Rename(p+".tmp", p); err != nil {
-		return "", err
+		return "", nil, err
 	}
-	return p, nil
+	return p, func() {
+		os.RemoveAll(p)
+	}, nil
 }
diff --git a/executor/runcexecutor/executor.go b/executor/runcexecutor/executor.go
index 75bce875..123a02dc 100644
--- a/executor/runcexecutor/executor.go
+++ b/executor/runcexecutor/executor.go
@@ -103,12 +103,12 @@ func (w *runcExecutor) Exec(ctx context.Context, meta executor.Meta, root cache.
 		return err
 	}
 
-	hostsFile, err := oci.GetHostsFile(ctx, w.root, meta.ExtraHosts)
+	hostsFile, clean, err := oci.GetHostsFile(ctx, w.root, meta.ExtraHosts)
 	if err != nil {
 		return err
 	}
-	if len(meta.ExtraHosts) > 0 {
-		defer os.RemoveAll(hostsFile)
+	if clean != nil {
+		defer clean()
 	}
 
 	mountable, err := root.Mount(ctx, false)