github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/utils/memmod/dll_windows.go (about)

     1  //go:build windows
     2  // +build windows
     3  
     4  /* SPDX-License-Identifier: MIT
     5   *
     6   * Copyright (C) 2017-2022 WireGuard LLC. All Rights Reserved.
     7   */
     8  
     9  package memmod
    10  
    11  import (
    12  	"fmt"
    13  	"sync"
    14  	"sync/atomic"
    15  	"syscall"
    16  	"unsafe"
    17  
    18  	"golang.org/x/sys/windows"
    19  )
    20  
    21  type lazyDLL struct {
    22  	Name   string
    23  	Base   windows.Handle
    24  	mu     sync.Mutex
    25  	module *Module
    26  
    27  	dllBytes []byte
    28  }
    29  
    30  func (d *lazyDLL) Load() error {
    31  	if atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&d.module))) != nil {
    32  		return nil
    33  	}
    34  	d.mu.Lock()
    35  	defer d.mu.Unlock()
    36  	if d.module != nil {
    37  		return nil
    38  	}
    39  
    40  	module, err := LoadLibrary(d.dllBytes)
    41  	if err != nil {
    42  		return fmt.Errorf("Unable to load library: %w", err)
    43  	}
    44  	d.Base = windows.Handle(module.BaseAddr())
    45  
    46  	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&d.module)), unsafe.Pointer(module))
    47  	return nil
    48  }
    49  
    50  func (p *lazyProc) nameToAddr() (uintptr, error) {
    51  	return p.dll.module.ProcAddressByName(p.Name)
    52  }
    53  
    54  func NewLazyDLL(name string, dllBytes []byte) *lazyDLL {
    55  	return &lazyDLL{Name: name, dllBytes: dllBytes}
    56  }
    57  
    58  func (d *lazyDLL) NewProc(name string) *lazyProc {
    59  	return &lazyProc{dll: d, Name: name}
    60  }
    61  
    62  type lazyProc struct {
    63  	Name string
    64  	mu   sync.Mutex
    65  	dll  *lazyDLL
    66  	addr uintptr
    67  }
    68  
    69  func (p *lazyProc) Find() error {
    70  	if atomic.LoadPointer((*unsafe.Pointer)(unsafe.Pointer(&p.addr))) != nil {
    71  		return nil
    72  	}
    73  	p.mu.Lock()
    74  	defer p.mu.Unlock()
    75  	if p.addr != 0 {
    76  		return nil
    77  	}
    78  
    79  	err := p.dll.Load()
    80  	if err != nil {
    81  		return fmt.Errorf("Error loading %v DLL: %w", p.dll.Name, err)
    82  	}
    83  	addr, err := p.nameToAddr()
    84  	if err != nil {
    85  		return fmt.Errorf("Error getting %v address: %w", p.Name, err)
    86  	}
    87  
    88  	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&p.addr)), unsafe.Pointer(addr))
    89  	return nil
    90  }
    91  
    92  func (p *lazyProc) Addr() uintptr {
    93  	err := p.Find()
    94  	if err != nil {
    95  		panic(err)
    96  	}
    97  	return p.addr
    98  }
    99  
   100  func (p *lazyProc) Call(i ...uintptr) (r1, r2 uintptr, err syscall.Errno) {
   101  	return syscall.SyscallN(p.Addr(), i...)
   102  }
   103  
   104  func StrPtr(s string) (uintptr, error) {
   105  	b, err := syscall.BytePtrFromString(s)
   106  	if err != nil {
   107  		return 0, err
   108  	}
   109  	return uintptr(unsafe.Pointer(b)), nil
   110  }