From 6b3827b63586364131ef2289b6ba47e6ec916d71 Mon Sep 17 00:00:00 2001
From: Vadym Popov <vadym.popov@goteleport.com>
Date: Tue, 15 Oct 2024 11:06:50 -0400
Subject: [PATCH] [v14] Add filesystem lock for unix and windows platforms
 (#47196)

* Add filesystem lock for unix and windows platforms

* Add non-blocking flag

* Replace lock with gofrs/flock
---
 lib/utils/fs.go      | 22 ++++++++++++++++++-
 lib/utils/fs_test.go | 52 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/lib/utils/fs.go b/lib/utils/fs.go
index 36322877f83b3..49254628f0ac4 100644
--- a/lib/utils/fs.go
+++ b/lib/utils/fs.go
@@ -201,6 +201,16 @@ func FSTryWriteLock(filePath string) (unlock func() error, err error) {
 	return fileLock.Unlock, nil
 }
 
+// FSWriteLock tries to grab write lock and block if lock is already acquired by someone else.
+func FSWriteLock(filePath string) (unlock func() error, err error) {
+	fileLock := flock.New(getPlatformLockFilePath(filePath))
+	if err := fileLock.Lock(); err != nil {
+		return nil, trace.ConvertSystemError(err)
+	}
+
+	return fileLock.Unlock, nil
+}
+
 // FSTryWriteLockTimeout tries to grab write lock, it's doing it until locks is acquired, or timeout is expired,
 // or context is expired.
 func FSTryWriteLockTimeout(ctx context.Context, filePath string, timeout time.Duration) (unlock func() error, err error) {
@@ -214,7 +224,7 @@ func FSTryWriteLockTimeout(ctx context.Context, filePath string, timeout time.Du
 	return fileLock.Unlock, nil
 }
 
-// FSTryReadLock tries to grab write lock, returns ErrUnsuccessfulLockTry
+// FSTryReadLock tries to grab shared lock, returns ErrUnsuccessfulLockTry
 // if lock is already acquired by someone else
 func FSTryReadLock(filePath string) (unlock func() error, err error) {
 	fileLock := flock.New(getPlatformLockFilePath(filePath))
@@ -229,6 +239,16 @@ func FSTryReadLock(filePath string) (unlock func() error, err error) {
 	return fileLock.Unlock, nil
 }
 
+// FSReadLock tries to grab shared lock and block if lock is already acquired by someone else.
+func FSReadLock(filePath string) (unlock func() error, err error) {
+	fileLock := flock.New(getPlatformLockFilePath(filePath))
+	if err := fileLock.RLock(); err != nil {
+		return nil, trace.ConvertSystemError(err)
+	}
+
+	return fileLock.Unlock, nil
+}
+
 // FSTryReadLockTimeout tries to grab read lock, it's doing it until locks is acquired, or timeout is expired,
 // or context is expired.
 func FSTryReadLockTimeout(ctx context.Context, filePath string, timeout time.Duration) (unlock func() error, err error) {
diff --git a/lib/utils/fs_test.go b/lib/utils/fs_test.go
index dc352df644fff..e4d192f3d2d79 100644
--- a/lib/utils/fs_test.go
+++ b/lib/utils/fs_test.go
@@ -18,9 +18,11 @@ package utils
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"path/filepath"
 	"runtime"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -326,6 +328,56 @@ func TestLocks(t *testing.T) {
 	require.NoError(t, unlock())
 }
 
+// TestLockWithBlocking verifies that second lock call is blocked until first is released.
+func TestLockWithBlocking(t *testing.T) {
+	var locked atomic.Bool
+
+	lockFile := filepath.Join(os.TempDir(), ".lock")
+	t.Cleanup(func() {
+		require.NoError(t, os.Remove(lockFile))
+	})
+
+	// Acquire first lock should not return any error.
+	unlock, err := FSWriteLock(lockFile)
+	require.NoError(t, err)
+	locked.Store(true)
+
+	signal := make(chan struct{})
+	errChan := make(chan error)
+	go func() {
+		signal <- struct{}{}
+		unlock, err := FSWriteLock(lockFile)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		if locked.Load() {
+			errChan <- fmt.Errorf("first lock is still acquired, second lock must be blocking")
+			return
+		}
+		if err := unlock(); err != nil {
+			errChan <- err
+			return
+		}
+		signal <- struct{}{}
+	}()
+
+	<-signal
+	// We have to wait till next lock is reached to ensure we block execution of goroutine.
+	// Since this is system call we can't track if the function reach blocking state already.
+	time.Sleep(100 * time.Millisecond)
+	locked.Store(false)
+	require.NoError(t, unlock())
+
+	select {
+	case err := <-errChan:
+		require.NoError(t, err)
+	case <-signal:
+	case <-time.After(5 * time.Second):
+		require.Fail(t, "second lock is not released")
+	}
+}
+
 func TestOverwriteFile(t *testing.T) {
 	have := []byte("Sensitive Information")
 	fName := filepath.Join(t.TempDir(), "teleport-overwrite-file-test")