diff --git a/agent/agent.go b/agent/agent.go index 07a3e80..d9cbaec 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -207,6 +207,14 @@ func formatBytes(val uint64) string { } func handle(conn io.ReadWriter, msg []byte) error { + reader := bufio.NewReader(conn) + magic, err := binary.ReadVarint(reader) + if err != nil { + return err + } + if magic != internal.MAGIC { + return fmt.Errorf("gops: invalid magic number: %v", magic) + } switch msg[0] { case signal.StackTrace: return pprof.Lookup("goroutine").WriteTo(conn, 2) @@ -286,7 +294,7 @@ func handle(conn io.ReadWriter, msg []byte) error { time.Sleep(5 * time.Second) trace.Stop() case signal.SetGCPercent: - perc, err := binary.ReadVarint(bufio.NewReader(conn)) + perc, err := binary.ReadVarint(reader) if err != nil { return err } diff --git a/internal/cmd/shared.go b/internal/cmd/shared.go index fc11e6b..8e54ca1 100644 --- a/internal/cmd/shared.go +++ b/internal/cmd/shared.go @@ -306,6 +306,9 @@ func cmdLazy(addr net.TCPAddr, c byte, params ...byte) (io.Reader, error) { return nil, err } buf := []byte{c} + magic := make([]byte, binary.MaxVarintLen64) + binary.PutVarint(magic, internal.MAGIC) + buf = append(buf, magic...) buf = append(buf, params...) if _, err := conn.Write(buf); err != nil { return nil, err diff --git a/internal/internal.go b/internal/internal.go index 7e3492a..376b1d7 100644 --- a/internal/internal.go +++ b/internal/internal.go @@ -13,7 +13,10 @@ import ( "strings" ) -const gopsConfigDirEnvKey = "GOPS_CONFIG_DIR" +const ( + goipsConfigDirEnvKey = "GOPS_CONFIG_DIR" + MAGIC = 0xabcdefedcba +) func ConfigDir() (string, error) { if configDir := os.Getenv(gopsConfigDirEnvKey); configDir != "" {