Cookie + Exports: Store and return the cookie from the underlying context

This has the ability so that we can easily get the same cookie back in
the client package by creating the one from the context

2 files changed, 14 insertions, 4 deletions

diff --git a/exports/exports.go b/exports/exports.go
index 9c0088e..5c36446 100644
--- a/exports/exports.go
+++ b/exports/exports.go
@@ -1047,7 +1047,7 @@ func SetTokenHandler(getter C.TokenGetter, setter C.TokenSetter) *C.char {
 //export CookieNew
 func CookieNew() C.uintptr_t {
 	c := cookie.NewWithContext(context.Background())
-	return C.uintptr_t(cgo.NewHandle(&c))
+	return C.uintptr_t(cgo.NewHandle(c))
 }
 
 // CookieReply replies to a state transition using the cookie
diff --git a/types/cookie/cookie.go b/types/cookie/cookie.go
index 55fe938..1582a9c 100644
--- a/types/cookie/cookie.go
+++ b/types/cookie/cookie.go
@@ -19,11 +19,21 @@ type Cookie struct {
 	H         cgo.Handle
 }
 
+type contextt int8
+
+const CONTEXTK contextt = 0
+
 // NewWithContext creates a new cookie with a context
 // It stores the cancel and channel inside of the struct
-func NewWithContext(ctx context.Context) Cookie {
+func NewWithContext(ctx context.Context) *Cookie {
+	// if the context already has a handle, return that cookie
+	if h, ok := ctx.Value(CONTEXTK).(cgo.Handle); ok {
+		if ck, ok := h.Value().(*Cookie); ok {
+			return ck
+		}
+	}
 	ctx, cancel := context.WithCancel(ctx)
-	return Cookie{
+	return &Cookie{
 		c:         make(chan string),
 		ctx:       ctx,
 		ctxCancel: cancel,
@@ -77,5 +87,5 @@ func (c *Cookie) Send(data string) error {
 
 // Context gets the underlying context of the cookie
 func (c *Cookie) Context() context.Context {
-	return c.ctx
+	return context.WithValue(c.ctx, CONTEXTK, c.H)
 }