policy: close memoized policy filesystems

build/policy_loader.go

@@ -179,10 +179,10 @@ func (p *policyPathFS) Stat(name string) (fs.FileInfo, error) {
 }
 
 func (p *policyPathFS) Close() error {
-	if err := p.cwdFS.close(); err != nil {
+	if err := p.cwdFS.closeAll(); err != nil {
 		return err
 	}
-	return p.contextFS.close()
+	return p.contextFS.closeAll()
 }
 
 func (p *policyPathFS) resolve(name string) (fs.StatFS, string, error) {
@@ -274,6 +274,22 @@ func (m *memoizedPolicyFS) close() error {
 	return nil
 }
 
+func (m *memoizedPolicyFS) closeAll() error {
+	m.mu.Lock()
+	closeFn := m.closeFn
+	m.fs = nil
+	m.closeFn = nil
+	m.err = nil
+	m.loaded = false
+	m.refs = 0
+	m.mu.Unlock()
+
+	if closeFn != nil {
+		return closeFn()
+	}
+	return nil
+}
+
 func normalizeRemotePolicyPath(raw string) (string, error) {
 	clean := strings.TrimPrefix(path.Join("/", filepath.ToSlash(raw)), "/")
 	if clean == "." || clean == "" {

build/policy_loader_test.go

@@ -83,3 +83,40 @@ func TestMemoizedPolicyFSReinitializesAfterAllRefsClosed(t *testing.T) {
 	require.NoError(t, m.close())
 	require.Equal(t, 2, closeCalls)
 }
+
+func TestMemoizedPolicyFSCloseAll(t *testing.T) {
+	var initCalls int
+	var closeCalls int
+
+	m := &memoizedPolicyFS{
+		init: func() (fs.StatFS, func() error, error) {
+			initCalls++
+			root := fstest.MapFS{
+				"policy.rego": &fstest.MapFile{Data: []byte("package docker\n")},
+			}
+			return root, func() error {
+				closeCalls++
+				return nil
+			}, nil
+		},
+	}
+
+	first, err := m.get()
+	require.NoError(t, err)
+	require.NotNil(t, first)
+	second, err := m.get()
+	require.NoError(t, err)
+	require.NotNil(t, second)
+	require.Equal(t, 1, initCalls)
+
+	require.NoError(t, m.closeAll())
+	require.Equal(t, 1, closeCalls)
+
+	third, err := m.get()
+	require.NoError(t, err)
+	require.NotNil(t, third)
+	require.Equal(t, 2, initCalls)
+
+	require.NoError(t, m.closeAll())
+	require.Equal(t, 2, closeCalls)
+}