Add method AtomicWritePrivateFile in new diskutil package (#187)

Signed-off-by: Max Lambrecht <maxlambrecht@gmail.com>

pkg/common/diskutil/disk.go

@@ -0,0 +1,63 @@
+package diskutil
+
+import (
+	"os"
+	"path/filepath"
+)
+
+// Define the file mode for private files.
+const (
+	fileModePrivate = 0600
+)
+
+// AtomicWritePrivateFile writes data to a file atomically.
+// The file is created with private permissions (0600).
+func AtomicWritePrivateFile(path string, data []byte) error {
+	return atomicWrite(path, data, fileModePrivate)
+}
+
+func atomicWrite(path string, data []byte, mode os.FileMode) error {
+	tmpPath := path + ".tmp"
+
+	// Attempt to write to temporary file
+	if err := write(tmpPath, data, mode); err != nil {
+		return err
+	}
+
+	// If write is successful, rename temp file to actual path
+	return rename(tmpPath, path)
+}
+
+func write(path string, data []byte, mode os.FileMode) error {
+	// Attempt to open file, deferring close until function finishes
+	file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, mode)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	_, err = file.Write(data)
+	if err != nil {
+		return err
+	}
+
+	// Sync file contents to disk
+	return file.Sync()
+}
+
+func rename(oldPath, newPath string) error {
+	// Attempt to rename file
+	if err := os.Rename(oldPath, newPath); err != nil {
+		return err
+	}
+
+	// Open containing directory and defer close until function finishes
+	dir, err := os.Open(filepath.Dir(newPath))
+	if err != nil {
+		return err
+	}
+	defer dir.Close()
+
+	// Sync directory changes to disk
+	return dir.Sync()
+}

pkg/common/diskutil/disk_test.go

@@ -0,0 +1,61 @@
+package diskutil
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestWriteFile(t *testing.T) {
+	tempDir, err := os.MkdirTemp("", "galadriel-test")
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		require.NoError(t, os.RemoveAll(tempDir))
+	})
+
+	tests := []struct {
+		name            string
+		data            []byte
+		atomicWriteFunc func(string, []byte) error
+		expectMode      os.FileMode
+	}{
+		{
+			name:            "basic - AtomicWritePrivateFile",
+			data:            []byte("Hello, World"),
+			atomicWriteFunc: AtomicWritePrivateFile,
+			expectMode:      0600,
+		},
+		{
+			name:            "empty - AtomicWritePrivateFile",
+			data:            []byte{},
+			atomicWriteFunc: AtomicWritePrivateFile,
+			expectMode:      0600,
+		},
+		{
+			name:            "binary - AtomicWritePrivateFile",
+			data:            []byte{0xFF, 0, 0xFF, 0x3D, 0xD8, 0xA9, 0xDC, 0xF0, 0x9F, 0x92, 0xA9},
+			atomicWriteFunc: AtomicWritePrivateFile,
+			expectMode:      0600,
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			file := filepath.Join(tempDir, "file")
+			err := tt.atomicWriteFunc(file, tt.data)
+			require.NoError(t, err)
+
+			info, err := os.Stat(file)
+			require.NoError(t, err)
+			require.EqualValues(t, tt.expectMode, info.Mode())
+
+			content, err := os.ReadFile(file)
+			require.NoError(t, err)
+			require.Equal(t, tt.data, content)
+
+			require.NoError(t, os.Remove(file))
+		})
+	}
+}

pkg/common/keymanager/disk.go

@@ -13,10 +13,9 @@ import (
 	"os"
 
 	"github.com/HewlettPackard/galadriel/pkg/common/cryptoutil"
+	"github.com/HewlettPackard/galadriel/pkg/common/diskutil"
 )
 
-const keyFilePerm = 0600
-
 // Disk extends the base KeyManager to store keys in disk.
 type Disk struct {
 	base
@@ -114,7 +113,7 @@ func (d *Disk) saveKeysToDisk() error {
 		return fmt.Errorf("failed to serialize keys: %w", err)
 	}
 
-	if err := os.WriteFile(d.keysFilePath, data, keyFilePerm); err != nil {
+	if err := diskutil.AtomicWritePrivateFile(d.keysFilePath, data); err != nil {
 		return fmt.Errorf("failed to write keys to disk: %w", err)
 	}
 

pkg/harvester/galadrielclient/client.go

@@ -18,6 +18,7 @@ import (
 
 	"github.com/HewlettPackard/galadriel/pkg/common/api"
 	"github.com/HewlettPackard/galadriel/pkg/common/constants"
+	"github.com/HewlettPackard/galadriel/pkg/common/diskutil"
 	"github.com/HewlettPackard/galadriel/pkg/common/entity"
 	"github.com/HewlettPackard/galadriel/pkg/common/util"
 	"github.com/HewlettPackard/galadriel/pkg/server/api/harvester"
@@ -30,7 +31,6 @@ const (
 	jwtRotationInterval = 5 * time.Minute
 	onboardPath         = "/trust-domain/onboard"
 	tokenFile           = "jwt-token"
-	tokenFileMode       = 0600
 )
 
 var (
@@ -547,8 +547,7 @@ func (p *jwtStore) setToken(jwt string) {
 	p.mu.Unlock()
 
 	// Save the token to disk
-	err := p.saveToken()
-	if err != nil {
+	if err := p.saveToken(); err != nil {
 		p.logger.Errorf("Failed to save JWT token to disk: %v", err)
 	}
 }
@@ -581,5 +580,5 @@ func (p *jwtStore) saveToken() error {
 	p.mu.RLock()
 	defer p.mu.RUnlock()
 
-	return os.WriteFile(p.tokenFilePath, []byte(p.jwt), tokenFileMode)
+	return diskutil.AtomicWritePrivateFile(p.tokenFilePath, []byte(p.jwt))
 }