packer-cn/vendor/github.com/mitchellh/panicwrap/panicwrap.go

310 lines
8.3 KiB
Go
Raw Normal View History

// 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"
"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
// 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,
})
}
// 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 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, os.Interrupt)
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 {
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)
panicHeaders := [][]byte{
[]byte("panic:"),
[]byte("fatal error: fault"),
}
panicType := -1
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(panicHeaders[panicType])))
// 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
}
panicType = -1
flushIdx := n
for i, header := range panicHeaders {
idx := bytes.Index(buf[0:n], header)
if idx >= 0 {
panicType = i
flushIdx = idx
break
}
}
// Flush to stderr what isn't a panic
w.Write(buf[0:flushIdx])
if panicType == -1 {
// Not a panic so just continue along
continue
}
// We have a panic header. Write we assume is a panic os far.
panicBuf.Write(buf[flushIdx:n])
panicTimer = time.After(dur)
}
}