Fixed recvfd leaks

pkg/networkservice/chains/nsmgr/unix_test.go

@@ -14,13 +14,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//+build !windows
+// +build linux
 
 package nsmgr_test
 
 import (
 	"context"
-	"runtime"
 	"testing"
 	"time"
 
@@ -84,9 +83,6 @@ func Test_Local_NoURLUsecase(t *testing.T) {
 }
 
 func Test_MultiForwarderSendfd(t *testing.T) {
-	if runtime.GOOS != "linux" {
-		t.Skip("sendfd works only on linux")
-	}
 	t.Cleanup(func() { goleak.VerifyNone(t) })
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)

pkg/networkservice/common/mechanisms/recvfd/common.go

@@ -86,7 +86,7 @@ func recvFDAndSwapInodeToFile(ctx context.Context, fileMap *perConnectionFileMap
 	return err
 }
 
-func swapFileToInode(fileMap *perConnectionFileMap, parameters map[string]string, closeAllFiles bool) error {
+func swapFileToInode(fileMap *perConnectionFileMap, parameters map[string]string) error {
 	// Get the inodeURL from  parameters
 	fileURLStr, ok := parameters[common.InodeURL]
 	if !ok {
@@ -112,11 +112,10 @@ func swapFileToInode(fileMap *perConnectionFileMap, parameters map[string]string
 		// Swap the fileURL for the inodeURL in parameters
 		parameters[common.InodeURL] = inodeURL.String()
 
-		// If closeAllFiles == true, close any files we may have open for any other inodes
 		// This is used to clean up files sent by MechanismPreferences that were *not* selected to be the
 		// connection mechanism
 		for inodeURLStr, file := range fileMap.filesByInodeURL {
-			if closeAllFiles || inodeURLStr != inodeURL.String() {
+			if inodeURLStr != inodeURL.String() {
 				delete(fileMap.filesByInodeURL, inodeURLStr)
 				_ = file.Close()
 			}

pkg/networkservice/common/mechanisms/recvfd/server.go

@@ -69,7 +69,7 @@ func (r *recvFDServer) Request(ctx context.Context, request *networkservice.Netw
 	}
 
 	// Swap back from File to Inode in the InodeURL in the Parameters
-	err = swapFileToInode(fileMap, conn.GetMechanism().GetParameters(), false)
+	err = swapFileToInode(fileMap, conn.GetMechanism().GetParameters())
 	if err != nil {
 		return nil, err
 	}
@@ -78,18 +78,20 @@ func (r *recvFDServer) Request(ctx context.Context, request *networkservice.Netw
 }
 
 func (r *recvFDServer) Close(ctx context.Context, conn *networkservice.Connection) (*empty.Empty, error) {
+	// Clean up the fileMap no matter what happens
+	defer r.closeFiles(conn)
+
 	// Get the grpcfd.FDRecver
 	recv, ok := grpcfd.FromContext(ctx)
 	if !ok {
 		return next.Server(ctx).Close(ctx, conn)
 	}
+
 	// Get the fileMap
 	fileMap, _ := r.fileMaps.LoadOrStore(conn.GetId(), &perConnectionFileMap{
 		filesByInodeURL:    make(map[string]*os.File),
 		inodeURLbyFilename: make(map[string]*url.URL),
 	})
-	// Clean up the fileMap no matter what happens
-	defer r.fileMaps.Delete(conn.GetId())
 
 	// Recv the FD and Swap the Inode for a file in InodeURL in Parameters
 	err := recvFDAndSwapInodeToFile(ctx, fileMap, conn.GetMechanism().GetParameters(), recv)
@@ -104,6 +106,20 @@ func (r *recvFDServer) Close(ctx context.Context, conn *networkservice.Connectio
 	}
 
 	// Swap back from File to Inode in the InodeURL in the Parameters
-	err = swapFileToInode(fileMap, conn.GetMechanism().GetParameters(), true)
+	err = swapFileToInode(fileMap, conn.GetMechanism().GetParameters())
 	return &empty.Empty{}, err
 }
+
+func (r *recvFDServer) closeFiles(conn *networkservice.Connection) {
+	defer r.fileMaps.Delete(conn.GetId())
+
+	fileMap, _ := r.fileMaps.LoadOrStore(conn.GetId(), &perConnectionFileMap{
+		filesByInodeURL:    make(map[string]*os.File),
+		inodeURLbyFilename: make(map[string]*url.URL),
+	})
+
+	for inodeURLStr, file := range fileMap.filesByInodeURL {
+		delete(fileMap.filesByInodeURL, inodeURLStr)
+		_ = file.Close()
+	}
+}

pkg/networkservice/common/mechanisms/recvfd/server_test.go

@@ -0,0 +1,216 @@
+// Copyright (c) 2021 Doc.ai and/or its affiliates.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at:
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// +build linux
+
+package recvfd_test
+
+import (
+	"context"
+	"fmt"
+	"net/url"
+	"os"
+	"path"
+	"runtime"
+	"testing"
+	"time"
+
+	"github.com/edwarnicke/grpcfd"
+	"github.com/networkservicemesh/api/pkg/api/networkservice"
+	"github.com/networkservicemesh/api/pkg/api/networkservice/mechanisms/cls"
+	"github.com/networkservicemesh/api/pkg/api/networkservice/mechanisms/common"
+	"github.com/networkservicemesh/api/pkg/api/networkservice/mechanisms/kernel"
+	"github.com/stretchr/testify/suite"
+	"go.uber.org/goleak"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+
+	"github.com/networkservicemesh/sdk/pkg/networkservice/chains/client"
+	"github.com/networkservicemesh/sdk/pkg/networkservice/common/mechanisms/recvfd"
+	"github.com/networkservicemesh/sdk/pkg/networkservice/common/mechanisms/sendfd"
+	"github.com/networkservicemesh/sdk/pkg/networkservice/core/chain"
+	"github.com/networkservicemesh/sdk/pkg/networkservice/utils/checks/checkcontext"
+	"github.com/networkservicemesh/sdk/pkg/tools/grpcfdutils"
+	"github.com/networkservicemesh/sdk/pkg/tools/grpcutils"
+	"github.com/networkservicemesh/sdk/pkg/tools/sandbox"
+)
+
+type checkRecvfdTestSuite struct {
+	suite.Suite
+
+	tempDir               string
+	onFileClosedContexts  []context.Context
+	onFileClosedCallbacks map[string]func()
+
+	testClient networkservice.NetworkServiceClient
+}
+
+func (s *checkRecvfdTestSuite) SetupTest() {
+	t := s.T()
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	t.Cleanup(func() {
+		cancel()
+		goleak.VerifyNone(t)
+	})
+
+	s.tempDir = t.TempDir()
+
+	sock, err := os.Create(path.Join(s.tempDir, "test.sock"))
+	s.Require().NoError(err)
+
+	serveURL := &url.URL{Scheme: "unix", Path: sock.Name()}
+
+	testChain := chain.NewNetworkServiceServer(
+		checkcontext.NewServer(t, func(t *testing.T, c context.Context) {
+			injectErr := grpcfdutils.InjectOnFileReceivedCallback(c, func(fileName string, file *os.File) {
+				runtime.SetFinalizer(file, func(file *os.File) {
+					onFileClosedCallback, ok := s.onFileClosedCallbacks[fileName]
+					if ok {
+						onFileClosedCallback()
+					}
+				})
+			})
+
+			s.Require().NoError(injectErr)
+		}),
+		recvfd.NewServer())
+
+	startServer(ctx, s, &testChain, serveURL)
+	s.testClient = createClient(ctx, serveURL)
+}
+
+func TestRecvfd(t *testing.T) {
+	suite.Run(t, new(checkRecvfdTestSuite))
+}
+
+func startServer(ctx context.Context, s *checkRecvfdTestSuite, testServerChain *networkservice.NetworkServiceServer, serveURL *url.URL) {
+	grpcServer := grpc.NewServer(grpc.Creds(grpcfd.TransportCredentials(insecure.NewCredentials())))
+	networkservice.RegisterNetworkServiceServer(grpcServer, *testServerChain)
+
+	errCh := grpcutils.ListenAndServe(ctx, serveURL, grpcServer)
+
+	s.Require().Len(errCh, 0)
+}
+
+func createClient(ctx context.Context, u *url.URL) networkservice.NetworkServiceClient {
+	return client.NewClient(
+		ctx,
+		client.WithClientURL(sandbox.CloneURL(u)),
+		client.WithDialOptions(grpc.WithTransportCredentials(
+			grpcfd.TransportCredentials(insecure.NewCredentials())),
+		),
+		client.WithDialTimeout(time.Second),
+		client.WithoutRefresh(),
+		client.WithAdditionalFunctionality(sendfd.NewClient()))
+}
+
+func createFile(s *checkRecvfdTestSuite, fileName string) (inodeURLStr string, fileClosedContext context.Context, cancelFunc func()) {
+	f, err := os.Create(fileName)
+	s.Require().NoErrorf(err, "Failed to create and open a file: %v", err)
+
+	err = f.Close()
+	s.Require().NoErrorf(err, "Failed to close file: %v", err)
+
+	fileClosedContext, cancelFunc = context.WithCancel(context.Background())
+
+	inodeURL, err := grpcfd.FilenameToURL(fileName)
+	s.Require().NoError(err)
+
+	return inodeURL.String(), fileClosedContext, cancelFunc
+}
+
+func (s *checkRecvfdTestSuite) TestRecvfdClosesSingleFile() {
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+
+	testFileName := path.Join(s.tempDir, "TestRecvfdClosesSingleFile.test")
+
+	inodeURLStr, fileClosedContext, cancelFunc := createFile(s, testFileName)
+
+	s.onFileClosedCallbacks = map[string]func(){
+		inodeURLStr: cancelFunc,
+	}
+
+	request := &networkservice.NetworkServiceRequest{
+		MechanismPreferences: []*networkservice.Mechanism{
+			{
+				Cls:  cls.LOCAL,
+				Type: kernel.MECHANISM,
+				Parameters: map[string]string{
+					common.InodeURL: "file:" + testFileName,
+				},
+			},
+		},
+	}
+
+	conn, err := s.testClient.Request(ctx, request)
+	s.Require().NoError(err)
+
+	_, err = s.testClient.Close(ctx, conn)
+	s.Require().NoError(err)
+
+	s.Require().Eventually(func() bool {
+		runtime.GC()
+		return fileClosedContext.Err() != nil
+	}, time.Second, time.Millisecond*100)
+}
+
+func (s *checkRecvfdTestSuite) TestRecvfdClosesMultipleFiles() {
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+
+	const numFiles = 3
+	s.onFileClosedContexts = make([]context.Context, numFiles)
+	s.onFileClosedCallbacks = make(map[string]func(), numFiles)
+
+	request := &networkservice.NetworkServiceRequest{
+		MechanismPreferences: make([]*networkservice.Mechanism, numFiles),
+	}
+
+	var filePath string
+	for i := 0; i < numFiles; i++ {
+		filePath = path.Join(s.tempDir, fmt.Sprintf("TestRecvfdClosesMultipleFiles.test%d", i))
+
+		inodeURLStr, fileClosedContext, cancelFunc := createFile(s, filePath)
+		s.onFileClosedCallbacks[inodeURLStr] = cancelFunc
+		s.onFileClosedContexts[i] = fileClosedContext
+
+		request.MechanismPreferences = append(request.MechanismPreferences,
+			&networkservice.Mechanism{
+				Cls:  cls.LOCAL,
+				Type: kernel.MECHANISM,
+				Parameters: map[string]string{
+					common.InodeURL: "file:" + filePath,
+				},
+			})
+	}
+
+	conn, err := s.testClient.Request(ctx, request)
+	s.Require().NoError(err)
+
+	_, err = s.testClient.Close(ctx, conn)
+	s.Require().NoError(err)
+
+	for i := range s.onFileClosedContexts {
+		onClosedFileCtx := s.onFileClosedContexts[i]
+		s.Require().Eventually(func() bool {
+			runtime.GC()
+			return onClosedFileCtx.Err() != nil
+		}, time.Second, time.Millisecond*100)
+	}
+}

pkg/registry/common/recvfd/server.go

@@ -83,7 +83,7 @@ func (r *recvfdNseServer) Register(ctx context.Context, endpoint *registry.Netwo
 		r.fileMaps.Store(endpoint.GetName(), fileMap)
 	}
 	// Swap back from File to Inode in the InodeURL in the Parameters
-	err = swapFileToInode(fileMap, returnedEndpoint, false)
+	err = swapFileToInode(fileMap, returnedEndpoint)
 	if err != nil {
 		return nil, err
 	}
@@ -98,6 +98,10 @@ func (r *recvfdNseServer) Unregister(ctx context.Context, endpoint *registry.Net
 	if endpoint.GetName() == "" {
 		return nil, errors.New("invalid endpoint specified")
 	}
+
+	// Clean up the fileMap no matter what happens
+	defer r.closeFiles(endpoint)
+
 	// Get the grpcfd.FDRecver
 	recv, ok := grpcfd.FromContext(ctx)
 	if !ok {
@@ -108,8 +112,6 @@ func (r *recvfdNseServer) Unregister(ctx context.Context, endpoint *registry.Net
 		filesByInodeURL:    make(map[string]*os.File),
 		inodeURLbyFilename: make(map[string]*url.URL),
 	})
-	// Clean up the fileMap no matter what happens
-	defer r.fileMaps.Delete(endpoint.GetName())
 
 	// Recv the FD and Swap the Inode for a file in InodeURL in Parameters
 	endpoint = endpoint.Clone()
@@ -126,7 +128,7 @@ func (r *recvfdNseServer) Unregister(ctx context.Context, endpoint *registry.Net
 
 	// Swap back from File to Inode in the InodeURL in the Parameters
 	endpoint = endpoint.Clone()
-	err = swapFileToInode(fileMap, endpoint, true)
+	err = swapFileToInode(fileMap, endpoint)
 	if err != nil {
 		return nil, err
 	}
@@ -178,7 +180,7 @@ func recvFDAndSwapInodeToUnix(ctx context.Context, fileMap *perEndpointFileMap,
 	return err
 }
 
-func swapFileToInode(fileMap *perEndpointFileMap, endpoint *registry.NetworkServiceEndpoint, closeAllFiles bool) error {
+func swapFileToInode(fileMap *perEndpointFileMap, endpoint *registry.NetworkServiceEndpoint) error {
 	// Transform string to URL for correctness checking and ease of use
 	unixURL, err := url.Parse(endpoint.GetUrl())
 	if err != nil {
@@ -198,15 +200,28 @@ func swapFileToInode(fileMap *perEndpointFileMap, endpoint *registry.NetworkServ
 		// Swap the fileURL for the inodeURL in parameters
 		endpoint.Url = inodeURL.String()
 
-		// If closeAllFiles == true, close any files we may have open for any other inodes
 		// This is used to clean up files sent by MechanismPreferences that were *not* selected to be the
 		// connection mechanism
 		for inodeURLStr, file := range fileMap.filesByInodeURL {
-			if closeAllFiles || inodeURLStr != inodeURL.String() {
+			if inodeURLStr != inodeURL.String() {
 				delete(fileMap.filesByInodeURL, inodeURLStr)
 				_ = file.Close()
 			}
 		}
 	})
 	return nil
 }
+
+func (r *recvfdNseServer) closeFiles(endpoint *registry.NetworkServiceEndpoint) {
+	defer r.fileMaps.Delete(endpoint.GetName())
+
+	fileMap, _ := r.fileMaps.LoadOrStore(endpoint.GetName(), &perEndpointFileMap{
+		filesByInodeURL:    make(map[string]*os.File),
+		inodeURLbyFilename: make(map[string]*url.URL),
+	})
+
+	for inodeURLStr, file := range fileMap.filesByInodeURL {
+		delete(fileMap.filesByInodeURL, inodeURLStr)
+		_ = file.Close()
+	}
+}