feat: Use WMI to implement SMB API to reduce PowerShell overhead (library version)

pkg/cim/smb.go

@@ -0,0 +1,100 @@
+//go:build windows
+// +build windows
+
+package cim
+
+import (
+	"strings"
+
+	"github.com/microsoft/wmi/pkg/base/query"
+	cim "github.com/microsoft/wmi/pkg/wmiinstance"
+)
+
+// Refer to https://learn.microsoft.com/en-us/previous-versions/windows/desktop/smb/msft-smbmapping
+const (
+	SmbMappingStatusOK int32 = iota
+	SmbMappingStatusPaused
+	SmbMappingStatusDisconnected
+	SmbMappingStatusNetworkError
+	SmbMappingStatusConnecting
+	SmbMappingStatusReconnecting
+	SmbMappingStatusUnavailable
+
+	credentialDelimiter = ":"
+)
+
+// escapeQueryParameter escapes a parameter for WMI Queries
+func escapeQueryParameter(s string) string {
+	s = strings.ReplaceAll(s, "'", "''")
+	s = strings.ReplaceAll(s, "\\", "\\\\")
+	return s
+}
+
+func escapeUserName(userName string) string {
+	// refer to https://github.com/PowerShell/PowerShell/blob/9303de597da55963a6e26a8fe164d0b256ca3d4d/src/Microsoft.PowerShell.Commands.Management/cimSupport/cmdletization/cim/cimConverter.cs#L169-L170
+	userName = strings.ReplaceAll(userName, "\\", "\\\\")
+	userName = strings.ReplaceAll(userName, credentialDelimiter, "\\"+credentialDelimiter)
+	return userName
+}
+
+// QuerySmbGlobalMappingByRemotePath retrieves the SMB global mapping from its remote path.
+//
+// The equivalent WMI query is:
+//
+//	SELECT [selectors] FROM MSFT_SmbGlobalMapping
+//
+// Refer to https://pkg.go.dev/github.com/microsoft/wmi/server2019/root/microsoft/windows/smb#MSFT_SmbGlobalMapping
+// for the WMI class definition.
+func QuerySmbGlobalMappingByRemotePath(remotePath string) (*cim.WmiInstance, error) {
+	smbQuery := query.NewWmiQuery("MSFT_SmbGlobalMapping", "RemotePath", escapeQueryParameter(remotePath))
+	instances, err := QueryInstances(WMINamespaceSmb, smbQuery)
+	if err != nil {
+		return nil, err
+	}
+
+	return instances[0], err
+}
+
+// GetSmbGlobalMappingStatus returns the status of an SMB global mapping.
+func GetSmbGlobalMappingStatus(inst *cim.WmiInstance) (int32, error) {
+	statusProp, err := inst.GetProperty("Status")
+	if err != nil {
+		return SmbMappingStatusUnavailable, err
+	}
+
+	return statusProp.(int32), nil
+}
+
+// RemoveSmbGlobalMappingByRemotePath removes an SMB global mapping matching to the remote path.
+//
+// Refer to https://pkg.go.dev/github.com/microsoft/wmi/server2019/root/microsoft/windows/smb#MSFT_SmbGlobalMapping
+// for the WMI class definition.
+func RemoveSmbGlobalMappingByRemotePath(remotePath string) error {
+	smbQuery := query.NewWmiQuery("MSFT_SmbGlobalMapping", "RemotePath", escapeQueryParameter(remotePath))
+	instances, err := QueryInstances(WMINamespaceSmb, smbQuery)
+	if err != nil {
+		return err
+	}
+
+	_, err = instances[0].InvokeMethod("Remove", true)
+	return err
+}
+
+// NewSmbGlobalMapping creates a new SMB global mapping to the remote path.
+//
+// Refer to https://pkg.go.dev/github.com/microsoft/wmi/server2019/root/microsoft/windows/smb#MSFT_SmbGlobalMapping
+// for the WMI class definition.
+func NewSmbGlobalMapping(remotePath, username, password string, requirePrivacy bool) (int, error) {
+	params := map[string]interface{}{
+		"RemotePath":     remotePath,
+		"RequirePrivacy": requirePrivacy,
+	}
+	if username != "" {
+		// refer to https://github.com/PowerShell/PowerShell/blob/9303de597da55963a6e26a8fe164d0b256ca3d4d/src/Microsoft.PowerShell.Commands.Management/cimSupport/cmdletization/cim/cimConverter.cs#L166-L178
+		// on how SMB credential is handled in PowerShell
+		params["Credential"] = escapeUserName(username) + credentialDelimiter + password
+	}
+
+	result, _, err := InvokeCimMethod(WMINamespaceSmb, "MSFT_SmbGlobalMapping", "Create", params)
+	return result, err
+}

pkg/cim/wmi.go

@@ -20,6 +20,7 @@ import (
 const (
 	WMINamespaceCimV2   = "Root\\CimV2"
 	WMINamespaceStorage = "Root\\Microsoft\\Windows\\Storage"
+	WMINamespaceSmb     = "Root\\Microsoft\\Windows\\Smb"
 )
 
 type InstanceHandler func(instance *cim.WmiInstance) (bool, error)

pkg/smb/hostapi/hostapi.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/kubernetes-csi/csi-proxy/v2/pkg/cim"
 	"github.com/kubernetes-csi/csi-proxy/v2/pkg/utils"
 )
 
@@ -22,62 +23,65 @@ func New() HostAPI {
 	return smbAPI{}
 }
 
+func remotePathForQuery(remotePath string) string {
+	return strings.ReplaceAll(remotePath, "\\", "\\\\")
+}
+
 func (smbAPI) IsSMBMapped(remotePath string) (bool, error) {
-	cmdLine := `$(Get-SmbGlobalMapping -RemotePath $Env:smbremotepath -ErrorAction Stop).Status `
-	cmdEnv := fmt.Sprintf("smbremotepath=%s", remotePath)
-	out, err := utils.RunPowershellCmd(cmdLine, cmdEnv)
-	if err != nil {
-		return false, fmt.Errorf("error checking SMB mapping. cmd %s, output: %s, err: %v", remotePath, string(out), err)
-	}
+	var isMapped bool
+	err := cim.WithCOMThread(func() error {
+		inst, err := cim.QuerySmbGlobalMappingByRemotePath(remotePath)
+		if err != nil {
+			return err
+		}
 
-	if len(out) == 0 || !strings.EqualFold(strings.TrimSpace(string(out)), "OK") {
-		return false, nil
-	}
-	return true, nil
+		status, err := cim.GetSmbGlobalMappingStatus(inst)
+		if err != nil {
+			return err
+		}
+
+		isMapped = status == cim.SmbMappingStatusOK
+		return nil
+	})
+	return isMapped, cim.IgnoreNotFound(err)
 }
 
 // NewSMBLink - creates a directory symbolic link to the remote share.
 // The os.Symlink was having issue for cases where the destination was an SMB share - the container
-// runtime would complain stating "Access Denied". Because of this, we had to perform
-// this operation with powershell commandlet creating an directory softlink.
-// Since os.Symlink is currently being used in working code paths, no attempt is made in
-// alpha to merge the paths.
-// TODO (for beta release): Merge the link paths - os.Symlink and Powershell link path.
+// runtime would complain stating "Access Denied".
 func (smbAPI) NewSMBLink(remotePath, localPath string) error {
 	if !strings.HasSuffix(remotePath, "\\") {
 		// Golang has issues resolving paths mapped to file shares if they do not end in a trailing \
 		// so add one if needed.
 		remotePath = remotePath + "\\"
 	}
+	longRemotePath := utils.EnsureLongPath(remotePath)
+	longLocalPath := utils.EnsureLongPath(localPath)
 
-	cmdLine := `New-Item -ItemType SymbolicLink $Env:smblocalPath -Target $Env:smbremotepath`
-	output, err := utils.RunPowershellCmd(cmdLine, fmt.Sprintf("smbremotepath=%s", remotePath), fmt.Sprintf("smblocalpath=%s", localPath))
+	err := utils.CreateSymlink(longLocalPath, longRemotePath, true)
 	if err != nil {
-		return fmt.Errorf("error linking %s to %s. output: %s, err: %v", remotePath, localPath, string(output), err)
+		return fmt.Errorf("error linking %s to %s. err: %v", remotePath, localPath, err)
 	}
 
 	return nil
 }
 
 func (smbAPI) NewSMBGlobalMapping(remotePath, username, password string) error {
-	// use PowerShell Environment Variables to store user input string to prevent command line injection
-	// https://docs.microsoft.com/en-us/powershell/module/microsoft.powershell.core/about/about_environment_variables?view=powershell-5.1
-	cmdLine := fmt.Sprintf(`$PWord = ConvertTo-SecureString -String $Env:smbpassword -AsPlainText -Force` +
-		`;$Credential = New-Object -TypeName System.Management.Automation.PSCredential -ArgumentList $Env:smbuser, $PWord` +
-		`;New-SmbGlobalMapping -RemotePath $Env:smbremotepath -Credential $Credential -RequirePrivacy $true`)
-
-	if output, err := utils.RunPowershellCmd(cmdLine, fmt.Sprintf("smbuser=%s", username),
-		fmt.Sprintf("smbpassword=%s", password),
-		fmt.Sprintf("smbremotepath=%s", remotePath)); err != nil {
-		return fmt.Errorf("NewSMBGlobalMapping failed. output: %q, err: %v", string(output), err)
-	}
-	return nil
+	return cim.WithCOMThread(func() error {
+		result, err := cim.NewSmbGlobalMapping(remotePath, username, password, true)
+		if err != nil {
+			return fmt.Errorf("NewSmbGlobalMapping failed. result: %d, err: %v", result, err)
+		}
+		return nil
+	})
 }
 
 func (smbAPI) RemoveSMBGlobalMapping(remotePath string) error {
-	cmd := `Remove-SmbGlobalMapping -RemotePath $Env:smbremotepath -Force`
-	if output, err := utils.RunPowershellCmd(cmd, fmt.Sprintf("smbremotepath=%s", remotePath)); err != nil {
-		return fmt.Errorf("UnmountSMBShare failed. output: %q, err: %v", string(output), err)
-	}
-	return nil
+	return cim.WithCOMThread(func() error {
+		err := cim.RemoveSmbGlobalMappingByRemotePath(remotePath)
+		if err != nil {
+			return fmt.Errorf("error remove smb mapping '%s'. err: %v", remotePath, err)
+		}
+		return nil
+	})
 }

pkg/utils/utils.go

@@ -1,12 +1,12 @@
 package utils
 
 import (
-	"errors"
 	"fmt"
 	"os"
 	"os/exec"
 	"strings"
 
+	"github.com/pkg/errors"
 	"golang.org/x/sys/windows"
 	"k8s.io/klog/v2"
 )
@@ -76,3 +76,38 @@ func IsPathSymlink(path string) (bool, error) {
 	isSymlink := fi.Mode()&os.ModeSymlink != 0 || fi.Mode()&os.ModeIrregular != 0
 	return isSymlink, nil
 }
+
+func CreateSymlink(link, target string, isDir bool) error {
+	linkPtr, err := windows.UTF16PtrFromString(link)
+	if err != nil {
+		return err
+	}
+	targetPtr, err := windows.UTF16PtrFromString(target)
+	if err != nil {
+		return err
+	}
+
+	var flags uint32
+	if isDir {
+		flags = windows.SYMBOLIC_LINK_FLAG_DIRECTORY
+	}
+
+	err = windows.CreateSymbolicLink(
+		linkPtr,
+		targetPtr,
+		flags,
+	)
+	return err
+}
+
+// PathExists checks whether the given `path` exists.
+func PathExists(path string) (bool, error) {
+	_, err := os.Lstat(path)
+	if err == nil {
+		return true, nil
+	}
+	if os.IsNotExist(err) {
+		return false, nil
+	}
+	return false, err
+}