Fix a memory leak

3 files changed, 31 insertions, 10 deletions

diff --git a/callback.go b/callback.go
index f2d7210..672a36e 100644
--- a/callback.go
+++ b/callback.go
@@ -8,12 +8,13 @@ var cb struct {
 	c int
 }
 
+func init() {
+	cb.m = make(map[int]interface{})
+}
+
 func cbAdd(v interface{}) int {
 	cb.Lock()
 	defer cb.Unlock()
-	if cb.m == nil {
-		cb.m = make(map[int]interface{})
-	}
 	cb.c++
 	cb.m[cb.c] = v
 	return cb.c
@@ -22,9 +23,17 @@ func cbAdd(v interface{}) int {
 func cbGet(c int) interface{} {
 	cb.Lock()
 	defer cb.Unlock()
-	v := cb.m[c]
-	if v == nil {
+	if v, ok := cb.m[c]; ok {
+		return v
+	}
+	panic("Callback pointer not found")
+}
+
+func cbDelete(c int) {
+	cb.Lock()
+	defer cb.Unlock()
+	if _, ok := cb.m[c]; !ok {
 		panic("Callback pointer not found")
 	}
-	return v
+	delete(cb.m, c)
 }
diff --git a/callback_test.go b/callback_test.go
index ae1eefe..45a84be 100644
--- a/callback_test.go
+++ b/callback_test.go
@@ -11,6 +11,7 @@ func TestCallback_001(t *testing.T) {
 	if reflect.TypeOf(v) != reflect.TypeOf(TestCallback_001) {
 		t.Error("Received unexpected value")
 	}
+	cbDelete(c)
 }
 
 func TestCallback_002(t *testing.T) {
@@ -21,3 +22,12 @@ func TestCallback_002(t *testing.T) {
 	cbGet(c + 1)
 	t.Error("Expected a panic")
 }
+
+func TestCallback_003(t *testing.T) {
+	defer func() {
+		recover()
+	}()
+	c := cbAdd(TestCallback_003)
+	cbDelete(c + 1)
+	t.Error("Expected a panic")
+}
diff --git a/transaction.go b/transaction.go
index f136abd..458804d 100644
--- a/transaction.go
+++ b/transaction.go
@@ -52,11 +52,11 @@ func (f ConversationFunc) RespondPAM(s Style, msg string) (string, error) {
 
 // Constructs a new conversation object with a given handler and a newly
 // allocated pam_conv struct that uses this object as its appdata_ptr.
-func newConversation(handler ConversationHandler) (*C.struct_pam_conv, C.int) {
+func newConversation(handler ConversationHandler) (*C.struct_pam_conv, int, C.int) {
 	c := cbAdd(handler)
 	conv := &C.struct_pam_conv{}
 	C.init_pam_conv(conv, C.long(c))
-	return conv, C.PAM_SUCCESS
+	return conv, c, C.PAM_SUCCESS
 }
 
 // Go-side function for processing a single conversational message. Ultimately
@@ -86,11 +86,13 @@ type Transaction struct {
 	handle *C.pam_handle_t
 	conv   *C.struct_pam_conv
 	status C.int
+	c      int
 }
 
 // Finalize a PAM transaction.
 func transactionFinalizer(t *Transaction) {
 	C.pam_end(t.handle, t.status)
+	cbDelete(t.c)
 }
 
 // Start initiates a new PAM transaction. Service is treated identically to
@@ -100,10 +102,11 @@ func transactionFinalizer(t *Transaction) {
 // transaction provides an interface to the remainder of the API.
 func Start(service, user string, handler ConversationHandler) (*Transaction, error) {
 	t := &Transaction{}
-	t.conv, t.status = newConversation(handler)
+	t.conv, t.c, t.status = newConversation(handler)
 	if t.status != C.PAM_SUCCESS {
 		return nil, t
 	}
+	runtime.SetFinalizer(t, transactionFinalizer)
 	s := C.CString(service)
 	defer C.free(unsafe.Pointer(s))
 	var u *C.char
@@ -115,7 +118,6 @@ func Start(service, user string, handler ConversationHandler) (*Transaction, err
 	if t.status != C.PAM_SUCCESS {
 		return nil, t
 	}
-	runtime.SetFinalizer(t, transactionFinalizer)
 	return t, nil
 }