From e073bf2775d020779d1463c4e020ee157b5d2f07 Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@abcum.com>
Date: Mon, 12 Feb 2018 21:53:48 +0000
Subject: [PATCH] Enable persistent request variables

It is now possible to set a variable on a connection, and use that variable in multiple subsequent queries. This ensures that websockets can make use of state variables for detecting logged-in, and non-logged-in clients.
---
 db/db.go      |  7 +++++++
 db/db_test.go | 10 ++++++++--
 db/vars.go    |  4 +++-
 web/auth.go   | 16 +++++++++++++++-
 web/rpc.go    |  5 +++++
 5 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/db/db.go b/db/db.go
index 379fe8bf..a343925c 100644
--- a/db/db.go
+++ b/db/db.go
@@ -207,6 +207,13 @@ func Process(fib *fibre.Context, ast *sql.Query, vars map[string]interface{}) (o
 
 	ctx = context.WithValue(ctx, ctxKeyId, id)
 
+	// Assign any global connection variables
+	// to the context so that we can retrieve
+	// the variables from within any queries.
+
+	keep := fib.Get(varKeyKeep).(*data.Doc)
+	ctx = context.WithValue(ctx, ctxKeyKeep, keep)
+
 	// Assign the authentication data to the
 	// context so that we can log the auth kind
 	// and the auth variable data to the request.
diff --git a/db/db_test.go b/db/db_test.go
index 6765bdca..61457904 100644
--- a/db/db_test.go
+++ b/db/db_test.go
@@ -40,7 +40,9 @@ func setupDB() {
 
 func setupKV() *fibre.Context {
 
-	auth := &cnf.Auth{}
+	keep := new(data.Doc)
+
+	auth := new(cnf.Auth)
 	auth.Kind = cnf.AuthKV
 	auth.Possible.NS = "*"
 	auth.Selected.NS = "*"
@@ -52,6 +54,7 @@ func setupKV() *fibre.Context {
 
 	ctx := fibre.NewContext(req, res, nil)
 	ctx.Set("auth", auth)
+	ctx.Set("keep", keep)
 
 	return ctx
 
@@ -59,7 +62,9 @@ func setupKV() *fibre.Context {
 
 func setupSC() *fibre.Context {
 
-	auth := &cnf.Auth{}
+	keep := new(data.Doc)
+
+	auth := new(cnf.Auth)
 	auth.Kind = cnf.AuthSC
 	auth.Possible.NS = "*"
 	auth.Selected.NS = "*"
@@ -71,6 +76,7 @@ func setupSC() *fibre.Context {
 
 	ctx := fibre.NewContext(req, res, nil)
 	ctx.Set("auth", auth)
+	ctx.Set("keep", keep)
 
 	return ctx
 
diff --git a/db/vars.go b/db/vars.go
index 9a3440c6..f52bd360 100644
--- a/db/vars.go
+++ b/db/vars.go
@@ -41,6 +41,7 @@ const (
 	ctxKeyVars  = "vars"
 	ctxKeySubs  = "subs"
 	ctxKeySpec  = "spec"
+	ctxKeyKeep  = "keep"
 	ctxKeyAuth  = "auth"
 	ctxKeyKind  = "kind"
 	ctxKeyScope = "scope"
@@ -50,6 +51,7 @@ const (
 	varKeyId     = "id"
 	varKeyIp     = "ip"
 	varKeyAuth   = "auth"
+	varKeyKeep   = "keep"
 	varKeyThis   = "this"
 	varKeyScope  = "scope"
 	varKeyValue  = "value"
@@ -65,5 +67,5 @@ var (
 	queryNotExecuted      = errors.New("Query not executed")
 	queryIdentFailed      = errors.New("Found ident but no doc available")
 	featureNotImplemented = errors.New("Feature is not yet implemented")
-	paramSearchKeys       = []string{ctxKeySpec, ctxKeySubs, ctxKeyVars}
+	paramSearchKeys       = []string{ctxKeySpec, ctxKeySubs, ctxKeyVars, ctxKeyKeep}
 )
diff --git a/web/auth.go b/web/auth.go
index 592bf3b2..7f5570d5 100644
--- a/web/auth.go
+++ b/web/auth.go
@@ -31,6 +31,8 @@ import (
 	"github.com/abcum/surreal/sql"
 	"github.com/dgrijalva/jwt-go"
 	"github.com/gorilla/websocket"
+
+	"github.com/abcum/surreal/util/data"
 )
 
 var ignore = func() error {
@@ -47,6 +49,7 @@ const (
 	varKeyTb     = "TB"
 	varKeyId     = "ID"
 	varKeyAuth   = "auth"
+	varKeyKeep   = "keep"
 	varKeyUser   = "user"
 	varKeyPass   = "pass"
 	varKeyOrigin = "origin"
@@ -65,7 +68,18 @@ func auth() fibre.MiddlewareFunc {
 	return func(h fibre.HandlerFunc) fibre.HandlerFunc {
 		return func(c *fibre.Context) (err error) {
 
-			auth := &cnf.Auth{}
+			// Initialise any session level variables
+			// which will be valid across all requests
+			// which are made over this connection.
+
+			vars := new(data.Doc)
+			c.Set(varKeyKeep, vars)
+
+			// Initialise the connection authentication
+			// information which will store whether the
+			// connection has authenticated or not.
+
+			auth := new(cnf.Auth)
 			c.Set(varKeyAuth, auth)
 
 			// Start off with an authentication level
diff --git a/web/rpc.go b/web/rpc.go
index 7fa0f0cf..0fad8890 100644
--- a/web/rpc.go
+++ b/web/rpc.go
@@ -19,6 +19,7 @@ import (
 	"github.com/abcum/surreal/cnf"
 	"github.com/abcum/surreal/db"
 	"github.com/abcum/surreal/sql"
+	"github.com/abcum/surreal/util/data"
 	"github.com/abcum/surreal/util/rand"
 )
 
@@ -36,6 +37,10 @@ func (r *rpc) Auth(c *fibre.Context, auth string) (interface{}, error) {
 	return c.Get("auth").(*cnf.Auth).Data, checkBearer(c, auth, ignore)
 }
 
+func (r *rpc) Let(c *fibre.Context, key string, val interface{}) (interface{}, error) {
+	return c.Get("keep").(*data.Doc).Set(val, key)
+}
+
 func (r *rpc) Query(c *fibre.Context, sql string, vars map[string]interface{}) (interface{}, error) {
 	return db.Execute(c, sql, vars)
 }