1
0

panicwrap.go 9.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339
  1. // The panicwrap package provides functions for capturing and handling
  2. // panics in your application. It does this by re-executing the running
  3. // application and monitoring stderr for any panics. At the same time,
  4. // stdout/stderr/etc. are set to the same values so that data is shuttled
  5. // through properly, making the existence of panicwrap mostly transparent.
  6. //
  7. // Panics are only detected when the subprocess exits with a non-zero
  8. // exit status, since this is the only time panics are real. Otherwise,
  9. // "panic-like" output is ignored.
  10. package panicwrap
  11. import (
  12. "bytes"
  13. "errors"
  14. "github.com/kardianos/osext"
  15. "io"
  16. "os"
  17. "os/exec"
  18. "os/signal"
  19. "runtime"
  20. "syscall"
  21. "time"
  22. )
  23. const (
  24. DEFAULT_COOKIE_KEY = "cccf35992f8f3cd8d1d28f0109dd953e26664531"
  25. DEFAULT_COOKIE_VAL = "7c28215aca87789f95b406b8dd91aa5198406750"
  26. )
  27. // HandlerFunc is the type called when a panic is detected.
  28. type HandlerFunc func(string)
  29. // WrapConfig is the configuration for panicwrap when wrapping an existing
  30. // binary. To get started, in general, you only need the BasicWrap function
  31. // that will set this up for you. However, for more customizability,
  32. // WrapConfig and Wrap can be used.
  33. type WrapConfig struct {
  34. // Handler is the function called when a panic occurs.
  35. Handler HandlerFunc
  36. // The cookie key and value are used within environmental variables
  37. // to tell the child process that it is already executing so that
  38. // wrap doesn't re-wrap itself.
  39. CookieKey string
  40. CookieValue string
  41. // If true, the panic will not be mirrored to the configured writer
  42. // and will instead ONLY go to the handler. This lets you effectively
  43. // hide panics from the end user. This is not recommended because if
  44. // your handler fails, the panic is effectively lost.
  45. HidePanic bool
  46. // If true, panicwrap will boot a monitor sub-process and let the parent
  47. // run the app. This mode is useful for processes run under supervisors
  48. // like runit as signals get sent to the correct codebase. This is not
  49. // supported when GOOS=windows, and ignores c.Stderr and c.Stdout.
  50. Monitor bool
  51. // The amount of time that a process must exit within after detecting
  52. // a panic header for panicwrap to assume it is a panic. Defaults to
  53. // 300 milliseconds.
  54. DetectDuration time.Duration
  55. // The writer to send the stderr to. If this is nil, then it defaults
  56. // to os.Stderr.
  57. Writer io.Writer
  58. // The writer to send stdout to. If this is nil, then it defaults to
  59. // os.Stdout.
  60. Stdout io.Writer
  61. }
  62. // BasicWrap calls Wrap with the given handler function, using defaults
  63. // for everything else. See Wrap and WrapConfig for more information on
  64. // functionality and return values.
  65. func BasicWrap(f HandlerFunc) (int, error) {
  66. return Wrap(&WrapConfig{
  67. Handler: f,
  68. })
  69. }
  70. // BasicMonitor calls Wrap with Monitor set to true on supported platforms.
  71. // It forks your program and runs it again form the start. In one process
  72. // BasicMonitor never returns, it just listens on stderr of the other process,
  73. // and calls your handler when a panic is seen. In the other it either returns
  74. // nil to indicate that the panic monitoring is enabled, or an error to indicate
  75. // that something else went wrong.
  76. func BasicMonitor(f HandlerFunc) error {
  77. exitStatus, err := Wrap(&WrapConfig{
  78. Handler: f,
  79. Monitor: runtime.GOOS != "windows",
  80. })
  81. if err != nil {
  82. return err
  83. }
  84. if exitStatus >= 0 {
  85. os.Exit(exitStatus)
  86. }
  87. return nil
  88. }
  89. // Wrap wraps the current executable in a handler to catch panics. It
  90. // returns an error if there was an error during the wrapping process.
  91. // If the error is nil, then the int result indicates the exit status of the
  92. // child process. If the exit status is -1, then this is the child process,
  93. // and execution should continue as normal. Otherwise, this is the parent
  94. // process and the child successfully ran already, and you should exit the
  95. // process with the returned exit status.
  96. //
  97. // This function should be called very very early in your program's execution.
  98. // Ideally, this runs as the first line of code of main.
  99. //
  100. // Once this is called, the given WrapConfig shouldn't be modified or used
  101. // any further.
  102. func Wrap(c *WrapConfig) (int, error) {
  103. if c.Handler == nil {
  104. return -1, errors.New("Handler must be set")
  105. }
  106. if c.DetectDuration == 0 {
  107. c.DetectDuration = 300 * time.Millisecond
  108. }
  109. if c.Writer == nil {
  110. c.Writer = os.Stderr
  111. }
  112. if c.Monitor {
  113. return monitor(c)
  114. } else {
  115. return wrap(c)
  116. }
  117. }
  118. func wrap(c *WrapConfig) (int, error) {
  119. // If we're already wrapped, exit out.
  120. if Wrapped(c) {
  121. return -1, nil
  122. }
  123. // Get the path to our current executable
  124. exePath, err := osext.Executable()
  125. if err != nil {
  126. return -1, err
  127. }
  128. // Pipe the stderr so we can read all the data as we look for panics
  129. stderr_r, stderr_w := io.Pipe()
  130. // doneCh is closed when we're done, signaling any other goroutines
  131. // to end immediately.
  132. doneCh := make(chan struct{})
  133. // panicCh is the channel on which the panic text will actually be
  134. // sent.
  135. panicCh := make(chan string)
  136. // On close, make sure to finish off the copying of data to stderr
  137. defer func() {
  138. defer close(doneCh)
  139. stderr_w.Close()
  140. <-panicCh
  141. }()
  142. // Start the goroutine that will watch stderr for any panics
  143. go trackPanic(stderr_r, c.Writer, c.DetectDuration, panicCh)
  144. // Create the writer for stdout that we're going to use
  145. var stdout_w io.Writer = os.Stdout
  146. if c.Stdout != nil {
  147. stdout_w = c.Stdout
  148. }
  149. // Build a subcommand to re-execute ourselves. We make sure to
  150. // set the environmental variable to include our cookie. We also
  151. // set stdin/stdout to match the config. Finally, we pipe stderr
  152. // through ourselves in order to watch for panics.
  153. cmd := exec.Command(exePath, os.Args[1:]...)
  154. cmd.Env = append(os.Environ(), c.CookieKey+"="+c.CookieValue)
  155. cmd.Stdin = os.Stdin
  156. cmd.Stdout = stdout_w
  157. cmd.Stderr = stderr_w
  158. if err := cmd.Start(); err != nil {
  159. return 1, err
  160. }
  161. // Listen to signals and capture them forever. We allow the child
  162. // process to handle them in some way.
  163. sigCh := make(chan os.Signal)
  164. signal.Notify(sigCh, signalsToIgnore...)
  165. go func() {
  166. defer signal.Stop(sigCh)
  167. for {
  168. select {
  169. case <-doneCh:
  170. return
  171. case <-sigCh:
  172. }
  173. }
  174. }()
  175. if err := cmd.Wait(); err != nil {
  176. exitErr, ok := err.(*exec.ExitError)
  177. if !ok {
  178. // This is some other kind of subprocessing error.
  179. return 1, err
  180. }
  181. exitStatus := 1
  182. if status, ok := exitErr.Sys().(syscall.WaitStatus); ok && status.Exited() {
  183. exitStatus = status.ExitStatus()
  184. }
  185. // Close the writer end so that the tracker goroutine ends at some point
  186. stderr_w.Close()
  187. // Wait on the panic data
  188. panicTxt := <-panicCh
  189. if panicTxt != "" {
  190. if !c.HidePanic {
  191. c.Writer.Write([]byte(panicTxt))
  192. }
  193. c.Handler(panicTxt)
  194. }
  195. return exitStatus, nil
  196. }
  197. return 0, nil
  198. }
  199. // Wrapped checks if we're already wrapped according to the configuration
  200. // given.
  201. //
  202. // Wrapped is very cheap and can be used early to short-circuit some pre-wrap
  203. // logic your application may have.
  204. func Wrapped(c *WrapConfig) bool {
  205. if c.CookieKey == "" {
  206. c.CookieKey = DEFAULT_COOKIE_KEY
  207. }
  208. if c.CookieValue == "" {
  209. c.CookieValue = DEFAULT_COOKIE_VAL
  210. }
  211. // If the cookie key/value match our environment, then we are the
  212. // child, so just exit now and tell the caller that we're the child
  213. return os.Getenv(c.CookieKey) == c.CookieValue
  214. }
  215. // trackPanic monitors the given reader for a panic. If a panic is detected,
  216. // it is outputted on the result channel. This will close the channel once
  217. // it is complete.
  218. func trackPanic(r io.Reader, w io.Writer, dur time.Duration, result chan<- string) {
  219. defer close(result)
  220. var panicTimer <-chan time.Time
  221. panicBuf := new(bytes.Buffer)
  222. panicHeader := []byte("panic:")
  223. tempBuf := make([]byte, 2048)
  224. for {
  225. var buf []byte
  226. var n int
  227. if panicTimer == nil && panicBuf.Len() > 0 {
  228. // We're not tracking a panic but the buffer length is
  229. // greater than 0. We need to clear out that buffer, but
  230. // look for another panic along the way.
  231. // First, remove the previous panic header so we don't loop
  232. w.Write(panicBuf.Next(len(panicHeader)))
  233. // Next, assume that this is our new buffer to inspect
  234. n = panicBuf.Len()
  235. buf = make([]byte, n)
  236. copy(buf, panicBuf.Bytes())
  237. panicBuf.Reset()
  238. } else {
  239. var err error
  240. buf = tempBuf
  241. n, err = r.Read(buf)
  242. if n <= 0 && err == io.EOF {
  243. if panicBuf.Len() > 0 {
  244. // We were tracking a panic, assume it was a panic
  245. // and return that as the result.
  246. result <- panicBuf.String()
  247. }
  248. return
  249. }
  250. }
  251. if panicTimer != nil {
  252. // We're tracking what we think is a panic right now.
  253. // If the timer ended, then it is not a panic.
  254. isPanic := true
  255. select {
  256. case <-panicTimer:
  257. isPanic = false
  258. default:
  259. }
  260. // No matter what, buffer the text some more.
  261. panicBuf.Write(buf[0:n])
  262. if !isPanic {
  263. // It isn't a panic, stop tracking. Clean-up will happen
  264. // on the next iteration.
  265. panicTimer = nil
  266. }
  267. continue
  268. }
  269. flushIdx := n
  270. idx := bytes.Index(buf[0:n], panicHeader)
  271. if idx >= 0 {
  272. flushIdx = idx
  273. }
  274. // Flush to stderr what isn't a panic
  275. w.Write(buf[0:flushIdx])
  276. if idx < 0 {
  277. // Not a panic so just continue along
  278. continue
  279. }
  280. // We have a panic header. Write we assume is a panic os far.
  281. panicBuf.Write(buf[idx:n])
  282. panicTimer = time.After(dur)
  283. }
  284. }