From 89e2dbac1504d6df205938df750fa9eb15996afd Mon Sep 17 00:00:00 2001 From: Wim Date: Sat, 11 Mar 2023 18:14:49 +0100 Subject: Check client disconnect to exit for loop (api). Fixes #1983 (#2012) Also update to latest melody upstream --- vendor/github.com/olahol/melody/.gitignore | 5 + vendor/github.com/olahol/melody/CHANGELOG.md | 52 +++++ vendor/github.com/olahol/melody/LICENSE | 22 ++ vendor/github.com/olahol/melody/README.md | 157 +++++++++++++ vendor/github.com/olahol/melody/config.go | 22 ++ vendor/github.com/olahol/melody/doc.go | 22 ++ vendor/github.com/olahol/melody/envelope.go | 7 + vendor/github.com/olahol/melody/errors.go | 10 + vendor/github.com/olahol/melody/hub.go | 91 ++++++++ vendor/github.com/olahol/melody/melody.go | 321 +++++++++++++++++++++++++++ vendor/github.com/olahol/melody/session.go | 239 ++++++++++++++++++++ 11 files changed, 948 insertions(+) create mode 100644 vendor/github.com/olahol/melody/.gitignore create mode 100644 vendor/github.com/olahol/melody/CHANGELOG.md create mode 100644 vendor/github.com/olahol/melody/LICENSE create mode 100644 vendor/github.com/olahol/melody/README.md create mode 100644 vendor/github.com/olahol/melody/config.go create mode 100644 vendor/github.com/olahol/melody/doc.go create mode 100644 vendor/github.com/olahol/melody/envelope.go create mode 100644 vendor/github.com/olahol/melody/errors.go create mode 100644 vendor/github.com/olahol/melody/hub.go create mode 100644 vendor/github.com/olahol/melody/melody.go create mode 100644 vendor/github.com/olahol/melody/session.go (limited to 'vendor/github.com') diff --git a/vendor/github.com/olahol/melody/.gitignore b/vendor/github.com/olahol/melody/.gitignore new file mode 100644 index 00000000..31ea56b1 --- /dev/null +++ b/vendor/github.com/olahol/melody/.gitignore @@ -0,0 +1,5 @@ +.DS_Store +benchmark +*.swp +coverage.out +Makefile diff --git a/vendor/github.com/olahol/melody/CHANGELOG.md b/vendor/github.com/olahol/melody/CHANGELOG.md new file mode 100644 index 00000000..e0b19412 --- /dev/null +++ b/vendor/github.com/olahol/melody/CHANGELOG.md @@ -0,0 +1,52 @@ +## 2022-09-12 (v1.1.0) + +* Create Go module. +* Update examples. +* Fix concurrent panic (PR-65). +* Add `Sessions` to get all sessions (PR-53). +* Add `LocalAddr` and `RemoteAddr` (PR-55). + +## 2017-05-18 + +* Fix `HandleSentMessageBinary`. + +## 2017-04-11 + +* Allow any origin by default. +* Add `BroadcastMultiple`. + +## 2017-04-09 + +* Add control message support. +* Add `IsClosed` to Session. + +## 2017-02-10 + +* Return errors for some exposed methods. +* Add `HandleRequestWithKeys`. +* Add `HandleSentMessage` and `HandleSentMessageBinary`. + +## 2017-01-20 + +* Add `Len()` to fetch number of connected sessions. + +## 2016-12-09 + +* Add metadata management for sessions. + +## 2016-05-09 + +* Add method `HandlePong` to melody instance. + +## 2015-10-07 + +* Add broadcast methods for binary messages. + +## 2015-09-03 + +* Add `Close` method to melody instance. + +### 2015-06-10 + +* Support for binary messages. +* BroadcastOthers method. diff --git a/vendor/github.com/olahol/melody/LICENSE b/vendor/github.com/olahol/melody/LICENSE new file mode 100644 index 00000000..4ae857e8 --- /dev/null +++ b/vendor/github.com/olahol/melody/LICENSE @@ -0,0 +1,22 @@ +Copyright (c) 2015 Ola Holmström. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + + Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/vendor/github.com/olahol/melody/README.md b/vendor/github.com/olahol/melody/README.md new file mode 100644 index 00000000..53196b44 --- /dev/null +++ b/vendor/github.com/olahol/melody/README.md @@ -0,0 +1,157 @@ +# melody + +![Build Status](https://github.com/olahol/melody/actions/workflows/test.yml/badge.svg) +[![Codecov](https://img.shields.io/codecov/c/github/olahol/melody)](https://app.codecov.io/github/olahol/melody) +[![Go Report Card](https://goreportcard.com/badge/github.com/olahol/melody)](https://goreportcard.com/report/github.com/olahol/melody) +[![GoDoc](https://godoc.org/github.com/olahol/melody?status.svg)](https://godoc.org/github.com/olahol/melody) + +> :notes: Minimalist websocket framework for Go. + +Melody is websocket framework based on [github.com/gorilla/websocket](https://github.com/gorilla/websocket) +that abstracts away the tedious parts of handling websockets. It gets out of +your way so you can write real-time apps. Features include: + +* [x] Clear and easy interface similar to `net/http` or Gin. +* [x] A simple way to broadcast to all or selected connected sessions. +* [x] Message buffers making concurrent writing safe. +* [x] Automatic handling of sending ping/pong heartbeats that timeout broken sessions. +* [x] Store data on sessions. + +## Install + +```bash +go get github.com/olahol/melody +``` + +## [Example: chat](https://github.com/olahol/melody/tree/master/examples/chat) + +[![Chat](https://cdn.rawgit.com/olahol/melody/master/examples/chat/demo.gif "Demo")](https://github.com/olahol/melody/tree/master/examples/chat) + +```go +package main + +import ( + "net/http" + + "github.com/olahol/melody" +) + +func main() { + m := melody.New() + + http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + http.ServeFile(w, r, "index.html") + }) + + http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) { + m.HandleRequest(w, r) + }) + + m.HandleMessage(func(s *melody.Session, msg []byte) { + m.Broadcast(msg) + }) + + http.ListenAndServe(":5000", nil) +} +``` + +## [Example: gophers](https://github.com/olahol/melody/tree/master/examples/gophers) + +[![Gophers](https://cdn.rawgit.com/olahol/melody/master/examples/gophers/demo.gif "Demo")](https://github.com/olahol/melody/tree/master/examples/gophers) + +```go +package main + +import ( + "net/http" + "strings" + + "github.com/google/uuid" + "github.com/olahol/melody" +) + +type GopherInfo struct { + ID, X, Y string +} + +func main() { + m := melody.New() + + http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + http.ServeFile(w, r, "index.html") + }) + + http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) { + m.HandleRequest(w, r) + }) + + m.HandleConnect(func(s *melody.Session) { + ss, _ := m.Sessions() + + for _, o := range ss { + value, exists := o.Get("info") + + if !exists { + continue + } + + info := value.(*GopherInfo) + + s.Write([]byte("set " + info.ID + " " + info.X + " " + info.Y)) + } + + id := uuid.NewString() + s.Set("info", &GopherInfo{id, "0", "0"}) + + s.Write([]byte("iam " + id)) + }) + + m.HandleDisconnect(func(s *melody.Session) { + value, exists := s.Get("info") + + if !exists { + return + } + + info := value.(*GopherInfo) + + m.BroadcastOthers([]byte("dis "+info.ID), s) + }) + + m.HandleMessage(func(s *melody.Session, msg []byte) { + p := strings.Split(string(msg), " ") + value, exists := s.Get("info") + + if len(p) != 2 || !exists { + return + } + + info := value.(*GopherInfo) + info.X = p[0] + info.Y = p[1] + + m.BroadcastOthers([]byte("set "+info.ID+" "+info.X+" "+info.Y), s) + }) + + http.ListenAndServe(":5000", nil) +} +``` + +### [More examples](https://github.com/olahol/melody/tree/master/examples) + +## [Documentation](https://godoc.org/github.com/olahol/melody) + +## Contributors + + + + + +## FAQ + +If you are getting a `403` when trying to connect to your websocket you can [change allow all origin hosts](http://godoc.org/github.com/gorilla/websocket#hdr-Origin_Considerations): + +```go +m := melody.New() +m.Upgrader.CheckOrigin = func(r *http.Request) bool { return true } +``` diff --git a/vendor/github.com/olahol/melody/config.go b/vendor/github.com/olahol/melody/config.go new file mode 100644 index 00000000..81ebd057 --- /dev/null +++ b/vendor/github.com/olahol/melody/config.go @@ -0,0 +1,22 @@ +package melody + +import "time" + +// Config melody configuration struct. +type Config struct { + WriteWait time.Duration // Milliseconds until write times out. + PongWait time.Duration // Timeout for waiting on pong. + PingPeriod time.Duration // Milliseconds between pings. + MaxMessageSize int64 // Maximum size in bytes of a message. + MessageBufferSize int // The max amount of messages that can be in a sessions buffer before it starts dropping them. +} + +func newConfig() *Config { + return &Config{ + WriteWait: 10 * time.Second, + PongWait: 60 * time.Second, + PingPeriod: (60 * time.Second * 9) / 10, + MaxMessageSize: 512, + MessageBufferSize: 256, + } +} diff --git a/vendor/github.com/olahol/melody/doc.go b/vendor/github.com/olahol/melody/doc.go new file mode 100644 index 00000000..b54b6d4a --- /dev/null +++ b/vendor/github.com/olahol/melody/doc.go @@ -0,0 +1,22 @@ +// Copyright 2015 Ola Holmström. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package melody implements a framework for dealing with WebSockets. +// +// Example +// +// A broadcasting echo server: +// +// func main() { +// m := melody.New() +// http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) { +// m.HandleRequest(w, r) +// }) +// m.HandleMessage(func(s *melody.Session, msg []byte) { +// m.Broadcast(msg) +// }) +// http.ListenAndServe(":5000", nil) +// } + +package melody diff --git a/vendor/github.com/olahol/melody/envelope.go b/vendor/github.com/olahol/melody/envelope.go new file mode 100644 index 00000000..baa55a3a --- /dev/null +++ b/vendor/github.com/olahol/melody/envelope.go @@ -0,0 +1,7 @@ +package melody + +type envelope struct { + t int + msg []byte + filter filterFunc +} diff --git a/vendor/github.com/olahol/melody/errors.go b/vendor/github.com/olahol/melody/errors.go new file mode 100644 index 00000000..d8be6157 --- /dev/null +++ b/vendor/github.com/olahol/melody/errors.go @@ -0,0 +1,10 @@ +package melody + +import "errors" + +var ( + ErrClosed = errors.New("melody instance is closed") + ErrSessionClosed = errors.New("session is closed") + ErrWriteClosed = errors.New("tried to write to closed a session") + ErrMessageBufferFull = errors.New("session message buffer is full") +) diff --git a/vendor/github.com/olahol/melody/hub.go b/vendor/github.com/olahol/melody/hub.go new file mode 100644 index 00000000..4a8cffac --- /dev/null +++ b/vendor/github.com/olahol/melody/hub.go @@ -0,0 +1,91 @@ +package melody + +import ( + "sync" +) + +type hub struct { + sessions map[*Session]bool + broadcast chan *envelope + register chan *Session + unregister chan *Session + exit chan *envelope + open bool + rwmutex *sync.RWMutex +} + +func newHub() *hub { + return &hub{ + sessions: make(map[*Session]bool), + broadcast: make(chan *envelope), + register: make(chan *Session), + unregister: make(chan *Session), + exit: make(chan *envelope), + open: true, + rwmutex: &sync.RWMutex{}, + } +} + +func (h *hub) run() { +loop: + for { + select { + case s := <-h.register: + h.rwmutex.Lock() + h.sessions[s] = true + h.rwmutex.Unlock() + case s := <-h.unregister: + if _, ok := h.sessions[s]; ok { + h.rwmutex.Lock() + delete(h.sessions, s) + h.rwmutex.Unlock() + } + case m := <-h.broadcast: + h.rwmutex.RLock() + for s := range h.sessions { + if m.filter != nil { + if m.filter(s) { + s.writeMessage(m) + } + } else { + s.writeMessage(m) + } + } + h.rwmutex.RUnlock() + case m := <-h.exit: + h.rwmutex.Lock() + for s := range h.sessions { + s.writeMessage(m) + delete(h.sessions, s) + s.Close() + } + h.open = false + h.rwmutex.Unlock() + break loop + } + } +} + +func (h *hub) closed() bool { + h.rwmutex.RLock() + defer h.rwmutex.RUnlock() + return !h.open +} + +func (h *hub) len() int { + h.rwmutex.RLock() + defer h.rwmutex.RUnlock() + + return len(h.sessions) +} + +func (h *hub) all() []*Session { + h.rwmutex.RLock() + defer h.rwmutex.RUnlock() + + s := make([]*Session, 0, len(h.sessions)) + for k := range h.sessions { + s = append(s, k) + } + return s +} diff --git a/vendor/github.com/olahol/melody/melody.go b/vendor/github.com/olahol/melody/melody.go new file mode 100644 index 00000000..ca2c5dd1 --- /dev/null +++ b/vendor/github.com/olahol/melody/melody.go @@ -0,0 +1,321 @@ +package melody + +import ( + "net/http" + "sync" + + "github.com/gorilla/websocket" +) + +// Close codes defined in RFC 6455, section 11.7. +// Duplicate of codes from gorilla/websocket for convenience. +const ( + CloseNormalClosure = 1000 + CloseGoingAway = 1001 + CloseProtocolError = 1002 + CloseUnsupportedData = 1003 + CloseNoStatusReceived = 1005 + CloseAbnormalClosure = 1006 + CloseInvalidFramePayloadData = 1007 + ClosePolicyViolation = 1008 + CloseMessageTooBig = 1009 + CloseMandatoryExtension = 1010 + CloseInternalServerErr = 1011 + CloseServiceRestart = 1012 + CloseTryAgainLater = 1013 + CloseTLSHandshake = 1015 +) + +// Duplicate of codes from gorilla/websocket for convenience. +var validReceivedCloseCodes = map[int]bool{ + // see http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number + + CloseNormalClosure: true, + CloseGoingAway: true, + CloseProtocolError: true, + CloseUnsupportedData: true, + CloseNoStatusReceived: false, + CloseAbnormalClosure: false, + CloseInvalidFramePayloadData: true, + ClosePolicyViolation: true, + CloseMessageTooBig: true, + CloseMandatoryExtension: true, + CloseInternalServerErr: true, + CloseServiceRestart: true, + CloseTryAgainLater: true, + CloseTLSHandshake: false, +} + +type handleMessageFunc func(*Session, []byte) +type handleErrorFunc func(*Session, error) +type handleCloseFunc func(*Session, int, string) error +type handleSessionFunc func(*Session) +type filterFunc func(*Session) bool + +// Melody implements a websocket manager. +type Melody struct { + Config *Config + Upgrader *websocket.Upgrader + messageHandler handleMessageFunc + messageHandlerBinary handleMessageFunc + messageSentHandler handleMessageFunc + messageSentHandlerBinary handleMessageFunc + errorHandler handleErrorFunc + closeHandler handleCloseFunc + connectHandler handleSessionFunc + disconnectHandler handleSessionFunc + pongHandler handleSessionFunc + hub *hub +} + +// New creates a new melody instance with default Upgrader and Config. +func New() *Melody { + upgrader := &websocket.Upgrader{ + ReadBufferSize: 1024, + WriteBufferSize: 1024, + CheckOrigin: func(r *http.Request) bool { return true }, + } + + hub := newHub() + + go hub.run() + + return &Melody{ + Config: newConfig(), + Upgrader: upgrader, + messageHandler: func(*Session, []byte) {}, + messageHandlerBinary: func(*Session, []byte) {}, + messageSentHandler: func(*Session, []byte) {}, + messageSentHandlerBinary: func(*Session, []byte) {}, + errorHandler: func(*Session, error) {}, + closeHandler: nil, + connectHandler: func(*Session) {}, + disconnectHandler: func(*Session) {}, + pongHandler: func(*Session) {}, + hub: hub, + } +} + +// HandleConnect fires fn when a session connects. +func (m *Melody) HandleConnect(fn func(*Session)) { + m.connectHandler = fn +} + +// HandleDisconnect fires fn when a session disconnects. +func (m *Melody) HandleDisconnect(fn func(*Session)) { + m.disconnectHandler = fn +} + +// HandlePong fires fn when a pong is received from a session. +func (m *Melody) HandlePong(fn func(*Session)) { + m.pongHandler = fn +} + +// HandleMessage fires fn when a text message comes in. +func (m *Melody) HandleMessage(fn func(*Session, []byte)) { + m.messageHandler = fn +} + +// HandleMessageBinary fires fn when a binary message comes in. +func (m *Melody) HandleMessageBinary(fn func(*Session, []byte)) { + m.messageHandlerBinary = fn +} + +// HandleSentMessage fires fn when a text message is successfully sent. +func (m *Melody) HandleSentMessage(fn func(*Session, []byte)) { + m.messageSentHandler = fn +} + +// HandleSentMessageBinary fires fn when a binary message is successfully sent. +func (m *Melody) HandleSentMessageBinary(fn func(*Session, []byte)) { + m.messageSentHandlerBinary = fn +} + +// HandleError fires fn when a session has an error. +func (m *Melody) HandleError(fn func(*Session, error)) { + m.errorHandler = fn +} + +// HandleClose sets the handler for close messages received from the session. +// The code argument to h is the received close code or CloseNoStatusReceived +// if the close message is empty. The default close handler sends a close frame +// back to the session. +// +// The application must read the connection to process close messages as +// described in the section on Control Frames above. +// +// The connection read methods return a CloseError when a close frame is +// received. Most applications should handle close messages as part of their +// normal error handling. Applications should only set a close handler when the +// application must perform some action before sending a close frame back to +// the session. +func (m *Melody) HandleClose(fn func(*Session, int, string) error) { + if fn != nil { + m.closeHandler = fn + } +} + +// HandleRequest upgrades http requests to websocket connections and dispatches them to be handled by the melody instance. +func (m *Melody) HandleRequest(w http.ResponseWriter, r *http.Request) error { + return m.HandleRequestWithKeys(w, r, nil) +} + +// HandleRequestWithKeys does the same as HandleRequest but populates session.Keys with keys. +func (m *Melody) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, keys map[string]interface{}) error { + if m.hub.closed() { + return ErrClosed + } + + conn, err := m.Upgrader.Upgrade(w, r, w.Header()) + + if err != nil { + return err + } + + session := &Session{ + Request: r, + Keys: keys, + conn: conn, + output: make(chan *envelope, m.Config.MessageBufferSize), + outputDone: make(chan struct{}), + melody: m, + open: true, + rwmutex: &sync.RWMutex{}, + } + + m.hub.register <- session + + m.connectHandler(session) + + go session.writePump() + + session.readPump() + + if !m.hub.closed() { + m.hub.unregister <- session + } + + session.close() + + m.disconnectHandler(session) + + return nil +} + +// Broadcast broadcasts a text message to all sessions. +func (m *Melody) Broadcast(msg []byte) error { + if m.hub.closed() { + return ErrClosed + } + + message := &envelope{t: websocket.TextMessage, msg: msg} + m.hub.broadcast <- message + + return nil +} + +// BroadcastFilter broadcasts a text message to all sessions that fn returns true for. +func (m *Melody) BroadcastFilter(msg []byte, fn func(*Session) bool) error { + if m.hub.closed() { + return ErrClosed + } + + message := &envelope{t: websocket.TextMessage, msg: msg, filter: fn} + m.hub.broadcast <- message + + return nil +} + +// BroadcastOthers broadcasts a text message to all sessions except session s. +func (m *Melody) BroadcastOthers(msg []byte, s *Session) error { + return m.BroadcastFilter(msg, func(q *Session) bool { + return s != q + }) +} + +// BroadcastMultiple broadcasts a text message to multiple sessions given in the sessions slice. +func (m *Melody) BroadcastMultiple(msg []byte, sessions []*Session) error { + for _, sess := range sessions { + if writeErr := sess.Write(msg); writeErr != nil { + return writeErr + } + } + return nil +} + +// BroadcastBinary broadcasts a binary message to all sessions. +func (m *Melody) BroadcastBinary(msg []byte) error { + if m.hub.closed() { + return ErrClosed + } + + message := &envelope{t: websocket.BinaryMessage, msg: msg} + m.hub.broadcast <- message + + return nil +} + +// BroadcastBinaryFilter broadcasts a binary message to all sessions that fn returns true for. +func (m *Melody) BroadcastBinaryFilter(msg []byte, fn func(*Session) bool) error { + if m.hub.closed() { + return ErrClosed + } + + message := &envelope{t: websocket.BinaryMessage, msg: msg, filter: fn} + m.hub.broadcast <- message + + return nil +} + +// BroadcastBinaryOthers broadcasts a binary message to all sessions except session s. +func (m *Melody) BroadcastBinaryOthers(msg []byte, s *Session) error { + return m.BroadcastBinaryFilter(msg, func(q *Session) bool { + return s != q + }) +} + +// Sessions returns all sessions. An error is returned if the melody session is closed. +func (m *Melody) Sessions() ([]*Session, error) { + if m.hub.closed() { + return nil, ErrClosed + } + return m.hub.all(), nil +} + +// Close closes the melody instance and all connected sessions. +func (m *Melody) Close() error { + if m.hub.closed() { + return ErrClosed + } + + m.hub.exit <- &envelope{t: websocket.CloseMessage, msg: []byte{}} + + return nil +} + +// CloseWithMsg closes the melody instance with the given close payload and all connected sessions. +// Use the FormatCloseMessage function to format a proper close message payload. +func (m *Melody) CloseWithMsg(msg []byte) error { + if m.hub.closed() { + return ErrClosed + } + + m.hub.exit <- &envelope{t: websocket.CloseMessage, msg: msg} + + return nil +} + +// Len return the number of connected sessions. +func (m *Melody) Len() int { + return m.hub.len() +} + +// IsClosed returns the status of the melody instance. +func (m *Melody) IsClosed() bool { + return m.hub.closed() +} + +// FormatCloseMessage formats closeCode and text as a WebSocket close message. +func FormatCloseMessage(closeCode int, text string) []byte { + return websocket.FormatCloseMessage(closeCode, text) +} diff --git a/vendor/github.com/olahol/melody/session.go b/vendor/github.com/olahol/melody/session.go new file mode 100644 index 00000000..879a724d --- /dev/null +++ b/vendor/github.com/olahol/melody/session.go @@ -0,0 +1,239 @@ +package melody + +import ( + "net" + "net/http" + "sync" + "time" + + "github.com/gorilla/websocket" +) + +// Session wrapper around websocket connections. +type Session struct { + Request *http.Request + Keys map[string]interface{} + conn *websocket.Conn + output chan *envelope + outputDone chan struct{} + melody *Melody + open bool + rwmutex *sync.RWMutex +} + +func (s *Session) writeMessage(message *envelope) { + if s.closed() { + s.melody.errorHandler(s, ErrWriteClosed) + return + } + + select { + case s.output <- message: + default: + s.melody.errorHandler(s, ErrMessageBufferFull) + } +} + +func (s *Session) writeRaw(message *envelope) error { + if s.closed() { + return ErrWriteClosed + } + + s.conn.SetWriteDeadline(time.Now().Add(s.melody.Config.WriteWait)) + err := s.conn.WriteMessage(message.t, message.msg) + + if err != nil { + return err + } + + return nil +} + +func (s *Session) closed() bool { + s.rwmutex.RLock() + defer s.rwmutex.RUnlock() + + return !s.open +} + +func (s *Session) close() { + s.rwmutex.Lock() + open := s.open + s.open = false + s.rwmutex.Unlock() + if open { + s.conn.Close() + close(s.outputDone) + } +} + +func (s *Session) ping() { + s.writeRaw(&envelope{t: websocket.PingMessage, msg: []byte{}}) +} + +func (s *Session) writePump() { + ticker := time.NewTicker(s.melody.Config.PingPeriod) + defer ticker.Stop() + +loop: + for { + select { + case msg := <-s.output: + err := s.writeRaw(msg) + + if err != nil { + s.melody.errorHandler(s, err) + break loop + } + + if msg.t == websocket.CloseMessage { + break loop + } + + if msg.t == websocket.TextMessage { + s.melody.messageSentHandler(s, msg.msg) + } + + if msg.t == websocket.BinaryMessage { + s.melody.messageSentHandlerBinary(s, msg.msg) + } + case <-ticker.C: + s.ping() + case _, ok := <-s.outputDone: + if !ok { + break loop + } + } + } + + s.close() +} + +func (s *Session) readPump() { + s.conn.SetReadLimit(s.melody.Config.MaxMessageSize) + s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait)) + + s.conn.SetPongHandler(func(string) error { + s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait)) + s.melody.pongHandler(s) + return nil + }) + + if s.melody.closeHandler != nil { + s.conn.SetCloseHandler(func(code int, text string) error { + return s.melody.closeHandler(s, code, text) + }) + } + + for { + t, message, err := s.conn.ReadMessage() + + if err != nil { + s.melody.errorHandler(s, err) + break + } + + if t == websocket.TextMessage { + s.melody.messageHandler(s, message) + } + + if t == websocket.BinaryMessage { + s.melody.messageHandlerBinary(s, message) + } + } +} + +// Write writes message to session. +func (s *Session) Write(msg []byte) error { + if s.closed() { + return ErrSessionClosed + } + + s.writeMessage(&envelope{t: websocket.TextMessage, msg: msg}) + + return nil +} + +// WriteBinary writes a binary message to session. +func (s *Session) WriteBinary(msg []byte) error { + if s.closed() { + return ErrSessionClosed + } + + s.writeMessage(&envelope{t: websocket.BinaryMessage, msg: msg}) + + return nil +} + +// Close closes session. +func (s *Session) Close() error { + if s.closed() { + return ErrSessionClosed + } + + s.writeMessage(&envelope{t: websocket.CloseMessage, msg: []byte{}}) + + return nil +} + +// CloseWithMsg closes the session with the provided payload. +// Use the FormatCloseMessage function to format a proper close message payload. +func (s *Session) CloseWithMsg(msg []byte) error { + if s.closed() { + return ErrSessionClosed + } + + s.writeMessage(&envelope{t: websocket.CloseMessage, msg: msg}) + + return nil +} + +// Set is used to store a new key/value pair exclusively for this session. +// It also lazy initializes s.Keys if it was not used previously. +func (s *Session) Set(key string, value interface{}) { + s.rwmutex.Lock() + defer s.rwmutex.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.rwmutex.RLock() + defer s.rwmutex.RUnlock() + + 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") +} + +// IsClosed returns the status of the connection. +func (s *Session) IsClosed() bool { + return s.closed() +} + +// LocalAddr returns the local addr of the connection. +func (s *Session) LocalAddr() net.Addr { + return s.conn.LocalAddr() +} + +// RemoteAddr returns the remote addr of the connection. +func (s *Session) RemoteAddr() net.Addr { + return s.conn.RemoteAddr() +} -- cgit v1.2.3