From 3548517cf2c6bc4e6ee5bcae6b276fb8b236164e Mon Sep 17 00:00:00 2001
From: Alex Suraci <suraci.alex@gmail.com>
Date: Wed, 26 Oct 2022 18:31:14 -0400
Subject: [PATCH] allow dynamic LocalDir access

Changes the filesync attachable to accept an interface instead of a
static allowlist of dirs. This way a single session can support syncing
directories not known ahead of time.

Signed-off-by: Alex Suraci <suraci.alex@gmail.com>
---
 client/solve.go                   |  8 ++++----
 session/filesync/filesync.go      | 28 ++++++++++++++++++----------
 session/filesync/filesync_test.go |  2 +-
 3 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/client/solve.go b/client/solve.go
index dbad6635c..71c51ba29 100644
--- a/client/solve.go
+++ b/client/solve.go
@@ -355,7 +355,7 @@ func (c *Client) solve(ctx context.Context, def *llb.Definition, runGateway runG
 	return res, nil
 }
 
-func prepareSyncedDirs(def *llb.Definition, localDirs map[string]string) ([]filesync.SyncedDir, error) {
+func prepareSyncedDirs(def *llb.Definition, localDirs map[string]string) (filesync.StaticDirSource, error) {
 	for _, d := range localDirs {
 		fi, err := os.Stat(d)
 		if err != nil {
@@ -371,10 +371,10 @@ func prepareSyncedDirs(def *llb.Definition, localDirs map[string]string) ([]file
 		return fsutil.MapResultKeep
 	}
 
-	dirs := make([]filesync.SyncedDir, 0, len(localDirs))
+	dirs := make(filesync.StaticDirSource, len(localDirs))
 	if def == nil {
 		for name, d := range localDirs {
-			dirs = append(dirs, filesync.SyncedDir{Name: name, Dir: d, Map: resetUIDAndGID})
+			dirs[name] = filesync.SyncedDir{Dir: d, Map: resetUIDAndGID}
 		}
 	} else {
 		for _, dt := range def.Def {
@@ -389,7 +389,7 @@ func prepareSyncedDirs(def *llb.Definition, localDirs map[string]string) ([]file
 					if !ok {
 						return nil, errors.Errorf("local directory %s not enabled", name)
 					}
-					dirs = append(dirs, filesync.SyncedDir{Name: name, Dir: d, Map: resetUIDAndGID})
+					dirs[name] = filesync.SyncedDir{Dir: d, Map: resetUIDAndGID}
 				}
 			}
 		}
diff --git a/session/filesync/filesync.go b/session/filesync/filesync.go
index be8cf9b8b..e31354262 100644
--- a/session/filesync/filesync.go
+++ b/session/filesync/filesync.go
@@ -27,27 +27,35 @@ const (
 )
 
 type fsSyncProvider struct {
-	dirs   map[string]SyncedDir
+	dirs   DirSource
 	p      progressCb
 	doneCh chan error
 }
 
 type SyncedDir struct {
-	Name     string
 	Dir      string
 	Excludes []string
 	Map      func(string, *fstypes.Stat) fsutil.MapResult
 }
 
+type DirSource interface {
+	LookupDir(string) (SyncedDir, bool)
+}
+
+type StaticDirSource map[string]SyncedDir
+
+var _ DirSource = StaticDirSource{}
+
+func (dirs StaticDirSource) LookupDir(name string) (SyncedDir, bool) {
+	dir, found := dirs[name]
+	return dir, found
+}
+
 // NewFSSyncProvider creates a new provider for sending files from client
-func NewFSSyncProvider(dirs []SyncedDir) session.Attachable {
-	p := &fsSyncProvider{
-		dirs: map[string]SyncedDir{},
+func NewFSSyncProvider(dirs DirSource) session.Attachable {
+	return &fsSyncProvider{
+		dirs: dirs,
 	}
-	for _, d := range dirs {
-		p.dirs[d.Name] = d
-	}
-	return p
 }
 
 func (sp *fsSyncProvider) Register(server *grpc.Server) {
@@ -81,7 +89,7 @@ func (sp *fsSyncProvider) handle(method string, stream grpc.ServerStream) (retEr
 		dirName = name[0]
 	}
 
-	dir, ok := sp.dirs[dirName]
+	dir, ok := sp.dirs.LookupDir(dirName)
 	if !ok {
 		return InvalidSessionError{status.Errorf(codes.NotFound, "no access allowed to dir %q", dirName)}
 	}
diff --git a/session/filesync/filesync_test.go b/session/filesync/filesync_test.go
index ad3150daf..424ffe31f 100644
--- a/session/filesync/filesync_test.go
+++ b/session/filesync/filesync_test.go
@@ -32,7 +32,7 @@ func TestFileSyncIncludePatterns(t *testing.T) {
 	m, err := session.NewManager()
 	require.NoError(t, err)
 
-	fs := NewFSSyncProvider([]SyncedDir{{Name: "test0", Dir: tmpDir}})
+	fs := NewFSSyncProvider(StaticDirSource{"test0": {Dir: tmpDir}})
 	s.Allow(fs)
 
 	dialer := session.Dialer(testutil.TestStream(testutil.Handler(m.HandleConn)))