Add SetToken method

Author: alexedwards

data.go

@@ -603,6 +603,20 @@ func (s *SessionManager) Token(ctx context.Context) string {
 	return sd.token
 }
 
+// SetToken changes the token for the session to a known value. Please take care
+// when using this function to ensure that the token you are setting is an
+// unguessable value from a trusted source. Most applications will not need to
+// use this method.
+func (s *SessionManager) SetToken(ctx context.Context, token string) {
+	sd := s.getSessionDataFromContext(ctx)
+
+	sd.mu.Lock()
+	defer sd.mu.Unlock()
+
+	sd.token = token
+	sd.status = Modified
+}
+
 func (s *SessionManager) addSessionDataToContext(ctx context.Context, sd *sessionData) context.Context {
 	return context.WithValue(ctx, s.ContextKey, sd)
 }

session_test.go

@@ -348,3 +348,39 @@ func TestIterate(t *testing.T) {
 		t.Fatal("didn't get expected error")
 	}
 }
+
+func TestSetToken(t *testing.T) {
+	t.Parallel()
+
+	sessionManager := New()
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/put", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		sessionManager.SetToken(r.Context(), "my-custom-token")
+		sessionManager.Put(r.Context(), "foo", "bar")
+	}))
+	mux.HandleFunc("/get", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		v := sessionManager.Get(r.Context(), "foo")
+		if v == nil {
+			http.Error(w, "foo does not exist in session", 500)
+			return
+		}
+		w.Write([]byte(v.(string)))
+	}))
+
+	ts := newTestServer(t, sessionManager.LoadAndSave(mux))
+	defer ts.Close()
+
+	header, _ := ts.execute(t, "/put")
+	cookie := header.Get("Set-Cookie")
+	token := extractTokenFromCookie(cookie)
+
+	if token != "my-custom-token" {
+		t.Errorf("want %q; got %q", "my-custom-token", token)
+	}
+
+	_, body := ts.execute(t, "/get")
+	if body != "bar" {
+		t.Errorf("want %q; got %q", "bar", body)
+	}
+}