| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339 |
- // The panicwrap package provides functions for capturing and handling
- // panics in your application. It does this by re-executing the running
- // application and monitoring stderr for any panics. At the same time,
- // stdout/stderr/etc. are set to the same values so that data is shuttled
- // through properly, making the existence of panicwrap mostly transparent.
- //
- // Panics are only detected when the subprocess exits with a non-zero
- // exit status, since this is the only time panics are real. Otherwise,
- // "panic-like" output is ignored.
- package panicwrap
- import (
- "bytes"
- "errors"
- "github.com/kardianos/osext"
- "io"
- "os"
- "os/exec"
- "os/signal"
- "runtime"
- "syscall"
- "time"
- )
- const (
- DEFAULT_COOKIE_KEY = "cccf35992f8f3cd8d1d28f0109dd953e26664531"
- DEFAULT_COOKIE_VAL = "7c28215aca87789f95b406b8dd91aa5198406750"
- )
- // HandlerFunc is the type called when a panic is detected.
- type HandlerFunc func(string)
- // WrapConfig is the configuration for panicwrap when wrapping an existing
- // binary. To get started, in general, you only need the BasicWrap function
- // that will set this up for you. However, for more customizability,
- // WrapConfig and Wrap can be used.
- type WrapConfig struct {
- // Handler is the function called when a panic occurs.
- Handler HandlerFunc
- // The cookie key and value are used within environmental variables
- // to tell the child process that it is already executing so that
- // wrap doesn't re-wrap itself.
- CookieKey string
- CookieValue string
- // If true, the panic will not be mirrored to the configured writer
- // and will instead ONLY go to the handler. This lets you effectively
- // hide panics from the end user. This is not recommended because if
- // your handler fails, the panic is effectively lost.
- HidePanic bool
- // If true, panicwrap will boot a monitor sub-process and let the parent
- // run the app. This mode is useful for processes run under supervisors
- // like runit as signals get sent to the correct codebase. This is not
- // supported when GOOS=windows, and ignores c.Stderr and c.Stdout.
- Monitor bool
- // The amount of time that a process must exit within after detecting
- // a panic header for panicwrap to assume it is a panic. Defaults to
- // 300 milliseconds.
- DetectDuration time.Duration
- // The writer to send the stderr to. If this is nil, then it defaults
- // to os.Stderr.
- Writer io.Writer
- // The writer to send stdout to. If this is nil, then it defaults to
- // os.Stdout.
- Stdout io.Writer
- }
- // BasicWrap calls Wrap with the given handler function, using defaults
- // for everything else. See Wrap and WrapConfig for more information on
- // functionality and return values.
- func BasicWrap(f HandlerFunc) (int, error) {
- return Wrap(&WrapConfig{
- Handler: f,
- })
- }
- // BasicMonitor calls Wrap with Monitor set to true on supported platforms.
- // It forks your program and runs it again form the start. In one process
- // BasicMonitor never returns, it just listens on stderr of the other process,
- // and calls your handler when a panic is seen. In the other it either returns
- // nil to indicate that the panic monitoring is enabled, or an error to indicate
- // that something else went wrong.
- func BasicMonitor(f HandlerFunc) error {
- exitStatus, err := Wrap(&WrapConfig{
- Handler: f,
- Monitor: runtime.GOOS != "windows",
- })
- if err != nil {
- return err
- }
- if exitStatus >= 0 {
- os.Exit(exitStatus)
- }
- return nil
- }
- // Wrap wraps the current executable in a handler to catch panics. It
- // returns an error if there was an error during the wrapping process.
- // If the error is nil, then the int result indicates the exit status of the
- // child process. If the exit status is -1, then this is the child process,
- // and execution should continue as normal. Otherwise, this is the parent
- // process and the child successfully ran already, and you should exit the
- // process with the returned exit status.
- //
- // This function should be called very very early in your program's execution.
- // Ideally, this runs as the first line of code of main.
- //
- // Once this is called, the given WrapConfig shouldn't be modified or used
- // any further.
- func Wrap(c *WrapConfig) (int, error) {
- if c.Handler == nil {
- return -1, errors.New("Handler must be set")
- }
- if c.DetectDuration == 0 {
- c.DetectDuration = 300 * time.Millisecond
- }
- if c.Writer == nil {
- c.Writer = os.Stderr
- }
- if c.Monitor {
- return monitor(c)
- } else {
- return wrap(c)
- }
- }
- func wrap(c *WrapConfig) (int, error) {
- // If we're already wrapped, exit out.
- if Wrapped(c) {
- return -1, nil
- }
- // Get the path to our current executable
- exePath, err := osext.Executable()
- if err != nil {
- return -1, err
- }
- // Pipe the stderr so we can read all the data as we look for panics
- stderr_r, stderr_w := io.Pipe()
- // doneCh is closed when we're done, signaling any other goroutines
- // to end immediately.
- doneCh := make(chan struct{})
- // panicCh is the channel on which the panic text will actually be
- // sent.
- panicCh := make(chan string)
- // On close, make sure to finish off the copying of data to stderr
- defer func() {
- defer close(doneCh)
- stderr_w.Close()
- <-panicCh
- }()
- // Start the goroutine that will watch stderr for any panics
- go trackPanic(stderr_r, c.Writer, c.DetectDuration, panicCh)
- // Create the writer for stdout that we're going to use
- var stdout_w io.Writer = os.Stdout
- if c.Stdout != nil {
- stdout_w = c.Stdout
- }
- // Build a subcommand to re-execute ourselves. We make sure to
- // set the environmental variable to include our cookie. We also
- // set stdin/stdout to match the config. Finally, we pipe stderr
- // through ourselves in order to watch for panics.
- cmd := exec.Command(exePath, os.Args[1:]...)
- cmd.Env = append(os.Environ(), c.CookieKey+"="+c.CookieValue)
- cmd.Stdin = os.Stdin
- cmd.Stdout = stdout_w
- cmd.Stderr = stderr_w
- if err := cmd.Start(); err != nil {
- return 1, err
- }
- // Listen to signals and capture them forever. We allow the child
- // process to handle them in some way.
- sigCh := make(chan os.Signal)
- signal.Notify(sigCh, signalsToIgnore...)
- go func() {
- defer signal.Stop(sigCh)
- for {
- select {
- case <-doneCh:
- return
- case <-sigCh:
- }
- }
- }()
- if err := cmd.Wait(); err != nil {
- exitErr, ok := err.(*exec.ExitError)
- if !ok {
- // This is some other kind of subprocessing error.
- return 1, err
- }
- exitStatus := 1
- if status, ok := exitErr.Sys().(syscall.WaitStatus); ok && status.Exited() {
- exitStatus = status.ExitStatus()
- }
- // Close the writer end so that the tracker goroutine ends at some point
- stderr_w.Close()
- // Wait on the panic data
- panicTxt := <-panicCh
- if panicTxt != "" {
- if !c.HidePanic {
- c.Writer.Write([]byte(panicTxt))
- }
- c.Handler(panicTxt)
- }
- return exitStatus, nil
- }
- return 0, nil
- }
- // Wrapped checks if we're already wrapped according to the configuration
- // given.
- //
- // Wrapped is very cheap and can be used early to short-circuit some pre-wrap
- // logic your application may have.
- func Wrapped(c *WrapConfig) bool {
- if c.CookieKey == "" {
- c.CookieKey = DEFAULT_COOKIE_KEY
- }
- if c.CookieValue == "" {
- c.CookieValue = DEFAULT_COOKIE_VAL
- }
- // If the cookie key/value match our environment, then we are the
- // child, so just exit now and tell the caller that we're the child
- return os.Getenv(c.CookieKey) == c.CookieValue
- }
- // trackPanic monitors the given reader for a panic. If a panic is detected,
- // it is outputted on the result channel. This will close the channel once
- // it is complete.
- func trackPanic(r io.Reader, w io.Writer, dur time.Duration, result chan<- string) {
- defer close(result)
- var panicTimer <-chan time.Time
- panicBuf := new(bytes.Buffer)
- panicHeader := []byte("panic:")
- tempBuf := make([]byte, 2048)
- for {
- var buf []byte
- var n int
- if panicTimer == nil && panicBuf.Len() > 0 {
- // We're not tracking a panic but the buffer length is
- // greater than 0. We need to clear out that buffer, but
- // look for another panic along the way.
- // First, remove the previous panic header so we don't loop
- w.Write(panicBuf.Next(len(panicHeader)))
- // Next, assume that this is our new buffer to inspect
- n = panicBuf.Len()
- buf = make([]byte, n)
- copy(buf, panicBuf.Bytes())
- panicBuf.Reset()
- } else {
- var err error
- buf = tempBuf
- n, err = r.Read(buf)
- if n <= 0 && err == io.EOF {
- if panicBuf.Len() > 0 {
- // We were tracking a panic, assume it was a panic
- // and return that as the result.
- result <- panicBuf.String()
- }
- return
- }
- }
- if panicTimer != nil {
- // We're tracking what we think is a panic right now.
- // If the timer ended, then it is not a panic.
- isPanic := true
- select {
- case <-panicTimer:
- isPanic = false
- default:
- }
- // No matter what, buffer the text some more.
- panicBuf.Write(buf[0:n])
- if !isPanic {
- // It isn't a panic, stop tracking. Clean-up will happen
- // on the next iteration.
- panicTimer = nil
- }
- continue
- }
- flushIdx := n
- idx := bytes.Index(buf[0:n], panicHeader)
- if idx >= 0 {
- flushIdx = idx
- }
- // Flush to stderr what isn't a panic
- w.Write(buf[0:flushIdx])
- if idx < 0 {
- // Not a panic so just continue along
- continue
- }
- // We have a panic header. Write we assume is a panic os far.
- panicBuf.Write(buf[idx:n])
- panicTimer = time.After(dur)
- }
- }
|