本站源代码
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

147 lines
4.2KB

  1. //
  2. // Copyright (c) 2014 David Mzareulyan
  3. //
  4. // Permission is hereby granted, free of charge, to any person obtaining a copy of this software
  5. // and associated documentation files (the "Software"), to deal in the Software without restriction,
  6. // including without limitation the rights to use, copy, modify, merge, publish, distribute,
  7. // sublicense, and/or sell copies of the Software, and to permit persons to whom the Software
  8. // is furnished to do so, subject to the following conditions:
  9. //
  10. // The above copyright notice and this permission notice shall be included in all copies or substantial
  11. // portions of the Software.
  12. //
  13. // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
  14. // BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  15. // NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
  16. // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  17. // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  18. //
  19. // +build windows
  20. package sshagent
  21. // see https://github.com/Yasushi/putty/blob/master/windows/winpgntc.c#L155
  22. // see https://github.com/paramiko/paramiko/blob/master/paramiko/win_pageant.py
  23. import (
  24. "encoding/binary"
  25. "errors"
  26. "fmt"
  27. "sync"
  28. "syscall"
  29. "unsafe"
  30. )
  31. // Maximum size of message can be sent to pageant
  32. const MaxMessageLen = 8192
  33. var (
  34. ErrPageantNotFound = errors.New("pageant process not found")
  35. ErrSendMessage = errors.New("error sending message")
  36. ErrMessageTooLong = errors.New("message too long")
  37. ErrInvalidMessageFormat = errors.New("invalid message format")
  38. ErrResponseTooLong = errors.New("response too long")
  39. )
  40. const (
  41. agentCopydataID = 0x804e50ba
  42. wmCopydata = 74
  43. )
  44. type copyData struct {
  45. dwData uintptr
  46. cbData uint32
  47. lpData unsafe.Pointer
  48. }
  49. var (
  50. lock sync.Mutex
  51. winFindWindow = winAPI("user32.dll", "FindWindowW")
  52. winGetCurrentThreadID = winAPI("kernel32.dll", "GetCurrentThreadId")
  53. winSendMessage = winAPI("user32.dll", "SendMessageW")
  54. )
  55. func winAPI(dllName, funcName string) func(...uintptr) (uintptr, uintptr, error) {
  56. proc := syscall.MustLoadDLL(dllName).MustFindProc(funcName)
  57. return func(a ...uintptr) (uintptr, uintptr, error) { return proc.Call(a...) }
  58. }
  59. // Available returns true if Pageant is running
  60. func Available() bool { return pageantWindow() != 0 }
  61. // Query sends message msg to Pageant and returns response or error.
  62. // 'msg' is raw agent request with length prefix
  63. // Response is raw agent response with length prefix
  64. func query(msg []byte) ([]byte, error) {
  65. if len(msg) > MaxMessageLen {
  66. return nil, ErrMessageTooLong
  67. }
  68. msgLen := binary.BigEndian.Uint32(msg[:4])
  69. if len(msg) != int(msgLen)+4 {
  70. return nil, ErrInvalidMessageFormat
  71. }
  72. lock.Lock()
  73. defer lock.Unlock()
  74. paWin := pageantWindow()
  75. if paWin == 0 {
  76. return nil, ErrPageantNotFound
  77. }
  78. thID, _, _ := winGetCurrentThreadID()
  79. mapName := fmt.Sprintf("PageantRequest%08x", thID)
  80. pMapName, _ := syscall.UTF16PtrFromString(mapName)
  81. mmap, err := syscall.CreateFileMapping(syscall.InvalidHandle, nil, syscall.PAGE_READWRITE, 0, MaxMessageLen+4, pMapName)
  82. if err != nil {
  83. return nil, err
  84. }
  85. defer syscall.CloseHandle(mmap)
  86. ptr, err := syscall.MapViewOfFile(mmap, syscall.FILE_MAP_WRITE, 0, 0, 0)
  87. if err != nil {
  88. return nil, err
  89. }
  90. defer syscall.UnmapViewOfFile(ptr)
  91. mmSlice := (*(*[MaxMessageLen]byte)(unsafe.Pointer(ptr)))[:]
  92. copy(mmSlice, msg)
  93. mapNameBytesZ := append([]byte(mapName), 0)
  94. cds := copyData{
  95. dwData: agentCopydataID,
  96. cbData: uint32(len(mapNameBytesZ)),
  97. lpData: unsafe.Pointer(&(mapNameBytesZ[0])),
  98. }
  99. resp, _, _ := winSendMessage(paWin, wmCopydata, 0, uintptr(unsafe.Pointer(&cds)))
  100. if resp == 0 {
  101. return nil, ErrSendMessage
  102. }
  103. respLen := binary.BigEndian.Uint32(mmSlice[:4])
  104. if respLen > MaxMessageLen-4 {
  105. return nil, ErrResponseTooLong
  106. }
  107. respData := make([]byte, respLen+4)
  108. copy(respData, mmSlice)
  109. return respData, nil
  110. }
  111. func pageantWindow() uintptr {
  112. nameP, _ := syscall.UTF16PtrFromString("Pageant")
  113. h, _, _ := winFindWindow(uintptr(unsafe.Pointer(nameP)), uintptr(unsafe.Pointer(nameP)))
  114. return h
  115. }
上海开阖软件有限公司 沪ICP备12045867号-1