Use WMI to implement SMB API to reduce PowerShell overhead

Author: laozc

pkg/cim/smb.go

@@ -0,0 +1,53 @@
+//go:build windows
+// +build windows
+
+package cim
+
+import (
+	"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
+)
+
+// 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", remotePath)
+	instances, err := QueryInstances(WMINamespaceSmb, smbQuery)
+	if err != nil {
+		return nil, err
+	}
+
+	return instances[0], err
+}
+
+// RemoveSmbGlobalMappingByRemotePath removes a 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", remotePath)
+	instances, err := QueryInstances(WMINamespaceSmb, smbQuery)
+	if err != nil {
+		return err
+	}
+
+	_, err = instances[0].InvokeMethod("Remove", true)
+	return err
+}

pkg/cim/wmi.go

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

pkg/os/smb/api.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/kubernetes-csi/csi-proxy/pkg/cim"
 	"github.com/kubernetes-csi/csi-proxy/pkg/utils"
 )
 
@@ -26,18 +27,22 @@ func New(requirePrivacy bool) *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)
+	inst, err := cim.QuerySmbGlobalMappingByRemotePath(remotePathForQuery(remotePath))
 	if err != nil {
-		return false, fmt.Errorf("error checking smb mapping. cmd %s, output: %s, err: %v", remotePath, string(out), err)
+		return false, cim.IgnoreNotFound(err)
 	}
 
-	if len(out) == 0 || !strings.EqualFold(strings.TrimSpace(string(out)), "OK") {
-		return false, nil
+	status, err := inst.GetProperty("Status")
+	if err != nil {
+		return false, err
 	}
-	return true, nil
+
+	return status.(int32) == cim.SmbMappingStatusOK, nil
 }
 
 // NewSmbLink - creates a directory symbolic link to the remote share.
@@ -48,7 +53,6 @@ func (*SmbAPI) IsSmbMapped(remotePath string) (bool, error) {
 // alpha to merge the paths.
 // TODO (for beta release): Merge the link paths - os.Symlink and Powershell link path.
 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.
@@ -78,12 +82,26 @@ func (api *SmbAPI) NewSmbGlobalMapping(remotePath, username, password string) er
 		return fmt.Errorf("NewSmbGlobalMapping failed. output: %q, err: %v", string(output), err)
 	}
 	return nil
+	//TODO: move to use WMI when the credentials could be correctly handled
+	//params := map[string]interface{}{
+	//	"RemotePath":     remotePath,
+	//	"RequirePrivacy": api.RequirePrivacy,
+	//}
+	//if username != "" {
+	//	params["Credential"] = fmt.Sprintf("%s:%s", username, password)
+	//}
+	//result, _, err := cim.InvokeCimMethod(cim.WMINamespaceSmb, "MSFT_SmbGlobalMapping", "Create", params)
+	//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)
+	err := cim.RemoveSmbGlobalMappingByRemotePath(remotePathForQuery(remotePath))
+	if err != nil {
+		return fmt.Errorf("error remove smb mapping '%s'. err: %v", remotePath, err)
 	}
+
 	return nil
 }