Added Release() method to the to unmount and return stub to freeStubs list

Signed-off-by: bpopovschi <[email protected]>

Author: Zyqsempai

agent/drive_handler.go

@@ -177,22 +177,11 @@ func (dh driveHandler) MountDrive(ctx context.Context, req *drivemount.MountDriv
 	logger = logger.WithField("drive_path", drive.Path())
 
 	// Do a basic check that we won't be mounting over any important system directories
-	resolvedDest, err := evalAnySymlinks(req.DestinationPath)
-	if err != nil {
-		return nil, errors.Wrapf(err,
-			"failed to evaluate any symlinks in drive mount destination %q", req.DestinationPath)
-	}
-
-	for _, systemDir := range bannedSystemDirs {
-		if isOrUnderDir(resolvedDest, systemDir) {
-			return nil, errors.Errorf(
-				"drive mount destination %q resolves to path %q under banned system directory %q",
-				req.DestinationPath, resolvedDest, systemDir,
-			)
-		}
+	if err := checkImportantDirs(req.DestinationPath); err != nil {
+		return nil, err
 	}
 
-	err = os.MkdirAll(req.DestinationPath, 0700)
+	err := os.MkdirAll(req.DestinationPath, 0700)
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to create drive mount destination %q", req.DestinationPath)
 	}
@@ -229,6 +218,45 @@ func (dh driveHandler) MountDrive(ctx context.Context, req *drivemount.MountDriv
 		drive.Path(), req.DestinationPath)
 }
 
+func (dh driveHandler) UnmountDrive(ctx context.Context, req *drivemount.UnmountDriveRequest) (*empty.Empty, error) {
+	drive, ok := dh.GetDrive(req.DriveID)
+	if !ok {
+		return nil, fmt.Errorf("drive %q could not be found", req.DriveID)
+	}
+
+	// Do a basic check that we won't be mounting over any important system directories
+	if err := checkImportantDirs(req.DestinationPath); err != nil {
+		return nil, err
+	}
+
+	err := mount.Unmount(req.DestinationPath, 0)
+	if err == nil {
+		return &empty.Empty{}, nil
+	}
+
+	return nil, errors.Errorf("failed to unmount the drive from %q to %q",
+		drive.Path(), req.DestinationPath)
+}
+
+func checkImportantDirs(path string) error {
+	resolvedDest, err := evalAnySymlinks(path)
+	if err != nil {
+		return errors.Wrapf(err,
+			"failed to evaluate any symlinks in drive ummount destination %q", path)
+	}
+
+	for _, systemDir := range bannedSystemDirs {
+		if isOrUnderDir(resolvedDest, systemDir) {
+			return errors.Errorf(
+				"drive mount destination %q resolves to path %q under banned system directory %q",
+				path, resolvedDest, systemDir,
+			)
+		}
+	}
+
+	return nil
+}
+
 // evalAnySymlinks is similar to filepath.EvalSymlinks, except it will not return an error if part of the
 // provided path does not exist. It will evaluate symlinks present in the path up to a component that doesn't
 // exist, at which point it will just append the rest of the provided path to what has been resolved so far.

proto/service/drivemount/drivemount.proto

@@ -6,11 +6,17 @@ option go_package = "drivemount";
 
 service DriveMounter {
     rpc MountDrive(MountDriveRequest) returns (google.protobuf.Empty);
+    rpc UnmountDrive(UnmountDriveRequest) returns (google.protobuf.Empty);
 }
 
 message MountDriveRequest {
+     string DriveID = 1;
+     string DestinationPath = 2;
+     string FilesytemType = 3;
+     repeated string Options = 4;
+}
+
+message UnmountDriveRequest {
     string DriveID = 1;
     string DestinationPath = 2;
-    string FilesytemType = 3;
-    repeated string Options = 4;
 }