Store data on a per sessions basis, based on the Gin interface.

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 2016-12-09
+
+* Add metadata management for sessions.
+
 ## 2016-05-09
 
 * Add method `HandlePong` to melody instance.

README.md

@@ -14,6 +14,7 @@ your way so you can write real-time apps. Features include:
 * [x] A simple way to broadcast to all or selected connected sessions.
 * [x] Message buffers making concurrent writing safe.
 * [x] Automatic handling of ping/pong and session timeouts.
+* [x] Store data on sessions.
 
 ## Install
 
@@ -170,6 +171,8 @@ func main() {
 * Ola Holmström (@olahol)
 * Shogo Iwano (@shiwano)
 * Matt Caldwell (@mattcaldwell)
+* Heikki Uljas (@huljas)
+* Robbie Trencheny (@robbiet480)
 
 ## FAQ
 

melody.go

@@ -3,6 +3,7 @@ package melody
 import (
 	"github.com/gorilla/websocket"
 	"net/http"
+	"sync"
 )
 
 type handleMessageFunc func(*Session, []byte)
@@ -88,10 +89,11 @@ func (m *Melody) HandleRequest(w http.ResponseWriter, r *http.Request) {
 
 	session := &Session{
 		Request: r,
-		Params:  make(map[string]interface{}),
+		Keys:    nil,
 		conn:    conn,
 		output:  make(chan *envelope, m.Config.MessageBufferSize),
 		melody:  m,
+		lock:    &sync.Mutex{},
 	}
 
 	m.hub.register <- session

melody_test.go

@@ -5,6 +5,7 @@ import (
 	"github.com/gorilla/websocket"
 	"net/http"
 	"net/http/httptest"
+	"strconv"
 	"strings"
 	"testing"
 	"testing/quick"
@@ -141,6 +142,58 @@ func TestHandlers(t *testing.T) {
 	NewDialer(server.URL)
 }
 
+func TestMetadata(t *testing.T) {
+	echo := NewTestServer()
+	echo.m.HandleConnect(func(session *Session) {
+		session.Set("stamp", time.Now().UnixNano())
+	})
+	echo.m.HandleMessage(func(session *Session, msg []byte) {
+		stamp := session.MustGet("stamp").(int64)
+		session.Write([]byte(strconv.Itoa(int(stamp))))
+	})
+	server := httptest.NewServer(echo)
+	defer server.Close()
+
+	fn := func(msg string) bool {
+		conn, err := NewDialer(server.URL)
+		defer conn.Close()
+
+		if err != nil {
+			t.Error(err)
+			return false
+		}
+
+		conn.WriteMessage(websocket.TextMessage, []byte(msg))
+
+		_, ret, err := conn.ReadMessage()
+
+		if err != nil {
+			t.Error(err)
+			return false
+		}
+
+		stamp, err := strconv.Atoi(string(ret))
+
+		if err != nil {
+			t.Error(err)
+			return false
+		}
+
+		diff := int(time.Now().UnixNano()) - stamp
+
+		if diff <= 0 {
+			t.Errorf("diff should be above 0 %d", diff)
+			return false
+		}
+
+		return true
+	}
+
+	if err := quick.Check(fn, nil); err != nil {
+		t.Error(err)
+	}
+}
+
 func TestUpgrader(t *testing.T) {
 	broadcast := NewTestServer()
 	broadcast.m.HandleMessage(func(session *Session, msg []byte) {

session.go

@@ -4,16 +4,18 @@ import (
 	"errors"
 	"github.com/gorilla/websocket"
 	"net/http"
+	"sync"
 	"time"
 )
 
 // Session is wrapper around websocket connections.
 type Session struct {
 	Request *http.Request
-	Params  map[string]interface{}
+	Keys    map[string]interface{}
 	conn    *websocket.Conn
 	output  chan *envelope
 	melody  *Melody
+	lock    *sync.Mutex
 }
 
 func (s *Session) writeMessage(message *envelope) {
@@ -119,3 +121,38 @@ func (s *Session) WriteBinary(msg []byte) {
 func (s *Session) Close() {
 	s.writeMessage(&envelope{t: websocket.CloseMessage, msg: []byte{}})
 }
+
+// Set is used to store a new key/value pair exclusivelly for this session.
+// It also lazy initializes s.Keys if it was not used previously.
+func (s *Session) Set(key string, value interface{}) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	if s.Keys == nil {
+		s.Keys = make(map[string]interface{})
+	}
+
+	s.Keys[key] = value
+}
+
+// Get returns the value for the given key, ie: (value, true).
+// If the value does not exists it returns (nil, false)
+func (s *Session) Get(key string) (value interface{}, exists bool) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	if s.Keys != nil {
+		value, exists = s.Keys[key]
+	}
+
+	return
+}
+
+// MustGet returns the value for the given key if it exists, otherwise it panics.
+func (s *Session) MustGet(key string) interface{} {
+	if value, exists := s.Get(key); exists {
+		return value
+	}
+
+	panic("Key \"" + key + "\" does not exist")
+}