From 77c710b60a5feb037e0a1b5554664383d9b617ee Mon Sep 17 00:00:00 2001 From: Alex Crane Date: Fri, 11 Sep 2020 21:13:40 +0100 Subject: [PATCH] Use HeapAlloc for allocating memory used in API calls where pointers to within the allocated memory are returned. It is believed that it is unsafe to use Go allocated memory for this --- printer.go | 115 ++++++++++++++++++++++++++++++++++++++++------------- zapi.go | 35 ++++++++++++++++ 2 files changed, 122 insertions(+), 28 deletions(-) diff --git a/printer.go b/printer.go index 106d646..c69b984 100644 --- a/printer.go +++ b/printer.go @@ -221,9 +221,25 @@ const ( JOB_CONTROL_LAST_PAGE_EJECTED = 7 JOB_CONTROL_RETAIN = 8 JOB_CONTROL_RELEASE = 9 + + HEAP_NO_SERIALIZE = 0x1 + HEAP_GENERATE_EXCEPTIONS = 0x4 + HEAP_ZERO_MEMORY = 0x8 +) + +var ( + ErrNoMem = errors.New("memory allocation failed") + ErrNoNotification = errors.New("no notification information") + processHeap syscall.Handle ) -var ErrNoNotification = errors.New("no notification information") +func init() { + var err error + processHeap, err = GetProcessHeap() + if err != nil { + panic(err) + } +} //sys GetDefaultPrinter(buf *uint16, bufN *uint32) (err error) = winspool.GetDefaultPrinterW //sys ClosePrinter(h syscall.Handle) (err error) = winspool.ClosePrinter @@ -243,6 +259,9 @@ var ErrNoNotification = errors.New("no notification information") //sys FindNextPrinterChangeNotification(h syscall.Handle, cause *uint16, options *PRINTER_NOTIFY_OPTIONS, info **PRINTER_NOTIFY_INFO) (err error) = winspool.FindNextPrinterChangeNotification //sys FindClosePrinterChangeNotification(h syscall.Handle) (err error) = winspool.FindClosePrinterChangeNotification //sys FreePrinterNotifyInfo(info *PRINTER_NOTIFY_INFO) (err error) = winspool.FreePrinterNotifyInfo +//sys GetProcessHeap() (h syscall.Handle, err error) = kernel32.GetProcessHeap +//sys HeapAlloc(h syscall.Handle, flags uint32, size uint) (p unsafe.Pointer) = kernel32.HeapAlloc +//sys HeapFree(h syscall.Handle, flags uint32, p unsafe.Pointer) (err error) = kernel32.HeapFree func Default() (string, error) { b := make([]uint16, 3) @@ -265,19 +284,28 @@ func Default() (string, error) { func ReadNames() ([]string, error) { const flags = PRINTER_ENUM_LOCAL | PRINTER_ENUM_CONNECTIONS var needed, returned uint32 - buf := make([]byte, 1) - err := EnumPrinters(flags, nil, printerInfoLevel, &buf[0], uint32(len(buf)), &needed, &returned) - if err != nil { + var bufsize uint32 = 1024 // 1KB initial attempt + var buf unsafe.Pointer + for { + buf = HeapAlloc(processHeap, HEAP_ZERO_MEMORY, uint(bufsize)) + if buf == nil { + return nil, ErrNoMem + } + err := EnumPrinters(flags, nil, printerInfoLevel, (*byte)(buf), bufsize, &needed, &returned) + if err == nil { + break + } if err != syscall.ERROR_INSUFFICIENT_BUFFER { return nil, err } - buf = make([]byte, needed) - err = EnumPrinters(flags, nil, printerInfoLevel, &buf[0], uint32(len(buf)), &needed, &returned) - if err != nil { + if needed <= bufsize { return nil, err } + HeapFree(processHeap, 0, buf) + bufsize = needed } - ps := (*[1024]PRINTER_INFO_5)(unsafe.Pointer(&buf[0]))[:returned:returned] + defer HeapFree(processHeap, 0, buf) + ps := (*[1024]PRINTER_INFO_5)(buf)[:returned:returned] names := make([]string, 0, returned) for _, p := range ps { names = append(names, utf16PtrToString(p.PrinterName)) @@ -483,25 +511,33 @@ func (j *JOB_INFO_4) ToJobInfo() *JobInfo { // Jobs returns information about all print jobs on this printer func (p *Printer) Jobs() ([]JobInfo, error) { var bytesNeeded, jobsReturned uint32 - buf := make([]byte, 1) + var bufsize uint32 = 10 * 1024 // 10KB initial attempt + var buf unsafe.Pointer for { - err := EnumJobs(p.h, 0, 255, jobInfoLevel, &buf[0], uint32(len(buf)), &bytesNeeded, &jobsReturned) + buf = HeapAlloc(processHeap, HEAP_ZERO_MEMORY, uint(bufsize)) + if buf == nil { + return nil, ErrNoMem + } + err := EnumJobs(p.h, 0, 255, jobInfoLevel, (*byte)(buf), bufsize, &bytesNeeded, &jobsReturned) if err == nil { break } if err != syscall.ERROR_INSUFFICIENT_BUFFER { return nil, err } - if bytesNeeded <= uint32(len(buf)) { + if bytesNeeded <= bufsize { return nil, err } - buf = make([]byte, bytesNeeded) + HeapFree(processHeap, 0, buf) + bufsize = bytesNeeded } + + defer HeapFree(processHeap, 0, buf) if jobsReturned <= 0 { return nil, nil } pjs := make([]JobInfo, 0, jobsReturned) - ji := (*[2048]JOB_INFO_4)(unsafe.Pointer(&buf[0]))[:jobsReturned:jobsReturned] + ji := (*[256]JOB_INFO_4)(buf)[:jobsReturned:jobsReturned] for _, j := range ji { pjs = append(pjs, *j.ToJobInfo()) } @@ -510,22 +546,29 @@ func (p *Printer) Jobs() ([]JobInfo, error) { func (p *Printer) Job(jobId uint32) (*JobInfo, error) { var bytesNeeded uint32 - buf := make([]byte, 1024) + var bufsize uint32 = 1024 // 1KB initial attempt + var buf unsafe.Pointer for { - err := GetJob(p.h, jobId, jobInfoLevel, &buf[0], uint32(len(buf)), &bytesNeeded) + buf = HeapAlloc(processHeap, HEAP_ZERO_MEMORY, uint(bufsize)) + if buf == nil { + return nil, ErrNoMem + } + err := GetJob(p.h, jobId, jobInfoLevel, (*byte)(buf), bufsize, &bytesNeeded) if err == nil { break } if err != syscall.ERROR_INSUFFICIENT_BUFFER { return nil, err } - if bytesNeeded <= uint32(len(buf)) { + if bytesNeeded <= bufsize { return nil, err } - buf = make([]byte, bytesNeeded) + HeapFree(processHeap, 0, buf) + bufsize = bytesNeeded } - ji := (*JOB_INFO_4)(unsafe.Pointer(&buf[0])) + defer HeapFree(processHeap, 0, buf) + ji := (*JOB_INFO_4)(buf) return ji.ToJobInfo(), nil } @@ -540,21 +583,29 @@ func (p *Printer) SetJob(jobID uint32, jobInfo *JobInfo, command uint32) error { // DriverInfo returns information about a printer's driver. func (p *Printer) DriverInfo() (*DriverInfo, error) { var needed uint32 - b := make([]byte, 1024*10) + var bufsize uint32 = 10 * 1024 // 10KB initial attempt + var buf unsafe.Pointer for { - err := GetPrinterDriver(p.h, nil, 8, &b[0], uint32(len(b)), &needed) + buf = HeapAlloc(processHeap, HEAP_ZERO_MEMORY, uint(bufsize)) + if buf == nil { + return nil, ErrNoMem + } + err := GetPrinterDriver(p.h, nil, 8, (*byte)(buf), bufsize, &needed) if err == nil { break } if err != syscall.ERROR_INSUFFICIENT_BUFFER { return nil, err } - if needed <= uint32(len(b)) { + if needed <= bufsize { return nil, err } - b = make([]byte, needed) + HeapFree(processHeap, 0, buf) + bufsize = needed } - di := (*DRIVER_INFO_8)(unsafe.Pointer(&b[0])) + + defer HeapFree(processHeap, 0, buf) + di := (*DRIVER_INFO_8)(buf) return &DriverInfo{ Attributes: di.PrinterDriverAttributes, Name: utf16PtrToString(di.Name), @@ -566,21 +617,29 @@ func (p *Printer) DriverInfo() (*DriverInfo, error) { // PrinterInfo returns information about a printer func (p *Printer) PrinterInfo() (*Info, error) { var needed uint32 - b := make([]byte, 1024*10) + var bufsize uint32 = 1024 // 1KB initial attempt + var buf unsafe.Pointer for { - err := GetPrinter(p.h, 5, &b[0], uint32(len(b)), &needed) + buf = HeapAlloc(processHeap, HEAP_ZERO_MEMORY, uint(bufsize)) + if buf == nil { + return nil, ErrNoMem + } + err := GetPrinter(p.h, 5, (*byte)(buf), bufsize, &needed) if err == nil { break } if err != syscall.ERROR_INSUFFICIENT_BUFFER { return nil, err } - if needed <= uint32(len(b)) { + if needed <= bufsize { return nil, err } - b = make([]byte, needed) + HeapFree(processHeap, 0, buf) + bufsize = needed } - pi := (*PRINTER_INFO_5)(unsafe.Pointer(&b[0])) + + defer HeapFree(processHeap, 0, buf) + pi := (*PRINTER_INFO_5)(buf) return &Info{ PrinterName: utf16PtrToString(pi.PrinterName), PortName: utf16PtrToString(pi.PortName), diff --git a/zapi.go b/zapi.go index 83b06ab..6f646a6 100644 --- a/zapi.go +++ b/zapi.go @@ -38,6 +38,7 @@ func errnoErr(e syscall.Errno) error { var ( modwinspool = windows.NewLazySystemDLL("winspool.drv") + modkernel32 = windows.NewLazySystemDLL("kernel32.dll") procGetDefaultPrinterW = modwinspool.NewProc("GetDefaultPrinterW") procClosePrinter = modwinspool.NewProc("ClosePrinter") @@ -57,6 +58,9 @@ var ( procFindNextPrinterChangeNotification = modwinspool.NewProc("FindNextPrinterChangeNotification") procFindClosePrinterChangeNotification = modwinspool.NewProc("FindClosePrinterChangeNotification") procFreePrinterNotifyInfo = modwinspool.NewProc("FreePrinterNotifyInfo") + procGetProcessHeap = modkernel32.NewProc("GetProcessHeap") + procHeapAlloc = modkernel32.NewProc("HeapAlloc") + procHeapFree = modkernel32.NewProc("HeapFree") ) func GetDefaultPrinter(buf *uint16, bufN *uint32) (err error) { @@ -275,3 +279,34 @@ func FreePrinterNotifyInfo(info *PRINTER_NOTIFY_INFO) (err error) { } return } + +func GetProcessHeap() (h syscall.Handle, err error) { + r0, _, e1 := syscall.Syscall(procGetProcessHeap.Addr(), 0, 0, 0, 0) + h = syscall.Handle(r0) + if h == 0 { + if e1 != 0 { + err = errnoErr(e1) + } else { + err = syscall.EINVAL + } + } + return +} + +func HeapAlloc(h syscall.Handle, flags uint32, size uint) (p unsafe.Pointer) { + r0, _, _ := syscall.Syscall(procHeapAlloc.Addr(), 3, uintptr(h), uintptr(flags), uintptr(size)) + p = unsafe.Pointer(r0) + return +} + +func HeapFree(h syscall.Handle, flags uint32, p unsafe.Pointer) (err error) { + r1, _, e1 := syscall.Syscall(procHeapFree.Addr(), 3, uintptr(h), uintptr(flags), uintptr(p)) + if r1 == 0 { + if e1 != 0 { + err = errnoErr(e1) + } else { + err = syscall.EINVAL + } + } + return +}