github.com/fzfile/BaiduPCS-Go@v0.0.0-20200606205115-4408961cf336/pcsutil/prealloc/prealloc_windows.go (about)

     1  package prealloc
     2  
     3  import (
     4  	"golang.org/x/sys/windows"
     5  	"log"
     6  	"syscall"
     7  )
     8  
     9  var (
    10  	kernel32             = syscall.NewLazyDLL("kernel32.dll")
    11  	procSetFileValidData = kernel32.NewProc("SetFileValidData")
    12  )
    13  
    14  func initPrivilege() error {
    15  	current, err := windows.GetCurrentProcess()
    16  	if err != nil {
    17  		return &PreAllocError{
    18  			ProcName: "GetCurrentProcess",
    19  			Err:      err,
    20  		}
    21  	}
    22  
    23  	var hToken windows.Token
    24  	err = windows.OpenProcessToken(current, windows.TOKEN_ADJUST_PRIVILEGES|windows.TOKEN_QUERY, &hToken)
    25  	if err != nil {
    26  		return &PreAllocError{
    27  			ProcName: "OpenProcessToken",
    28  			Err:      err,
    29  		}
    30  	}
    31  
    32  	var (
    33  		seManageVolumeName, _ = windows.UTF16PtrFromString("SeManageVolumePrivilege")
    34  		tp                    = windows.Tokenprivileges{
    35  			PrivilegeCount: 1,
    36  			Privileges: [1]windows.LUIDAndAttributes{
    37  				windows.LUIDAndAttributes{
    38  					Luid:       windows.LUID{},
    39  					Attributes: windows.SE_PRIVILEGE_ENABLED,
    40  				},
    41  			},
    42  		}
    43  	)
    44  	err = windows.LookupPrivilegeValue(nil, seManageVolumeName, &tp.Privileges[0].Luid)
    45  	if err != nil {
    46  		return &PreAllocError{
    47  			ProcName: "LookupPrivilegeValue",
    48  			Err:      err,
    49  		}
    50  	}
    51  
    52  	err = windows.AdjustTokenPrivileges(hToken, false, &tp, 0, nil, nil)
    53  	if err != nil {
    54  		return &PreAllocError{
    55  			ProcName: "AdjustTokenPrivileges",
    56  			Err:      err,
    57  		}
    58  	}
    59  
    60  	return nil
    61  }
    62  
    63  // 初始化权限
    64  func init() {
    65  	err := initPrivilege()
    66  	if err != nil {
    67  		log.Printf("prealloc: init privileges error: %s\n", err) // 打印警告
    68  	}
    69  }
    70  
    71  // PreAlloc 预分配文件空间
    72  func PreAlloc(fd uintptr, length int64) error {
    73  	err := syscall.Ftruncate(syscall.Handle(fd), length)
    74  	if err != nil {
    75  		return &PreAllocError{
    76  			ProcName: "Ftruncate",
    77  			Err:      err,
    78  		}
    79  	}
    80  
    81  	r1, _, err := procSetFileValidData.Call(fd, uintptr(length))
    82  	if r1 == 0 {
    83  		return &PreAllocError{
    84  			ProcName: "SetFileValidData",
    85  			Err:      err,
    86  		}
    87  	}
    88  
    89  	return nil
    90  }