regstate.go 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288
  1. package regstate
  2. import (
  3. "encoding/json"
  4. "fmt"
  5. "net/url"
  6. "os"
  7. "path/filepath"
  8. "reflect"
  9. "syscall"
  10. "golang.org/x/sys/windows/registry"
  11. )
  12. //go:generate go run $GOROOT/src/syscall/mksyscall_windows.go -output zsyscall_windows.go regstate.go
  13. //sys regCreateKeyEx(key syscall.Handle, subkey *uint16, reserved uint32, class *uint16, options uint32, desired uint32, sa *syscall.SecurityAttributes, result *syscall.Handle, disposition *uint32) (regerrno error) = advapi32.RegCreateKeyExW
  14. const (
  15. _REG_OPTION_VOLATILE = 1
  16. _REG_CREATED_NEW_KEY = 1
  17. _REG_OPENED_EXISTING_KEY = 2
  18. )
  19. type Key struct {
  20. registry.Key
  21. Name string
  22. }
  23. var localMachine = &Key{registry.LOCAL_MACHINE, "HKEY_LOCAL_MACHINE"}
  24. var localUser = &Key{registry.CURRENT_USER, "HKEY_CURRENT_USER"}
  25. var rootPath = `SOFTWARE\Microsoft\runhcs`
  26. type NotFoundError struct {
  27. Id string
  28. }
  29. func (err *NotFoundError) Error() string {
  30. return fmt.Sprintf("ID '%s' was not found", err.Id)
  31. }
  32. func IsNotFoundError(err error) bool {
  33. _, ok := err.(*NotFoundError)
  34. return ok
  35. }
  36. type NoStateError struct {
  37. ID string
  38. Key string
  39. }
  40. func (err *NoStateError) Error() string {
  41. return fmt.Sprintf("state '%s' is not present for ID '%s'", err.Key, err.ID)
  42. }
  43. func createVolatileKey(k *Key, path string, access uint32) (newk *Key, openedExisting bool, err error) {
  44. var (
  45. h syscall.Handle
  46. d uint32
  47. )
  48. fullpath := filepath.Join(k.Name, path)
  49. err = regCreateKeyEx(syscall.Handle(k.Key), syscall.StringToUTF16Ptr(path), 0, nil, _REG_OPTION_VOLATILE, access, nil, &h, &d)
  50. if err != nil {
  51. return nil, false, &os.PathError{Op: "RegCreateKeyEx", Path: fullpath, Err: err}
  52. }
  53. return &Key{registry.Key(h), fullpath}, d == _REG_OPENED_EXISTING_KEY, nil
  54. }
  55. func hive(perUser bool) *Key {
  56. r := localMachine
  57. if perUser {
  58. r = localUser
  59. }
  60. return r
  61. }
  62. func Open(root string, perUser bool) (*Key, error) {
  63. k, _, err := createVolatileKey(hive(perUser), rootPath, registry.ALL_ACCESS)
  64. if err != nil {
  65. return nil, err
  66. }
  67. defer k.Close()
  68. k2, _, err := createVolatileKey(k, url.PathEscape(root), registry.ALL_ACCESS)
  69. if err != nil {
  70. return nil, err
  71. }
  72. return k2, nil
  73. }
  74. func RemoveAll(root string, perUser bool) error {
  75. k, err := hive(perUser).open(rootPath)
  76. if err != nil {
  77. return err
  78. }
  79. defer k.Close()
  80. r, err := k.open(url.PathEscape(root))
  81. if err != nil {
  82. return err
  83. }
  84. defer r.Close()
  85. ids, err := r.Enumerate()
  86. if err != nil {
  87. return err
  88. }
  89. for _, id := range ids {
  90. err = r.Remove(id)
  91. if err != nil {
  92. return err
  93. }
  94. }
  95. r.Close()
  96. return k.Remove(root)
  97. }
  98. func (k *Key) Close() error {
  99. err := k.Key.Close()
  100. k.Key = 0
  101. return err
  102. }
  103. func (k *Key) Enumerate() ([]string, error) {
  104. escapedIDs, err := k.ReadSubKeyNames(0)
  105. if err != nil {
  106. return nil, err
  107. }
  108. var ids []string
  109. for _, e := range escapedIDs {
  110. id, err := url.PathUnescape(e)
  111. if err == nil {
  112. ids = append(ids, id)
  113. }
  114. }
  115. return ids, nil
  116. }
  117. func (k *Key) open(name string) (*Key, error) {
  118. fullpath := filepath.Join(k.Name, name)
  119. nk, err := registry.OpenKey(k.Key, name, registry.ALL_ACCESS)
  120. if err != nil {
  121. return nil, &os.PathError{Op: "RegOpenKey", Path: fullpath, Err: err}
  122. }
  123. return &Key{nk, fullpath}, nil
  124. }
  125. func (k *Key) openid(id string) (*Key, error) {
  126. escaped := url.PathEscape(id)
  127. fullpath := filepath.Join(k.Name, escaped)
  128. nk, err := k.open(escaped)
  129. if perr, ok := err.(*os.PathError); ok && perr.Err == syscall.ERROR_FILE_NOT_FOUND {
  130. return nil, &NotFoundError{id}
  131. }
  132. if err != nil {
  133. return nil, &os.PathError{Op: "RegOpenKey", Path: fullpath, Err: err}
  134. }
  135. return nk, nil
  136. }
  137. func (k *Key) Remove(id string) error {
  138. escaped := url.PathEscape(id)
  139. err := registry.DeleteKey(k.Key, escaped)
  140. if err != nil {
  141. if err == syscall.ERROR_FILE_NOT_FOUND {
  142. return &NotFoundError{id}
  143. }
  144. return &os.PathError{Op: "RegDeleteKey", Path: filepath.Join(k.Name, escaped), Err: err}
  145. }
  146. return nil
  147. }
  148. func (k *Key) set(id string, create bool, key string, state interface{}) error {
  149. var sk *Key
  150. var err error
  151. if create {
  152. var existing bool
  153. eid := url.PathEscape(id)
  154. sk, existing, err = createVolatileKey(k, eid, registry.ALL_ACCESS)
  155. if err != nil {
  156. return err
  157. }
  158. defer sk.Close()
  159. if existing {
  160. sk.Close()
  161. return fmt.Errorf("container %s already exists", id)
  162. }
  163. } else {
  164. sk, err = k.openid(id)
  165. if err != nil {
  166. return err
  167. }
  168. defer sk.Close()
  169. }
  170. switch reflect.TypeOf(state).Kind() {
  171. case reflect.Bool:
  172. v := uint32(0)
  173. if state.(bool) {
  174. v = 1
  175. }
  176. err = sk.SetDWordValue(key, v)
  177. case reflect.Int:
  178. err = sk.SetQWordValue(key, uint64(state.(int)))
  179. case reflect.String:
  180. err = sk.SetStringValue(key, state.(string))
  181. default:
  182. var js []byte
  183. js, err = json.Marshal(state)
  184. if err != nil {
  185. return err
  186. }
  187. err = sk.SetBinaryValue(key, js)
  188. }
  189. if err != nil {
  190. if err == syscall.ERROR_FILE_NOT_FOUND {
  191. return &NoStateError{id, key}
  192. }
  193. return &os.PathError{Op: "RegSetValueEx", Path: sk.Name + ":" + key, Err: err}
  194. }
  195. return nil
  196. }
  197. func (k *Key) Create(id, key string, state interface{}) error {
  198. return k.set(id, true, key, state)
  199. }
  200. func (k *Key) Set(id, key string, state interface{}) error {
  201. return k.set(id, false, key, state)
  202. }
  203. func (k *Key) Clear(id, key string) error {
  204. sk, err := k.openid(id)
  205. if err != nil {
  206. return err
  207. }
  208. defer sk.Close()
  209. err = sk.DeleteValue(key)
  210. if err != nil {
  211. if err == syscall.ERROR_FILE_NOT_FOUND {
  212. return &NoStateError{id, key}
  213. }
  214. return &os.PathError{Op: "RegDeleteValue", Path: sk.Name + ":" + key, Err: err}
  215. }
  216. return nil
  217. }
  218. func (k *Key) Get(id, key string, state interface{}) error {
  219. sk, err := k.openid(id)
  220. if err != nil {
  221. return err
  222. }
  223. defer sk.Close()
  224. var js []byte
  225. switch reflect.TypeOf(state).Elem().Kind() {
  226. case reflect.Bool:
  227. var v uint64
  228. v, _, err = sk.GetIntegerValue(key)
  229. if err == nil {
  230. *state.(*bool) = v != 0
  231. }
  232. case reflect.Int:
  233. var v uint64
  234. v, _, err = sk.GetIntegerValue(key)
  235. if err == nil {
  236. *state.(*int) = int(v)
  237. }
  238. case reflect.String:
  239. var v string
  240. v, _, err = sk.GetStringValue(key)
  241. if err == nil {
  242. *state.(*string) = string(v)
  243. }
  244. default:
  245. js, _, err = sk.GetBinaryValue(key)
  246. }
  247. if err != nil {
  248. if err == syscall.ERROR_FILE_NOT_FOUND {
  249. return &NoStateError{id, key}
  250. }
  251. return &os.PathError{Op: "RegQueryValueEx", Path: sk.Name + ":" + key, Err: err}
  252. }
  253. if js != nil {
  254. err = json.Unmarshal(js, state)
  255. }
  256. return err
  257. }