Ensure that host file parsing doesn't drop entries when parsing. (#1261)

Co-authored-by: Jakub Novák <[email protected]>

Author: djeebus

packages/envd/internal/api/init.go

@@ -7,6 +7,8 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"net/netip"
+	"os"
 	"time"
 
 	"github.com/rs/zerolog"
@@ -15,8 +17,11 @@ import (
 
 	"github.com/e2b-dev/infra/packages/envd/internal/host"
 	"github.com/e2b-dev/infra/packages/envd/internal/logs"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
+const hostsFilePermissions = 0o644
+
 var ErrAccessTokenAlreadySet = errors.New("access token is already set")
 
 const (
@@ -119,19 +124,68 @@ func (a *API) SetupHyperloop(address string) {
 	a.hyperloopLock.Lock()
 	defer a.hyperloopLock.Unlock()
 
-	hosts, err := txeh.NewHosts(&txeh.HostsConfig{ReadFilePath: "/etc/hosts", WriteFilePath: "/etc/hosts"})
+	if err := rewriteHostsFile(address, "/etc/hosts"); err != nil {
+		a.logger.Error().Err(err).Msg("failed to modify hosts file")
+	} else {
+		a.envVars.Store("E2B_EVENTS_ADDRESS", fmt.Sprintf("http://%s", address))
+	}
+}
+
+const eventsHost = "events.e2b.local"
+
+func rewriteHostsFile(address, path string) error {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return fmt.Errorf("failed to read hosts file: %w", err)
+	}
+
+	// the txeh library drops an entry if the file does not end with a newline
+	if len(data) > 0 && data[len(data)-1] != '\n' {
+		data = append(data, '\n')
+	}
+
+	hosts, err := txeh.NewHosts(&txeh.HostsConfig{RawText: utils.ToPtr(string(data))})
 	if err != nil {
-		a.logger.Error().Msgf("Failed to create hosts: %v", err)
-		return
+		return fmt.Errorf("failed to create hosts: %w", err)
 	}
 
 	// Update /etc/hosts to point events.e2b.local to the hyperloop IP
 	// This will remove any existing entries for events.e2b.local first
-	hosts.AddHost(address, "events.e2b.local")
-	err = hosts.Save()
+	ipFamily, err := getIPFamily(address)
+	if err != nil {
+		return fmt.Errorf("failed to get ip family: %w", err)
+	}
+
+	if ok, current, _ := hosts.HostAddressLookup(eventsHost, ipFamily); ok && current == address {
+		return nil // nothing to be done
+	}
+
+	hosts.AddHost(address, eventsHost)
+
+	if err = os.WriteFile(path, []byte(hosts.RenderHostsFile()), hostsFilePermissions); err != nil {
+		return fmt.Errorf("failed to save hosts file: %w", err)
+	}
+
+	return nil
+}
+
+var (
+	ErrInvalidAddress       = errors.New("invalid IP address")
+	ErrUnknownAddressFormat = errors.New("unknown IP address format")
+)
+
+func getIPFamily(address string) (txeh.IPFamily, error) {
+	addressIP, err := netip.ParseAddr(address)
 	if err != nil {
-		a.logger.Error().Msgf("Failed to add events host entry: %v", err)
+		return txeh.IPFamilyV4, fmt.Errorf("failed to parse IP address: %w", err)
 	}
 
-	a.envVars.Store("E2B_EVENTS_ADDRESS", fmt.Sprintf("http://%s", address))
+	switch {
+	case addressIP.Is4():
+		return txeh.IPFamilyV4, nil
+	case addressIP.Is6():
+		return txeh.IPFamilyV6, nil
+	default:
+		return txeh.IPFamilyV4, fmt.Errorf("%w: %s", ErrUnknownAddressFormat, address)
+	}
 }

packages/envd/internal/api/init_test.go

@@ -0,0 +1,47 @@
+package api
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSimpleCases(t *testing.T) {
+	testCases := map[string]func(string) string{
+		"both newlines":               func(s string) string { return s },
+		"no newline prefix":           func(s string) string { return strings.TrimPrefix(s, "\n") },
+		"no newline suffix":           func(s string) string { return strings.TrimSuffix(s, "\n") },
+		"no newline prefix or suffix": strings.TrimSpace,
+	}
+
+	for name, preprocessor := range testCases {
+		t.Run(name, func(t *testing.T) {
+			tempDir := t.TempDir()
+
+			value := `
+# comment
+127.0.0.1        one.host
+127.0.0.2        two.host
+`
+			value = preprocessor(value)
+			inputPath := filepath.Join(tempDir, "hosts")
+			err := os.WriteFile(inputPath, []byte(value), hostsFilePermissions)
+			require.NoError(t, err)
+
+			err = rewriteHostsFile("127.0.0.3", inputPath)
+			require.NoError(t, err)
+
+			data, err := os.ReadFile(inputPath)
+			require.NoError(t, err)
+
+			assert.Equal(t, `# comment
+127.0.0.1        one.host
+127.0.0.2        two.host
+127.0.0.3        events.e2b.local`, strings.TrimSpace(string(data)))
+		})
+	}
+}

packages/envd/main.go

@@ -38,7 +38,7 @@ const (
 )
 
 var (
-	Version = "0.3.5"
+	Version = "0.3.6"
 
 	commitSHA string