From 12bd4761fc66ac946e16fcc2a32b1e0b066f6177 Mon Sep 17 00:00:00 2001
From: secracon <secracon@gmail.com>
Date: Sat, 8 Dec 2018 22:45:19 +0100
Subject: [PATCH] Use golang context pkg instead of gorilla/context to fix
 memory leaks (#175)

* - use golang context pkg instead of gorilla/context to fix memory leaks
* - add test case for checking request context content upon shallow copy
* - update docs, readme.md and travis.yml
---
 .travis.yml      |  4 ----
 README.md        | 12 ------------
 doc.go           |  8 --------
 sessions.go      |  8 ++++----
 sessions_test.go | 31 +++++++++++++++++++++++++++++++
 5 files changed, 35 insertions(+), 28 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 899d4cb..85c28cd 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -3,10 +3,6 @@ sudo: false
 
 matrix:
   include:
-    - go: 1.3.x
-    - go: 1.4.x
-    - go: 1.5.x
-    - go: 1.6.x
     - go: 1.7.x
     - go: 1.8.x
     - go: 1.9.x
diff --git a/README.md b/README.md
index 00b43ca..98c993d 100644
--- a/README.md
+++ b/README.md
@@ -51,18 +51,6 @@ secret key used to authenticate the session. Inside the handler, we call
 some session values in session.Values, which is a `map[interface{}]interface{}`.
 And finally we call `session.Save()` to save the session in the response.
 
-Important Note: If you aren't using gorilla/mux, you need to wrap your handlers
-with
-[`context.ClearHandler`](https://www.gorillatoolkit.org/pkg/context#ClearHandler)
-or else you will leak memory! An easy way to do this is to wrap the top-level
-mux when calling http.ListenAndServe:
-
-```go
-	http.ListenAndServe(":8080", context.ClearHandler(http.DefaultServeMux))
-```
-
-The ClearHandler function is provided by the gorilla/context package.
-
 More examples are available [on the Gorilla
 website](https://www.gorillatoolkit.org/pkg/sessions).
 
diff --git a/doc.go b/doc.go
index 7db6729..64f858c 100644
--- a/doc.go
+++ b/doc.go
@@ -59,14 +59,6 @@ session.Save(r, w), and either display an error message or otherwise handle it.
 Save must be called before writing to the response, otherwise the session
 cookie will not be sent to the client.
 
-Important Note: If you aren't using gorilla/mux, you need to wrap your handlers
-with context.ClearHandler as or else you will leak memory! An easy way to do this
-is to wrap the top-level mux when calling http.ListenAndServe:
-
-    http.ListenAndServe(":8080", context.ClearHandler(http.DefaultServeMux))
-
-The ClearHandler function is provided by the gorilla/context package.
-
 That's all you need to know for the basic usage. Let's take a look at other
 options, starting with flash messages.
 
diff --git a/sessions.go b/sessions.go
index a821d31..c052b28 100644
--- a/sessions.go
+++ b/sessions.go
@@ -5,12 +5,11 @@
 package sessions
 
 import (
+	"context"
 	"encoding/gob"
 	"fmt"
 	"net/http"
 	"time"
-
-	"github.com/gorilla/context"
 )
 
 // Default flashes key.
@@ -108,7 +107,8 @@ const registryKey contextKey = 0
 
 // GetRegistry returns a registry instance for the current request.
 func GetRegistry(r *http.Request) *Registry {
-	registry := context.Get(r, registryKey)
+	var ctx = r.Context()
+	registry := ctx.Value(registryKey)
 	if registry != nil {
 		return registry.(*Registry)
 	}
@@ -116,7 +116,7 @@ func GetRegistry(r *http.Request) *Registry {
 		request:  r,
 		sessions: make(map[string]sessionInfo),
 	}
-	context.Set(r, registryKey, newRegistry)
+	*r = *r.WithContext(context.WithValue(ctx, registryKey, newRegistry))
 	return newRegistry
 }
 
diff --git a/sessions_test.go b/sessions_test.go
index 81fbf9a..a734f67 100644
--- a/sessions_test.go
+++ b/sessions_test.go
@@ -153,6 +153,37 @@ func TestFlashes(t *testing.T) {
 	if custom.Type != 42 || custom.Message != "foo" {
 		t.Errorf("Expected %#v, got %#v", FlashMessage{42, "foo"}, custom)
 	}
+
+	// Round 5 ----------------------------------------------------------------
+	// Check if a request shallow copy resets the request context data store.
+
+	req, _ = http.NewRequest("GET", "http://localhost:8080/", nil)
+
+	// Get a session.
+	if session, err = store.Get(req, "session-key"); err != nil {
+		t.Fatalf("Error getting session: %v", err)
+	}
+
+	// Put a test value into the session data store.
+	session.Values["test"] = "test-value"
+
+	// Create a shallow copy of the request.
+	req = req.WithContext(req.Context())
+
+	// Get the session again.
+	if session, err = store.Get(req, "session-key"); err != nil {
+		t.Fatalf("Error getting session: %v", err)
+	}
+
+	// Check if the previous inserted value still exists.
+	if session.Values["test"] == nil {
+		t.Fatalf("Session test value is lost in the request context!")
+	}
+
+	// Check if the previous inserted value has the same value.
+	if session.Values["test"] != "test-value" {
+		t.Fatalf("Session test value is changed in the request context!")
+	}
 }
 
 func TestCookieStoreMapPanic(t *testing.T) {