From f43007c55e11180bd3b206b901aa4a4df6e6630c Mon Sep 17 00:00:00 2001
From: Derek Parker <parkerderek86@gmail.com>
Date: Mon, 10 Aug 2015 08:55:57 -0500
Subject: [PATCH] Fix: Handle inability to find return addr

Fixes a code path where stacktrace returns < 2 locations and
thread.ReturnAddress would panic. Now returns an error.
---
 _fixtures/testreturnaddress.go | 15 ++++++++++
 proc/proc_test.go              | 50 ++++++++++++++--------------------
 proc/stack.go                  | 11 ++++++--
 3 files changed, 45 insertions(+), 31 deletions(-)
 create mode 100644 _fixtures/testreturnaddress.go

diff --git a/_fixtures/testreturnaddress.go b/_fixtures/testreturnaddress.go
new file mode 100644
index 00000000..d514396c
--- /dev/null
+++ b/_fixtures/testreturnaddress.go
@@ -0,0 +1,15 @@
+package main
+
+import "sync"
+
+func doSomething(wg *sync.WaitGroup) {
+	wg.Done()
+}
+
+func main() {
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go doSomething(&wg)
+	go doSomething(&wg)
+	wg.Wait()
+}
diff --git a/proc/proc_test.go b/proc/proc_test.go
index b5d65b40..4f7225c8 100644
--- a/proc/proc_test.go
+++ b/proc/proc_test.go
@@ -2,7 +2,6 @@ package proc
 
 import (
 	"bytes"
-	"encoding/binary"
 	"fmt"
 	"net"
 	"net/http"
@@ -380,48 +379,22 @@ func TestRuntimeBreakpoint(t *testing.T) {
 
 func TestFindReturnAddress(t *testing.T) {
 	withTestProcess("testnextprog", t, func(p *Process, fixture protest.Fixture) {
-		var (
-			fdes = p.frameEntries
-			gsd  = p.goSymTable
-		)
-
-		start, _, err := gsd.LineToPC(fixture.Source, 24)
+		start, _, err := p.goSymTable.LineToPC(fixture.Source, 24)
 		if err != nil {
 			t.Fatal(err)
 		}
-
 		_, err = p.SetBreakpoint(start)
 		if err != nil {
 			t.Fatal(err)
 		}
-
 		err = p.Continue()
 		if err != nil {
 			t.Fatal(err)
 		}
-
-		regs, err := p.Registers()
+		addr, err := p.CurrentThread.ReturnAddress()
 		if err != nil {
 			t.Fatal(err)
 		}
-
-		fde, err := fdes.FDEForPC(start)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		ret := fde.ReturnAddressOffset(start)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		addr := uint64(int64(regs.SP()) + ret)
-		data, err := p.CurrentThread.readMemory(uintptr(addr), 8)
-		if err != nil {
-			t.Fatal(err)
-		}
-		addr = binary.LittleEndian.Uint64(data)
-
 		_, l, _ := p.goSymTable.PCToLine(addr)
 		if l != 40 {
 			t.Fatalf("return address not found correctly, expected line 40")
@@ -429,6 +402,25 @@ func TestFindReturnAddress(t *testing.T) {
 	})
 }
 
+func TestFindReturnAddressTopOfStackFn(t *testing.T) {
+	withTestProcess("testreturnaddress", t, func(p *Process, fixture protest.Fixture) {
+		fnName := "runtime.rt0_go"
+		fn := p.goSymTable.LookupFunc(fnName)
+		if fn == nil {
+			t.Fatalf("could not find function %s", fnName)
+		}
+		if _, err := p.SetBreakpoint(fn.Entry); err != nil {
+			t.Fatal(err)
+		}
+		if err := p.Continue(); err != nil {
+			t.Fatal(err)
+		}
+		if _, err := p.CurrentThread.ReturnAddress(); err == nil {
+			t.Fatal("expected error to be returned")
+		}
+	})
+}
+
 func TestSwitchThread(t *testing.T) {
 	withTestProcess("testnextprog", t, func(p *Process, fixture protest.Fixture) {
 		// With invalid thread id
diff --git a/proc/stack.go b/proc/stack.go
index 05ed9cd6..2959d3f0 100644
--- a/proc/stack.go
+++ b/proc/stack.go
@@ -1,6 +1,9 @@
 package proc
 
-import "encoding/binary"
+import (
+	"encoding/binary"
+	"fmt"
+)
 
 // Takes an offset from RSP and returns the address of the
 // instruction the current function is going to return to.
@@ -9,6 +12,9 @@ func (thread *Thread) ReturnAddress() (uint64, error) {
 	if err != nil {
 		return 0, err
 	}
+	if len(locations) < 2 {
+		return 0, fmt.Errorf("could not find return address for %s", locations[0].Fn.BaseName())
+	}
 	return locations[1].PC, nil
 }
 
@@ -75,7 +81,8 @@ func (dbp *Process) stacktrace(pc, sp uint64, depth int) ([]Location, error) {
 			break
 		}
 		locations = append(locations, Location{PC: ret, File: f, Line: l, Fn: fn})
-		if fn.Name == "runtime.goexit" {
+		// Look for "top of stack" functions.
+		if fn.Name == "runtime.rt0_go" {
 			break
 		}
 	}