pool_sanitize.go 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121
  1. // +build pool_sanitize
  2. package pbytes
  3. import (
  4. "reflect"
  5. "runtime"
  6. "sync/atomic"
  7. "syscall"
  8. "unsafe"
  9. "golang.org/x/sys/unix"
  10. )
  11. const magic = uint64(0x777742)
  12. type guard struct {
  13. magic uint64
  14. size int
  15. owners int32
  16. }
  17. const guardSize = int(unsafe.Sizeof(guard{}))
  18. type Pool struct {
  19. min, max int
  20. }
  21. func New(min, max int) *Pool {
  22. return &Pool{min, max}
  23. }
  24. // Get returns probably reused slice of bytes with at least capacity of c and
  25. // exactly len of n.
  26. func (p *Pool) Get(n, c int) []byte {
  27. if n > c {
  28. panic("requested length is greater than capacity")
  29. }
  30. pageSize := syscall.Getpagesize()
  31. pages := (c+guardSize)/pageSize + 1
  32. size := pages * pageSize
  33. bts := alloc(size)
  34. g := (*guard)(unsafe.Pointer(&bts[0]))
  35. *g = guard{
  36. magic: magic,
  37. size: size,
  38. owners: 1,
  39. }
  40. return bts[guardSize : guardSize+n]
  41. }
  42. func (p *Pool) GetCap(c int) []byte { return p.Get(0, c) }
  43. func (p *Pool) GetLen(n int) []byte { return Get(n, n) }
  44. // Put returns given slice to reuse pool.
  45. func (p *Pool) Put(bts []byte) {
  46. hdr := *(*reflect.SliceHeader)(unsafe.Pointer(&bts))
  47. ptr := hdr.Data - uintptr(guardSize)
  48. g := (*guard)(unsafe.Pointer(ptr))
  49. if g.magic != magic {
  50. panic("unknown slice returned to the pool")
  51. }
  52. if n := atomic.AddInt32(&g.owners, -1); n < 0 {
  53. panic("multiple Put() detected")
  54. }
  55. // Disable read and write on bytes memory pages. This will cause panic on
  56. // incorrect access to returned slice.
  57. mprotect(ptr, false, false, g.size)
  58. runtime.SetFinalizer(&bts, func(b *[]byte) {
  59. mprotect(ptr, true, true, g.size)
  60. free(*(*[]byte)(unsafe.Pointer(&reflect.SliceHeader{
  61. Data: ptr,
  62. Len: g.size,
  63. Cap: g.size,
  64. })))
  65. })
  66. }
  67. func alloc(n int) []byte {
  68. b, err := unix.Mmap(-1, 0, n, unix.PROT_READ|unix.PROT_WRITE|unix.PROT_EXEC, unix.MAP_SHARED|unix.MAP_ANONYMOUS)
  69. if err != nil {
  70. panic(err.Error())
  71. }
  72. return b
  73. }
  74. func free(b []byte) {
  75. if err := unix.Munmap(b); err != nil {
  76. panic(err.Error())
  77. }
  78. }
  79. func mprotect(ptr uintptr, r, w bool, size int) {
  80. // Need to avoid "EINVAL addr is not a valid pointer,
  81. // or not a multiple of PAGESIZE."
  82. start := ptr & ^(uintptr(syscall.Getpagesize() - 1))
  83. prot := uintptr(syscall.PROT_EXEC)
  84. switch {
  85. case r && w:
  86. prot |= syscall.PROT_READ | syscall.PROT_WRITE
  87. case r:
  88. prot |= syscall.PROT_READ
  89. case w:
  90. prot |= syscall.PROT_WRITE
  91. }
  92. _, _, err := syscall.Syscall(syscall.SYS_MPROTECT,
  93. start, uintptr(size), prot,
  94. )
  95. if err != 0 {
  96. panic(err.Error())
  97. }
  98. }