├── .github
└── workflows
│ └── test.yml
├── .gitignore
├── CHANGELOG.md
├── LICENSE
├── README.md
├── config.go
├── doc.go
├── envelope.go
├── errors.go
├── examples
├── chat-echo
│ ├── index.html
│ └── main.go
├── chat-gin
│ ├── index.html
│ └── main.go
├── chat
│ ├── demo.gif
│ ├── index.html
│ └── main.go
├── filewatch
│ ├── file.txt
│ ├── index.html
│ └── main.go
├── go.mod
├── go.sum
├── gophers
│ ├── demo.gif
│ ├── index.html
│ └── main.go
└── multichat
│ ├── chan.html
│ ├── index.html
│ └── main.go
├── go.mod
├── go.sum
├── hub.go
├── melody.go
├── melody_test.go
└── session.go
/.github/workflows/test.yml:
--------------------------------------------------------------------------------
1 | name: test
2 | on:
3 | push:
4 | branches:
5 | - master
6 | jobs:
7 | test:
8 | strategy:
9 | matrix:
10 | go-version: ["1.20", "1.21", "1.22"]
11 | runs-on: "ubuntu-latest"
12 | steps:
13 | - uses: actions/checkout@v3
14 | with:
15 | fetch-depth: 2
16 | - uses: actions/setup-go@v3
17 | with:
18 | go-version: ${{ matrix.go-version }}
19 | - run: go get -t -v ./...
20 | - run: go test -race -coverprofile=coverage.out -covermode=atomic -timeout 60s
21 | - uses: codecov/codecov-action@v3
22 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | .DS_Store
2 | benchmark
3 | *.swp
4 | coverage.out
5 | Makefile
6 |
--------------------------------------------------------------------------------
/CHANGELOG.md:
--------------------------------------------------------------------------------
1 | ## 2022-09-12 (v1.1.0)
2 |
3 | * Create Go module.
4 | * Update examples.
5 | * Fix concurrent panic (PR-65).
6 | * Add `Sessions` to get all sessions (PR-53).
7 | * Add `LocalAddr` and `RemoteAddr` (PR-55).
8 |
9 | ## 2017-05-18
10 |
11 | * Fix `HandleSentMessageBinary`.
12 |
13 | ## 2017-04-11
14 |
15 | * Allow any origin by default.
16 | * Add `BroadcastMultiple`.
17 |
18 | ## 2017-04-09
19 |
20 | * Add control message support.
21 | * Add `IsClosed` to Session.
22 |
23 | ## 2017-02-10
24 |
25 | * Return errors for some exposed methods.
26 | * Add `HandleRequestWithKeys`.
27 | * Add `HandleSentMessage` and `HandleSentMessageBinary`.
28 |
29 | ## 2017-01-20
30 |
31 | * Add `Len()` to fetch number of connected sessions.
32 |
33 | ## 2016-12-09
34 |
35 | * Add metadata management for sessions.
36 |
37 | ## 2016-05-09
38 |
39 | * Add method `HandlePong` to melody instance.
40 |
41 | ## 2015-10-07
42 |
43 | * Add broadcast methods for binary messages.
44 |
45 | ## 2015-09-03
46 |
47 | * Add `Close` method to melody instance.
48 |
49 | ### 2015-06-10
50 |
51 | * Support for binary messages.
52 | * BroadcastOthers method.
53 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | Copyright (c) 2015 Ola Holmström. All rights reserved.
2 |
3 | Redistribution and use in source and binary forms, with or without
4 | modification, are permitted provided that the following conditions are met:
5 |
6 | Redistributions of source code must retain the above copyright notice, this
7 | list of conditions and the following disclaimer.
8 |
9 | Redistributions in binary form must reproduce the above copyright notice,
10 | this list of conditions and the following disclaimer in the documentation
11 | and/or other materials provided with the distribution.
12 |
13 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
14 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
15 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
16 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
17 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
18 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
19 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
20 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
21 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
22 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
23 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # melody
2 |
3 | 
4 | [](https://app.codecov.io/github/olahol/melody)
5 | [](https://goreportcard.com/report/github.com/olahol/melody)
6 | [](https://godoc.org/github.com/olahol/melody)
7 |
8 | > :notes: Minimalist websocket framework for Go.
9 |
10 | Melody is websocket framework based on [github.com/gorilla/websocket](https://github.com/gorilla/websocket)
11 | that abstracts away the tedious parts of handling websockets. It gets out of
12 | your way so you can write real-time apps. Features include:
13 |
14 | * [x] Clear and easy interface similar to `net/http` or Gin.
15 | * [x] A simple way to broadcast to all or selected connected sessions.
16 | * [x] Message buffers making concurrent writing safe.
17 | * [x] Automatic handling of sending ping/pong heartbeats that timeout broken sessions.
18 | * [x] Store data on sessions.
19 |
20 | ## Install
21 |
22 | ```bash
23 | go get github.com/olahol/melody
24 | ```
25 |
26 | ## [Example: chat](https://github.com/olahol/melody/tree/master/examples/chat)
27 |
28 | [](https://github.com/olahol/melody/tree/master/examples/chat)
29 |
30 | ```go
31 | package main
32 |
33 | import (
34 | "net/http"
35 |
36 | "github.com/olahol/melody"
37 | )
38 |
39 | func main() {
40 | m := melody.New()
41 |
42 | http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
43 | http.ServeFile(w, r, "index.html")
44 | })
45 |
46 | http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
47 | m.HandleRequest(w, r)
48 | })
49 |
50 | m.HandleMessage(func(s *melody.Session, msg []byte) {
51 | m.Broadcast(msg)
52 | })
53 |
54 | http.ListenAndServe(":5000", nil)
55 | }
56 | ```
57 |
58 | ## [Example: gophers](https://github.com/olahol/melody/tree/master/examples/gophers)
59 |
60 | [](https://github.com/olahol/melody/tree/master/examples/gophers)
61 |
62 | ```go
63 | package main
64 |
65 | import (
66 | "fmt"
67 | "net/http"
68 | "sync/atomic"
69 |
70 | "github.com/olahol/melody"
71 | )
72 |
73 | var idCounter atomic.Int64
74 |
75 | func main() {
76 | m := melody.New()
77 |
78 | http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
79 | http.ServeFile(w, r, "index.html")
80 | })
81 |
82 | http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
83 | m.HandleRequest(w, r)
84 | })
85 |
86 | m.HandleConnect(func(s *melody.Session) {
87 | id := idCounter.Add(1)
88 |
89 | s.Set("id", id)
90 |
91 | s.Write([]byte(fmt.Sprintf("iam %d", id)))
92 | })
93 |
94 | m.HandleDisconnect(func(s *melody.Session) {
95 | if id, ok := s.Get("id"); ok {
96 | m.BroadcastOthers([]byte(fmt.Sprintf("dis %d", id)), s)
97 | }
98 | })
99 |
100 | m.HandleMessage(func(s *melody.Session, msg []byte) {
101 | if id, ok := s.Get("id"); ok {
102 | m.BroadcastOthers([]byte(fmt.Sprintf("set %d %s", id, msg)), s)
103 | }
104 | })
105 |
106 | http.ListenAndServe(":5000", nil)
107 | }
108 | ```
109 |
110 | ### [More examples](https://github.com/olahol/melody/tree/master/examples)
111 |
112 | ## [Documentation](https://godoc.org/github.com/olahol/melody)
113 |
114 | ## Contributors
115 |
116 |
117 | 
118 |
119 |
120 | ## FAQ
121 |
122 | 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):
123 |
124 | ```go
125 | m := melody.New()
126 | m.Upgrader.CheckOrigin = func(r *http.Request) bool { return true }
127 | ```
128 |
--------------------------------------------------------------------------------
/config.go:
--------------------------------------------------------------------------------
1 | package melody
2 |
3 | import "time"
4 |
5 | // Config melody configuration struct.
6 | type Config struct {
7 | WriteWait time.Duration // Duration until write times out.
8 | PongWait time.Duration // Timeout for waiting on pong.
9 | PingPeriod time.Duration // Duration between pings.
10 | MaxMessageSize int64 // Maximum size in bytes of a message.
11 | MessageBufferSize int // The max amount of messages that can be in a sessions buffer before it starts dropping them.
12 | ConcurrentMessageHandling bool // Handle messages from sessions concurrently.
13 | }
14 |
15 | func newConfig() *Config {
16 | return &Config{
17 | WriteWait: 10 * time.Second,
18 | PongWait: 60 * time.Second,
19 | PingPeriod: 54 * time.Second,
20 | MaxMessageSize: 512,
21 | MessageBufferSize: 256,
22 | }
23 | }
24 |
--------------------------------------------------------------------------------
/doc.go:
--------------------------------------------------------------------------------
1 | // Copyright 2015 Ola Holmström. All rights reserved.
2 | // Use of this source code is governed by a BSD-style
3 | // license that can be found in the LICENSE file.
4 |
5 | // Package melody implements a framework for dealing with WebSockets.
6 | //
7 | // Example
8 | //
9 | // A broadcasting echo server:
10 | //
11 | // func main() {
12 | // m := melody.New()
13 | // http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
14 | // m.HandleRequest(w, r)
15 | // })
16 | // m.HandleMessage(func(s *melody.Session, msg []byte) {
17 | // m.Broadcast(msg)
18 | // })
19 | // http.ListenAndServe(":5000", nil)
20 | // }
21 |
22 | package melody
23 |
--------------------------------------------------------------------------------
/envelope.go:
--------------------------------------------------------------------------------
1 | package melody
2 |
3 | type envelope struct {
4 | t int
5 | msg []byte
6 | filter filterFunc
7 | }
8 |
--------------------------------------------------------------------------------
/errors.go:
--------------------------------------------------------------------------------
1 | package melody
2 |
3 | import "errors"
4 |
5 | var (
6 | ErrClosed = errors.New("melody instance is closed")
7 | ErrSessionClosed = errors.New("session is closed")
8 | ErrWriteClosed = errors.New("tried to write to closed a session")
9 | ErrMessageBufferFull = errors.New("session message buffer is full")
10 | )
11 |
--------------------------------------------------------------------------------
/examples/chat-echo/index.html:
--------------------------------------------------------------------------------
1 |
2 |
3 | Melody example: chatting
4 |
5 |
6 |
15 |
16 |
17 |
18 | Chat
19 |
20 |
21 |
22 |
23 |
49 |
50 |
51 |
--------------------------------------------------------------------------------
/examples/chat-echo/main.go:
--------------------------------------------------------------------------------
1 | package main
2 |
3 | import (
4 | "github.com/labstack/echo/v4"
5 | "github.com/labstack/echo/v4/middleware"
6 | "github.com/olahol/melody"
7 | "net/http"
8 | )
9 |
10 | func main() {
11 | e := echo.New()
12 | m := melody.New()
13 |
14 | e.Use(middleware.Logger())
15 | e.Use(middleware.Recover())
16 |
17 | e.GET("/", func(c echo.Context) error {
18 | http.ServeFile(c.Response().Writer, c.Request(), "index.html")
19 | return nil
20 | })
21 |
22 | e.GET("/ws", func(c echo.Context) error {
23 | m.HandleRequest(c.Response().Writer, c.Request())
24 | return nil
25 | })
26 |
27 | m.HandleMessage(func(s *melody.Session, msg []byte) {
28 | m.Broadcast(msg)
29 | })
30 |
31 | e.Logger.Fatal(e.Start(":5000"))
32 | }
33 |
--------------------------------------------------------------------------------
/examples/chat-gin/index.html:
--------------------------------------------------------------------------------
1 |
2 |
3 | Melody example: chatting
4 |
5 |
6 |
15 |
16 |
17 |
18 | Chat
19 |
20 |
21 |
22 |
23 |
49 |
50 |
51 |
--------------------------------------------------------------------------------
/examples/chat-gin/main.go:
--------------------------------------------------------------------------------
1 | package main
2 |
3 | import (
4 | "github.com/gin-gonic/gin"
5 | "github.com/olahol/melody"
6 | "net/http"
7 | )
8 |
9 | func main() {
10 | r := gin.Default()
11 | m := melody.New()
12 |
13 | r.GET("/", func(c *gin.Context) {
14 | http.ServeFile(c.Writer, c.Request, "index.html")
15 | })
16 |
17 | r.GET("/ws", func(c *gin.Context) {
18 | m.HandleRequest(c.Writer, c.Request)
19 | })
20 |
21 | m.HandleMessage(func(s *melody.Session, msg []byte) {
22 | m.Broadcast(msg)
23 | })
24 |
25 | r.Run(":5000")
26 | }
27 |
--------------------------------------------------------------------------------
/examples/chat/demo.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/olahol/melody/6bf86153260ddfd8ef1929323e0ee37cc3acd44a/examples/chat/demo.gif
--------------------------------------------------------------------------------
/examples/chat/index.html:
--------------------------------------------------------------------------------
1 |
2 |
3 | Melody example: chatting
4 |
5 |
6 |
15 |
16 |
17 |
18 | Chat
19 |
20 |
21 |
22 |
23 |
49 |
50 |
51 |
--------------------------------------------------------------------------------
/examples/chat/main.go:
--------------------------------------------------------------------------------
1 | package main
2 |
3 | import (
4 | "net/http"
5 |
6 | "github.com/olahol/melody"
7 | )
8 |
9 | func main() {
10 | m := melody.New()
11 |
12 | http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
13 | http.ServeFile(w, r, "index.html")
14 | })
15 |
16 | http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
17 | m.HandleRequest(w, r)
18 | })
19 |
20 | m.HandleMessage(func(s *melody.Session, msg []byte) {
21 | m.Broadcast(msg)
22 | })
23 |
24 | http.ListenAndServe(":5000", nil)
25 | }
26 |
--------------------------------------------------------------------------------
/examples/filewatch/file.txt:
--------------------------------------------------------------------------------
1 | Hello World!
2 |
--------------------------------------------------------------------------------
/examples/filewatch/index.html:
--------------------------------------------------------------------------------
1 |
2 |
3 | Melody example: file watching
4 |
5 |
6 |
15 |
16 |
17 |
18 | Watching a file
19 |
20 |
21 |
22 |
31 |
32 |
33 |
--------------------------------------------------------------------------------
/examples/filewatch/main.go:
--------------------------------------------------------------------------------
1 | package main
2 |
3 | import (
4 | "net/http"
5 | "os"
6 |
7 | "github.com/fsnotify/fsnotify"
8 | "github.com/olahol/melody"
9 | )
10 |
11 | func main() {
12 | file := "file.txt"
13 |
14 | m := melody.New()
15 | w, _ := fsnotify.NewWatcher()
16 |
17 | http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
18 | http.ServeFile(w, r, "index.html")
19 | })
20 |
21 | http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
22 | m.HandleRequest(w, r)
23 | })
24 |
25 | m.HandleConnect(func(s *melody.Session) {
26 | content, _ := os.ReadFile(file)
27 | s.Write(content)
28 | })
29 |
30 | go func() {
31 | for {
32 | ev := <-w.Events
33 | if ev.Op == fsnotify.Write {
34 | content, _ := os.ReadFile(ev.Name)
35 | m.Broadcast(content)
36 | }
37 | }
38 | }()
39 |
40 | w.Add(file)
41 |
42 | http.ListenAndServe(":5000", nil)
43 | }
44 |
--------------------------------------------------------------------------------
/examples/go.mod:
--------------------------------------------------------------------------------
1 | module github.com/olahol/melody/examples
2 |
3 | go 1.22
4 |
5 | replace github.com/olahol/melody => ../
6 |
7 | require (
8 | github.com/fsnotify/fsnotify v1.5.4
9 | github.com/gin-gonic/gin v1.9.1
10 | github.com/labstack/echo/v4 v4.9.0
11 | github.com/olahol/melody v0.0.0-00010101000000-000000000000
12 | )
13 |
14 | require (
15 | github.com/bytedance/sonic v1.9.1 // indirect
16 | github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect
17 | github.com/gabriel-vasile/mimetype v1.4.2 // indirect
18 | github.com/gin-contrib/sse v0.1.0 // indirect
19 | github.com/go-playground/locales v0.14.1 // indirect
20 | github.com/go-playground/universal-translator v0.18.1 // indirect
21 | github.com/go-playground/validator/v10 v10.14.0 // indirect
22 | github.com/goccy/go-json v0.10.2 // indirect
23 | github.com/golang-jwt/jwt v3.2.2+incompatible // indirect
24 | github.com/gorilla/websocket v1.5.0 // indirect
25 | github.com/json-iterator/go v1.1.12 // indirect
26 | github.com/klauspost/cpuid/v2 v2.2.4 // indirect
27 | github.com/labstack/gommon v0.3.1 // indirect
28 | github.com/leodido/go-urn v1.2.4 // indirect
29 | github.com/mattn/go-colorable v0.1.11 // indirect
30 | github.com/mattn/go-isatty v0.0.19 // indirect
31 | github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
32 | github.com/modern-go/reflect2 v1.0.2 // indirect
33 | github.com/pelletier/go-toml/v2 v2.0.8 // indirect
34 | github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
35 | github.com/ugorji/go/codec v1.2.11 // indirect
36 | github.com/valyala/bytebufferpool v1.0.0 // indirect
37 | github.com/valyala/fasttemplate v1.2.1 // indirect
38 | golang.org/x/arch v0.3.0 // indirect
39 | golang.org/x/crypto v0.9.0 // indirect
40 | golang.org/x/net v0.10.0 // indirect
41 | golang.org/x/sys v0.8.0 // indirect
42 | golang.org/x/text v0.9.0 // indirect
43 | golang.org/x/time v0.0.0-20201208040808-7e3f01d25324 // indirect
44 | google.golang.org/protobuf v1.30.0 // indirect
45 | gopkg.in/yaml.v3 v3.0.1 // indirect
46 | )
47 |
--------------------------------------------------------------------------------
/examples/go.sum:
--------------------------------------------------------------------------------
1 | github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM=
2 | github.com/bytedance/sonic v1.9.1 h1:6iJ6NqdoxCDr6mbY8h18oSO+cShGSMRGCEo7F2h0x8s=
3 | github.com/bytedance/sonic v1.9.1/go.mod h1:i736AoUSYt75HyZLoJW9ERYxcy6eaN6h4BZXU064P/U=
4 | github.com/chenzhuoyu/base64x v0.0.0-20211019084208-fb5309c8db06/go.mod h1:DH46F32mSOjUmXrMHnKwZdA8wcEefY7UVqBKYGjpdQY=
5 | github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 h1:qSGYFH7+jGhDF8vLC+iwCD4WpbV1EBDSzWkJODFLams=
6 | github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311/go.mod h1:b583jCggY9gE99b6G5LEC39OIiVsWj+R97kbl5odCEk=
7 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
8 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
9 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
10 | github.com/fsnotify/fsnotify v1.5.4 h1:jRbGcIw6P2Meqdwuo0H1p6JVLbL5DHKAKlYndzMwVZI=
11 | github.com/fsnotify/fsnotify v1.5.4/go.mod h1:OVB6XrOHzAwXMpEM7uPOzcehqUV2UqJxmVXmkdnm1bU=
12 | github.com/gabriel-vasile/mimetype v1.4.2 h1:w5qFW6JKBz9Y393Y4q372O9A7cUSequkh1Q7OhCmWKU=
13 | github.com/gabriel-vasile/mimetype v1.4.2/go.mod h1:zApsH/mKG4w07erKIaJPFiX0Tsq9BFQgN3qGY5GnNgA=
14 | github.com/gin-contrib/sse v0.1.0 h1:Y/yl/+YNO8GZSjAhjMsSuLt29uWRFHdHYUb5lYOV9qE=
15 | github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm+fLHvGI=
16 | github.com/gin-gonic/gin v1.9.1 h1:4idEAncQnU5cB7BeOkPtxjfCSye0AAm1R0RVIqJ+Jmg=
17 | github.com/gin-gonic/gin v1.9.1/go.mod h1:hPrL7YrpYKXt5YId3A/Tnip5kqbEAP+KLuI3SUcPTeU=
18 | github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s=
19 | github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4=
20 | github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA=
21 | github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY=
22 | github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY=
23 | github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY=
24 | github.com/go-playground/validator/v10 v10.14.0 h1:vgvQWe3XCz3gIeFDm/HnTIbj6UGmg/+t63MyGU2n5js=
25 | github.com/go-playground/validator/v10 v10.14.0/go.mod h1:9iXMNT7sEkjXb0I+enO7QXmzG6QCsPWY4zveKFVRSyU=
26 | github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
27 | github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
28 | github.com/golang-jwt/jwt v3.2.2+incompatible h1:IfV12K8xAKAnZqdXVzCZ+TOjboZ2keLg81eXfW3O+oY=
29 | github.com/golang-jwt/jwt v3.2.2+incompatible/go.mod h1:8pz2t5EyA70fFQQSrl6XZXzqecmYZeUEB8OUGHkxJ+I=
30 | github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
31 | github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
32 | github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
33 | github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
34 | github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
35 | github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
36 | github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
37 | github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
38 | github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
39 | github.com/klauspost/cpuid/v2 v2.2.4 h1:acbojRNwl3o09bUq+yDCtZFc1aiwaAAxtcn8YkZXnvk=
40 | github.com/klauspost/cpuid/v2 v2.2.4/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
41 | github.com/labstack/echo/v4 v4.9.0 h1:wPOF1CE6gvt/kmbMR4dGzWvHMPT+sAEUJOwOTtvITVY=
42 | github.com/labstack/echo/v4 v4.9.0/go.mod h1:xkCDAdFCIf8jsFQ5NnbK7oqaF/yU1A1X20Ltm0OvSks=
43 | github.com/labstack/gommon v0.3.1 h1:OomWaJXm7xR6L1HmEtGyQf26TEn7V6X88mktX9kee9o=
44 | github.com/labstack/gommon v0.3.1/go.mod h1:uW6kP17uPlLJsD3ijUYn3/M5bAxtlZhMI6m3MFxTMTM=
45 | github.com/leodido/go-urn v1.2.4 h1:XlAE/cm/ms7TE/VMVoduSpNBoyc2dOxHs5MZSwAN63Q=
46 | github.com/leodido/go-urn v1.2.4/go.mod h1:7ZrI8mTSeBSHl/UaRyKQW1qZeMgak41ANeCNaVckg+4=
47 | github.com/mattn/go-colorable v0.1.11 h1:nQ+aFkoE2TMGc0b68U2OKSexC+eq46+XwZzWXHRmPYs=
48 | github.com/mattn/go-colorable v0.1.11/go.mod h1:u5H1YNBxpqRaxsYJYSkiCWKzEfiAb1Gb520KVy5xxl4=
49 | github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
50 | github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=
51 | github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
52 | github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
53 | github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=
54 | github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
55 | github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
56 | github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
57 | github.com/pelletier/go-toml/v2 v2.0.8 h1:0ctb6s9mE31h0/lhu+J6OPmVeDxJn+kYnJc2jZR9tGQ=
58 | github.com/pelletier/go-toml/v2 v2.0.8/go.mod h1:vuYfssBdrU2XDZ9bYydBu6t+6a6PYNcZljzZR9VXg+4=
59 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
60 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
61 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
62 | github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
63 | github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
64 | github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
65 | github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
66 | github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
67 | github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
68 | github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
69 | github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
70 | github.com/stretchr/testify v1.8.3 h1:RP3t2pwF7cMEbC1dqtB6poj3niw/9gnV4Cjg5oW5gtY=
71 | github.com/stretchr/testify v1.8.3/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
72 | github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=
73 | github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
74 | github.com/ugorji/go/codec v1.2.11 h1:BMaWp1Bb6fHwEtbplGBGJ498wD+LKlNSl25MjdZY4dU=
75 | github.com/ugorji/go/codec v1.2.11/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg=
76 | github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
77 | github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
78 | github.com/valyala/fasttemplate v1.2.1 h1:TVEnxayobAdVkhQfrfes2IzOB6o+z4roRkPF52WA1u4=
79 | github.com/valyala/fasttemplate v1.2.1/go.mod h1:KHLXt3tVN2HBp8eijSv/kGJopbvo7S+qRAEEKiv+SiQ=
80 | golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
81 | golang.org/x/arch v0.3.0 h1:02VY4/ZcO/gBOH6PUaoiptASxtXU10jazRCP865E97k=
82 | golang.org/x/arch v0.3.0/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
83 | golang.org/x/crypto v0.9.0 h1:LF6fAI+IutBocDJ2OT0Q1g8plpYljMZ4+lty+dsqw3g=
84 | golang.org/x/crypto v0.9.0/go.mod h1:yrmDGqONDYtNj3tH8X9dzUun2m2lzPa9ngI6/RUPGR0=
85 | golang.org/x/net v0.10.0 h1:X2//UzNDwYmtCLn7To6G58Wr6f5ahEAQgKNzv9Y951M=
86 | golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
87 | golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
88 | golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
89 | golang.org/x/sys v0.0.0-20211103235746-7861aae1554b/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
90 | golang.org/x/sys v0.0.0-20220412211240-33da011f77ad/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
91 | golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
92 | golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
93 | golang.org/x/sys v0.8.0 h1:EBmGv8NaZBZTWvrbjNoL6HVt+IVy3QDQpJs7VRIw3tU=
94 | golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
95 | golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
96 | golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
97 | golang.org/x/time v0.0.0-20201208040808-7e3f01d25324 h1:Hir2P/De0WpUhtrKGGjvSb2YxUgyZ7EFOSLIcSSpiwE=
98 | golang.org/x/time v0.0.0-20201208040808-7e3f01d25324/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
99 | golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
100 | golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
101 | google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
102 | google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng=
103 | google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
104 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
105 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
106 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
107 | gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
108 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
109 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
110 | rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4=
111 |
--------------------------------------------------------------------------------
/examples/gophers/demo.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/olahol/melody/6bf86153260ddfd8ef1929323e0ee37cc3acd44a/examples/gophers/demo.gif
--------------------------------------------------------------------------------
/examples/gophers/index.html:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 | goofy gophers
5 |
20 |
21 |
22 |
69 |
70 |
71 |
--------------------------------------------------------------------------------
/examples/gophers/main.go:
--------------------------------------------------------------------------------
1 | package main
2 |
3 | import (
4 | "fmt"
5 | "net/http"
6 | "sync/atomic"
7 |
8 | "github.com/olahol/melody"
9 | )
10 |
11 | var idCounter atomic.Int64
12 |
13 | func main() {
14 | m := melody.New()
15 |
16 | http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
17 | http.ServeFile(w, r, "index.html")
18 | })
19 |
20 | http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
21 | m.HandleRequest(w, r)
22 | })
23 |
24 | m.HandleConnect(func(s *melody.Session) {
25 | id := idCounter.Add(1)
26 |
27 | s.Set("id", id)
28 |
29 | s.Write([]byte(fmt.Sprintf("iam %d", id)))
30 | })
31 |
32 | m.HandleDisconnect(func(s *melody.Session) {
33 | if id, ok := s.Get("id"); ok {
34 | m.BroadcastOthers([]byte(fmt.Sprintf("dis %d", id)), s)
35 | }
36 | })
37 |
38 | m.HandleMessage(func(s *melody.Session, msg []byte) {
39 | if id, ok := s.Get("id"); ok {
40 | m.BroadcastOthers([]byte(fmt.Sprintf("set %d %s", id, msg)), s)
41 | }
42 | })
43 |
44 | http.ListenAndServe(":5000", nil)
45 | }
46 |
--------------------------------------------------------------------------------
/examples/multichat/chan.html:
--------------------------------------------------------------------------------
1 |
2 |
3 | Melody example: chatting
4 |
5 |
6 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
52 |
53 |
54 |
--------------------------------------------------------------------------------
/examples/multichat/index.html:
--------------------------------------------------------------------------------
1 |
2 |
3 | Melody example: chatting
4 |
5 |
6 |
15 |
16 |
17 |
18 | Join a channel
19 |
20 |
21 |
22 |
33 |
34 |
35 |
--------------------------------------------------------------------------------
/examples/multichat/main.go:
--------------------------------------------------------------------------------
1 | package main
2 |
3 | import (
4 | "net/http"
5 |
6 | "github.com/olahol/melody"
7 | )
8 |
9 | func main() {
10 | m := melody.New()
11 |
12 | http.HandleFunc("GET /", func(w http.ResponseWriter, r *http.Request) {
13 | http.ServeFile(w, r, "index.html")
14 | })
15 |
16 | http.HandleFunc("GET /channel/{chan}", func(w http.ResponseWriter, r *http.Request) {
17 | http.ServeFile(w, r, "chan.html")
18 | })
19 |
20 | http.HandleFunc("GET /channel/{chan}/ws", func(w http.ResponseWriter, r *http.Request) {
21 | m.HandleRequest(w, r)
22 | })
23 |
24 | m.HandleMessage(func(s *melody.Session, msg []byte) {
25 | m.BroadcastFilter(msg, func(q *melody.Session) bool {
26 | return q.Request.URL.Path == s.Request.URL.Path
27 | })
28 | })
29 |
30 | http.ListenAndServe(":5000", nil)
31 | }
32 |
--------------------------------------------------------------------------------
/go.mod:
--------------------------------------------------------------------------------
1 | module github.com/olahol/melody
2 |
3 | go 1.19
4 |
5 | require (
6 | github.com/gorilla/websocket v1.5.0
7 | github.com/stretchr/testify v1.8.0
8 | )
9 |
10 | require (
11 | github.com/davecgh/go-spew v1.1.1 // indirect
12 | github.com/pmezard/go-difflib v1.0.0 // indirect
13 | gopkg.in/yaml.v3 v3.0.1 // indirect
14 | )
15 |
--------------------------------------------------------------------------------
/go.sum:
--------------------------------------------------------------------------------
1 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
2 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
3 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
4 | github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
5 | github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
6 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
7 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
8 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
9 | github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
10 | github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
11 | github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
12 | github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
13 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
14 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
15 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
16 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
17 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
18 |
--------------------------------------------------------------------------------
/hub.go:
--------------------------------------------------------------------------------
1 | package melody
2 |
3 | import (
4 | "sync"
5 | "sync/atomic"
6 | )
7 |
8 | type sessionSet struct {
9 | mu sync.RWMutex
10 | members map[*Session]struct{}
11 | }
12 |
13 | func (ss *sessionSet) add(s *Session) {
14 | ss.mu.Lock()
15 | defer ss.mu.Unlock()
16 |
17 | ss.members[s] = struct{}{}
18 | }
19 |
20 | func (ss *sessionSet) del(s *Session) {
21 | ss.mu.Lock()
22 | defer ss.mu.Unlock()
23 |
24 | delete(ss.members, s)
25 | }
26 |
27 | func (ss *sessionSet) clear() {
28 | ss.mu.Lock()
29 | defer ss.mu.Unlock()
30 |
31 | ss.members = make(map[*Session]struct{})
32 | }
33 |
34 | func (ss *sessionSet) each(cb func(*Session)) {
35 | ss.mu.RLock()
36 | defer ss.mu.RUnlock()
37 |
38 | for s := range ss.members {
39 | cb(s)
40 | }
41 | }
42 |
43 | func (ss *sessionSet) len() int {
44 | ss.mu.RLock()
45 | defer ss.mu.RUnlock()
46 |
47 | return len(ss.members)
48 | }
49 |
50 | func (ss *sessionSet) all() []*Session {
51 | ss.mu.RLock()
52 | defer ss.mu.RUnlock()
53 |
54 | s := make([]*Session, 0, len(ss.members))
55 | for k := range ss.members {
56 | s = append(s, k)
57 | }
58 |
59 | return s
60 | }
61 |
62 | type hub struct {
63 | sessions sessionSet
64 | broadcast chan envelope
65 | register chan *Session
66 | unregister chan *Session
67 | exit chan envelope
68 | open atomic.Bool
69 | }
70 |
71 | func newHub() *hub {
72 | return &hub{
73 | sessions: sessionSet{
74 | members: make(map[*Session]struct{}),
75 | },
76 | broadcast: make(chan envelope),
77 | register: make(chan *Session),
78 | unregister: make(chan *Session),
79 | exit: make(chan envelope),
80 | }
81 | }
82 |
83 | func (h *hub) run() {
84 | h.open.Store(true)
85 |
86 | loop:
87 | for {
88 | select {
89 | case s := <-h.register:
90 | h.sessions.add(s)
91 | case s := <-h.unregister:
92 | h.sessions.del(s)
93 | case m := <-h.broadcast:
94 | h.sessions.each(func(s *Session) {
95 | if m.filter == nil {
96 | s.writeMessage(m)
97 | } else if m.filter(s) {
98 | s.writeMessage(m)
99 | }
100 | })
101 | case m := <-h.exit:
102 | h.open.Store(false)
103 |
104 | h.sessions.each(func(s *Session) {
105 | s.writeMessage(m)
106 | s.Close()
107 | })
108 |
109 | h.sessions.clear()
110 |
111 | break loop
112 | }
113 | }
114 | }
115 |
116 | func (h *hub) closed() bool {
117 | return !h.open.Load()
118 | }
119 |
120 | func (h *hub) len() int {
121 | return h.sessions.len()
122 | }
123 |
124 | func (h *hub) all() []*Session {
125 | return h.sessions.all()
126 | }
127 |
--------------------------------------------------------------------------------
/melody.go:
--------------------------------------------------------------------------------
1 | package melody
2 |
3 | import (
4 | "net/http"
5 | "sync"
6 |
7 | "github.com/gorilla/websocket"
8 | )
9 |
10 | // Close codes defined in RFC 6455, section 11.7.
11 | // Duplicate of codes from gorilla/websocket for convenience.
12 | const (
13 | CloseNormalClosure = 1000
14 | CloseGoingAway = 1001
15 | CloseProtocolError = 1002
16 | CloseUnsupportedData = 1003
17 | CloseNoStatusReceived = 1005
18 | CloseAbnormalClosure = 1006
19 | CloseInvalidFramePayloadData = 1007
20 | ClosePolicyViolation = 1008
21 | CloseMessageTooBig = 1009
22 | CloseMandatoryExtension = 1010
23 | CloseInternalServerErr = 1011
24 | CloseServiceRestart = 1012
25 | CloseTryAgainLater = 1013
26 | CloseTLSHandshake = 1015
27 | )
28 |
29 | type handleMessageFunc func(*Session, []byte)
30 | type handleErrorFunc func(*Session, error)
31 | type handleCloseFunc func(*Session, int, string) error
32 | type handleSessionFunc func(*Session)
33 | type filterFunc func(*Session) bool
34 |
35 | // Melody implements a websocket manager.
36 | type Melody struct {
37 | Config *Config
38 | Upgrader *websocket.Upgrader
39 | messageHandler handleMessageFunc
40 | messageHandlerBinary handleMessageFunc
41 | messageSentHandler handleMessageFunc
42 | messageSentHandlerBinary handleMessageFunc
43 | errorHandler handleErrorFunc
44 | closeHandler handleCloseFunc
45 | connectHandler handleSessionFunc
46 | disconnectHandler handleSessionFunc
47 | pongHandler handleSessionFunc
48 | hub *hub
49 | }
50 |
51 | // New creates a new melody instance with default Upgrader and Config.
52 | func New() *Melody {
53 | upgrader := &websocket.Upgrader{
54 | ReadBufferSize: 1024,
55 | WriteBufferSize: 1024,
56 | CheckOrigin: func(r *http.Request) bool { return true },
57 | }
58 |
59 | hub := newHub()
60 |
61 | go hub.run()
62 |
63 | return &Melody{
64 | Config: newConfig(),
65 | Upgrader: upgrader,
66 | messageHandler: func(*Session, []byte) {},
67 | messageHandlerBinary: func(*Session, []byte) {},
68 | messageSentHandler: func(*Session, []byte) {},
69 | messageSentHandlerBinary: func(*Session, []byte) {},
70 | errorHandler: func(*Session, error) {},
71 | closeHandler: nil,
72 | connectHandler: func(*Session) {},
73 | disconnectHandler: func(*Session) {},
74 | pongHandler: func(*Session) {},
75 | hub: hub,
76 | }
77 | }
78 |
79 | // HandleConnect fires fn when a session connects.
80 | func (m *Melody) HandleConnect(fn func(*Session)) {
81 | m.connectHandler = fn
82 | }
83 |
84 | // HandleDisconnect fires fn when a session disconnects.
85 | func (m *Melody) HandleDisconnect(fn func(*Session)) {
86 | m.disconnectHandler = fn
87 | }
88 |
89 | // HandlePong fires fn when a pong is received from a session.
90 | func (m *Melody) HandlePong(fn func(*Session)) {
91 | m.pongHandler = fn
92 | }
93 |
94 | // HandleMessage fires fn when a text message comes in.
95 | // NOTE: by default Melody handles messages sequentially for each
96 | // session. This has the effect that a message handler exceeding the
97 | // read deadline (Config.PongWait, by default 1 minute) will time out
98 | // the session. Concurrent message handling can be turned on by setting
99 | // Config.ConcurrentMessageHandling to true.
100 | func (m *Melody) HandleMessage(fn func(*Session, []byte)) {
101 | m.messageHandler = fn
102 | }
103 |
104 | // HandleMessageBinary fires fn when a binary message comes in.
105 | func (m *Melody) HandleMessageBinary(fn func(*Session, []byte)) {
106 | m.messageHandlerBinary = fn
107 | }
108 |
109 | // HandleSentMessage fires fn when a text message is successfully sent.
110 | func (m *Melody) HandleSentMessage(fn func(*Session, []byte)) {
111 | m.messageSentHandler = fn
112 | }
113 |
114 | // HandleSentMessageBinary fires fn when a binary message is successfully sent.
115 | func (m *Melody) HandleSentMessageBinary(fn func(*Session, []byte)) {
116 | m.messageSentHandlerBinary = fn
117 | }
118 |
119 | // HandleError fires fn when a session has an error.
120 | func (m *Melody) HandleError(fn func(*Session, error)) {
121 | m.errorHandler = fn
122 | }
123 |
124 | // HandleClose sets the handler for close messages received from the session.
125 | // The code argument to h is the received close code or CloseNoStatusReceived
126 | // if the close message is empty. The default close handler sends a close frame
127 | // back to the session.
128 | //
129 | // The application must read the connection to process close messages as
130 | // described in the section on Control Frames above.
131 | //
132 | // The connection read methods return a CloseError when a close frame is
133 | // received. Most applications should handle close messages as part of their
134 | // normal error handling. Applications should only set a close handler when the
135 | // application must perform some action before sending a close frame back to
136 | // the session.
137 | func (m *Melody) HandleClose(fn func(*Session, int, string) error) {
138 | if fn != nil {
139 | m.closeHandler = fn
140 | }
141 | }
142 |
143 | // HandleRequest upgrades http requests to websocket connections and dispatches them to be handled by the melody instance.
144 | func (m *Melody) HandleRequest(w http.ResponseWriter, r *http.Request) error {
145 | return m.HandleRequestWithKeys(w, r, nil)
146 | }
147 |
148 | // HandleRequestWithKeys does the same as HandleRequest but populates session.Keys with keys.
149 | func (m *Melody) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, keys map[string]any) error {
150 | if m.hub.closed() {
151 | return ErrClosed
152 | }
153 |
154 | conn, err := m.Upgrader.Upgrade(w, r, w.Header())
155 |
156 | if err != nil {
157 | return err
158 | }
159 |
160 | session := &Session{
161 | Request: r,
162 | Keys: keys,
163 | conn: conn,
164 | output: make(chan envelope, m.Config.MessageBufferSize),
165 | outputDone: make(chan struct{}),
166 | melody: m,
167 | open: true,
168 | rwmutex: &sync.RWMutex{},
169 | }
170 |
171 | m.hub.register <- session
172 |
173 | m.connectHandler(session)
174 |
175 | go session.writePump()
176 |
177 | session.readPump()
178 |
179 | if !m.hub.closed() {
180 | m.hub.unregister <- session
181 | }
182 |
183 | session.close()
184 |
185 | m.disconnectHandler(session)
186 |
187 | return nil
188 | }
189 |
190 | // Broadcast broadcasts a text message to all sessions.
191 | func (m *Melody) Broadcast(msg []byte) error {
192 | if m.hub.closed() {
193 | return ErrClosed
194 | }
195 |
196 | message := envelope{t: websocket.TextMessage, msg: msg}
197 | m.hub.broadcast <- message
198 |
199 | return nil
200 | }
201 |
202 | // BroadcastFilter broadcasts a text message to all sessions that fn returns true for.
203 | func (m *Melody) BroadcastFilter(msg []byte, fn func(*Session) bool) error {
204 | if m.hub.closed() {
205 | return ErrClosed
206 | }
207 |
208 | message := envelope{t: websocket.TextMessage, msg: msg, filter: fn}
209 | m.hub.broadcast <- message
210 |
211 | return nil
212 | }
213 |
214 | // BroadcastOthers broadcasts a text message to all sessions except session s.
215 | func (m *Melody) BroadcastOthers(msg []byte, s *Session) error {
216 | return m.BroadcastFilter(msg, func(q *Session) bool {
217 | return s != q
218 | })
219 | }
220 |
221 | // BroadcastMultiple broadcasts a text message to multiple sessions given in the sessions slice.
222 | func (m *Melody) BroadcastMultiple(msg []byte, sessions []*Session) error {
223 | for _, sess := range sessions {
224 | if writeErr := sess.Write(msg); writeErr != nil {
225 | return writeErr
226 | }
227 | }
228 | return nil
229 | }
230 |
231 | // BroadcastBinary broadcasts a binary message to all sessions.
232 | func (m *Melody) BroadcastBinary(msg []byte) error {
233 | if m.hub.closed() {
234 | return ErrClosed
235 | }
236 |
237 | message := envelope{t: websocket.BinaryMessage, msg: msg}
238 | m.hub.broadcast <- message
239 |
240 | return nil
241 | }
242 |
243 | // BroadcastBinaryFilter broadcasts a binary message to all sessions that fn returns true for.
244 | func (m *Melody) BroadcastBinaryFilter(msg []byte, fn func(*Session) bool) error {
245 | if m.hub.closed() {
246 | return ErrClosed
247 | }
248 |
249 | message := envelope{t: websocket.BinaryMessage, msg: msg, filter: fn}
250 | m.hub.broadcast <- message
251 |
252 | return nil
253 | }
254 |
255 | // BroadcastBinaryOthers broadcasts a binary message to all sessions except session s.
256 | func (m *Melody) BroadcastBinaryOthers(msg []byte, s *Session) error {
257 | return m.BroadcastBinaryFilter(msg, func(q *Session) bool {
258 | return s != q
259 | })
260 | }
261 |
262 | // Sessions returns all sessions. An error is returned if the melody session is closed.
263 | func (m *Melody) Sessions() ([]*Session, error) {
264 | if m.hub.closed() {
265 | return nil, ErrClosed
266 | }
267 | return m.hub.all(), nil
268 | }
269 |
270 | // Close closes the melody instance and all connected sessions.
271 | func (m *Melody) Close() error {
272 | if m.hub.closed() {
273 | return ErrClosed
274 | }
275 |
276 | m.hub.exit <- envelope{t: websocket.CloseMessage, msg: []byte{}}
277 |
278 | return nil
279 | }
280 |
281 | // CloseWithMsg closes the melody instance with the given close payload and all connected sessions.
282 | // Use the FormatCloseMessage function to format a proper close message payload.
283 | func (m *Melody) CloseWithMsg(msg []byte) error {
284 | if m.hub.closed() {
285 | return ErrClosed
286 | }
287 |
288 | m.hub.exit <- envelope{t: websocket.CloseMessage, msg: msg}
289 |
290 | return nil
291 | }
292 |
293 | // Len return the number of connected sessions.
294 | func (m *Melody) Len() int {
295 | return m.hub.len()
296 | }
297 |
298 | // IsClosed returns the status of the melody instance.
299 | func (m *Melody) IsClosed() bool {
300 | return m.hub.closed()
301 | }
302 |
303 | // FormatCloseMessage formats closeCode and text as a WebSocket close message.
304 | func FormatCloseMessage(closeCode int, text string) []byte {
305 | return websocket.FormatCloseMessage(closeCode, text)
306 | }
307 |
--------------------------------------------------------------------------------
/melody_test.go:
--------------------------------------------------------------------------------
1 | package melody
2 |
3 | import (
4 | "bytes"
5 | "errors"
6 | "math/rand"
7 | "net/http"
8 | "net/http/httptest"
9 | "os"
10 | "strconv"
11 | "strings"
12 | "sync"
13 | "sync/atomic"
14 | "testing"
15 | "testing/quick"
16 | "time"
17 |
18 | "github.com/gorilla/websocket"
19 | "github.com/stretchr/testify/assert"
20 | )
21 |
22 | var TestMsg = []byte("test")
23 |
24 | type TestServer struct {
25 | withKeys bool
26 | m *Melody
27 | }
28 |
29 | func NewTestServerHandler(handler handleMessageFunc) *TestServer {
30 | m := New()
31 | m.HandleMessage(handler)
32 | return &TestServer{
33 | m: m,
34 | }
35 | }
36 |
37 | func NewTestServer() *TestServer {
38 | m := New()
39 | return &TestServer{
40 | m: m,
41 | }
42 | }
43 |
44 | func (s *TestServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
45 | if s.withKeys {
46 | s.m.HandleRequestWithKeys(w, r, make(map[string]any))
47 | } else {
48 | s.m.HandleRequest(w, r)
49 | }
50 | }
51 |
52 | func NewDialer(url string) (*websocket.Conn, error) {
53 | dialer := &websocket.Dialer{}
54 | conn, _, err := dialer.Dial(strings.Replace(url, "http", "ws", 1), nil)
55 | return conn, err
56 | }
57 |
58 | func MustNewDialer(url string) *websocket.Conn {
59 | conn, err := NewDialer(url)
60 |
61 | if err != nil {
62 | panic("could not dail websocket")
63 | }
64 |
65 | return conn
66 | }
67 |
68 | func TestEcho(t *testing.T) {
69 | ws := NewTestServerHandler(func(session *Session, msg []byte) {
70 | session.Write(msg)
71 | })
72 | server := httptest.NewServer(ws)
73 | defer server.Close()
74 |
75 | fn := func(msg string) bool {
76 | conn := MustNewDialer(server.URL)
77 | defer conn.Close()
78 |
79 | conn.WriteMessage(websocket.TextMessage, []byte(msg))
80 |
81 | _, ret, err := conn.ReadMessage()
82 |
83 | assert.Nil(t, err)
84 |
85 | assert.Equal(t, msg, string(ret))
86 |
87 | return true
88 | }
89 |
90 | err := quick.Check(fn, nil)
91 |
92 | assert.Nil(t, err)
93 | }
94 |
95 | func TestEchoBinary(t *testing.T) {
96 | ws := NewTestServerHandler(func(session *Session, msg []byte) {
97 | session.WriteBinary(msg)
98 | })
99 | server := httptest.NewServer(ws)
100 | defer server.Close()
101 |
102 | fn := func(msg string) bool {
103 | conn := MustNewDialer(server.URL)
104 | defer conn.Close()
105 |
106 | conn.WriteMessage(websocket.TextMessage, []byte(msg))
107 |
108 | _, ret, err := conn.ReadMessage()
109 |
110 | assert.Nil(t, err)
111 |
112 | assert.True(t, bytes.Equal([]byte(msg), ret))
113 |
114 | return true
115 | }
116 |
117 | err := quick.Check(fn, nil)
118 |
119 | assert.Nil(t, err)
120 | }
121 |
122 | func TestWriteClosedServer(t *testing.T) {
123 | done := make(chan bool)
124 |
125 | ws := NewTestServer()
126 |
127 | server := httptest.NewServer(ws)
128 | defer server.Close()
129 |
130 | ws.m.HandleConnect(func(s *Session) {
131 | s.Close()
132 | })
133 |
134 | ws.m.HandleDisconnect(func(s *Session) {
135 | err := s.Write(TestMsg)
136 |
137 | assert.NotNil(t, err)
138 | close(done)
139 | })
140 |
141 | conn := MustNewDialer(server.URL)
142 | conn.ReadMessage()
143 | defer conn.Close()
144 |
145 | <-done
146 | }
147 |
148 | func TestWriteClosedClient(t *testing.T) {
149 | done := make(chan bool)
150 |
151 | ws := NewTestServer()
152 |
153 | server := httptest.NewServer(ws)
154 | defer server.Close()
155 |
156 | ws.m.HandleDisconnect(func(s *Session) {
157 | err := s.Write(TestMsg)
158 |
159 | assert.NotNil(t, err)
160 | close(done)
161 | })
162 |
163 | conn := MustNewDialer(server.URL)
164 | conn.Close()
165 |
166 | <-done
167 | }
168 |
169 | func TestUpgrader(t *testing.T) {
170 | ws := NewTestServer()
171 | ws.m.HandleMessage(func(session *Session, msg []byte) {
172 | session.Write(msg)
173 | })
174 |
175 | server := httptest.NewServer(ws)
176 | defer server.Close()
177 |
178 | ws.m.Upgrader = &websocket.Upgrader{
179 | ReadBufferSize: 1024,
180 | WriteBufferSize: 1024,
181 | CheckOrigin: func(r *http.Request) bool { return false },
182 | }
183 |
184 | _, err := NewDialer(server.URL)
185 |
186 | assert.ErrorIs(t, err, websocket.ErrBadHandshake)
187 | }
188 |
189 | func TestBroadcast(t *testing.T) {
190 | n := 10
191 | msg := "test"
192 |
193 | test := func(h func(*TestServer), w func(*websocket.Conn)) {
194 |
195 | ws := NewTestServer()
196 |
197 | h(ws)
198 |
199 | server := httptest.NewServer(ws)
200 | defer server.Close()
201 |
202 | conn := MustNewDialer(server.URL)
203 | defer conn.Close()
204 |
205 | listeners := make([]*websocket.Conn, n)
206 | for i := range listeners {
207 | listener := MustNewDialer(server.URL)
208 | listeners[i] = listener
209 | defer listeners[i].Close()
210 | }
211 |
212 | w(conn)
213 |
214 | for _, listener := range listeners {
215 | _, ret, err := listener.ReadMessage()
216 |
217 | assert.Nil(t, err)
218 |
219 | assert.Equal(t, msg, string(ret))
220 | }
221 | }
222 |
223 | test(func(ws *TestServer) {
224 | ws.m.HandleMessage(func(s *Session, msg []byte) {
225 | ws.m.Broadcast(msg)
226 | })
227 | }, func(conn *websocket.Conn) {
228 | conn.WriteMessage(websocket.TextMessage, []byte(msg))
229 | })
230 |
231 | test(func(ws *TestServer) {
232 | ws.m.HandleMessageBinary(func(s *Session, msg []byte) {
233 | ws.m.BroadcastBinary(msg)
234 | })
235 | }, func(conn *websocket.Conn) {
236 | conn.WriteMessage(websocket.BinaryMessage, []byte(msg))
237 | })
238 |
239 | test(func(ws *TestServer) {
240 | ws.m.HandleMessage(func(s *Session, msg []byte) {
241 | ws.m.BroadcastFilter(msg, func(s *Session) bool {
242 | return true
243 | })
244 | })
245 | }, func(conn *websocket.Conn) {
246 | conn.WriteMessage(websocket.TextMessage, []byte(msg))
247 | })
248 |
249 | test(func(ws *TestServer) {
250 | ws.m.HandleMessageBinary(func(s *Session, msg []byte) {
251 | ws.m.BroadcastBinaryFilter(msg, func(s *Session) bool {
252 | return true
253 | })
254 | })
255 | }, func(conn *websocket.Conn) {
256 | conn.WriteMessage(websocket.BinaryMessage, []byte(msg))
257 | })
258 |
259 | test(func(ws *TestServer) {
260 | ws.m.HandleMessage(func(s *Session, msg []byte) {
261 | ws.m.BroadcastOthers(msg, s)
262 | })
263 | }, func(conn *websocket.Conn) {
264 | conn.WriteMessage(websocket.TextMessage, []byte(msg))
265 | })
266 |
267 | test(func(ws *TestServer) {
268 | ws.m.HandleMessageBinary(func(s *Session, msg []byte) {
269 | ws.m.BroadcastBinaryOthers(msg, s)
270 | })
271 | }, func(conn *websocket.Conn) {
272 | conn.WriteMessage(websocket.BinaryMessage, []byte(msg))
273 | })
274 |
275 | test(func(ws *TestServer) {
276 | ws.m.HandleMessage(func(s *Session, msg []byte) {
277 | ss, _ := ws.m.Sessions()
278 | ws.m.BroadcastMultiple(msg, ss)
279 | })
280 | }, func(conn *websocket.Conn) {
281 | conn.WriteMessage(websocket.TextMessage, []byte(msg))
282 | })
283 | }
284 |
285 | func TestClose(t *testing.T) {
286 | ws := NewTestServer()
287 |
288 | server := httptest.NewServer(ws)
289 | defer server.Close()
290 |
291 | n := 10
292 |
293 | conns := make([]*websocket.Conn, n)
294 | for i := range conns {
295 | conn := MustNewDialer(server.URL)
296 | conns[i] = conn
297 | defer conns[i].Close()
298 | }
299 |
300 | q := make(chan bool)
301 | ws.m.HandleDisconnect(func(s *Session) {
302 | q <- true
303 | })
304 |
305 | ws.m.Close()
306 |
307 | for _, conn := range conns {
308 | conn.ReadMessage()
309 | }
310 |
311 | assert.Zero(t, ws.m.Len())
312 |
313 | m := 0
314 | for range q {
315 | m += 1
316 | if m == n {
317 | break
318 | }
319 | }
320 | }
321 |
322 | func TestLen(t *testing.T) {
323 | rand.Seed(time.Now().UnixNano())
324 |
325 | connect := int(rand.Int31n(100))
326 | disconnect := rand.Float32()
327 | conns := make([]*websocket.Conn, connect)
328 |
329 | defer func() {
330 | for _, conn := range conns {
331 | if conn != nil {
332 | conn.Close()
333 | }
334 | }
335 | }()
336 |
337 | ws := NewTestServer()
338 |
339 | server := httptest.NewServer(ws)
340 | defer server.Close()
341 |
342 | disconnected := 0
343 | for i := 0; i < connect; i++ {
344 | conn := MustNewDialer(server.URL)
345 |
346 | if rand.Float32() < disconnect {
347 | conns[i] = nil
348 | disconnected++
349 | conn.Close()
350 | continue
351 | }
352 |
353 | conns[i] = conn
354 | }
355 |
356 | time.Sleep(time.Millisecond)
357 |
358 | connected := connect - disconnected
359 |
360 | assert.Equal(t, ws.m.Len(), connected)
361 | }
362 |
363 | func TestSessions(t *testing.T) {
364 | rand.Seed(time.Now().UnixNano())
365 |
366 | connect := int(rand.Int31n(100))
367 | disconnect := rand.Float32()
368 | conns := make([]*websocket.Conn, connect)
369 | defer func() {
370 | for _, conn := range conns {
371 | if conn != nil {
372 | conn.Close()
373 | }
374 | }
375 | }()
376 |
377 | ws := NewTestServer()
378 | server := httptest.NewServer(ws)
379 | defer server.Close()
380 |
381 | disconnected := 0
382 | for i := 0; i < connect; i++ {
383 | conn, err := NewDialer(server.URL)
384 |
385 | if err != nil {
386 | t.Error(err)
387 | }
388 |
389 | if rand.Float32() < disconnect {
390 | conns[i] = nil
391 | disconnected++
392 | conn.Close()
393 | continue
394 | }
395 |
396 | conns[i] = conn
397 | }
398 |
399 | time.Sleep(time.Millisecond)
400 |
401 | connected := connect - disconnected
402 |
403 | ss, err := ws.m.Sessions()
404 |
405 | assert.Nil(t, err)
406 |
407 | assert.Equal(t, len(ss), connected)
408 | }
409 |
410 | func TestPingPong(t *testing.T) {
411 | done := make(chan bool)
412 |
413 | ws := NewTestServer()
414 | ws.m.Config.PingPeriod = time.Millisecond
415 |
416 | ws.m.HandlePong(func(s *Session) {
417 | close(done)
418 | })
419 |
420 | server := httptest.NewServer(ws)
421 | defer server.Close()
422 |
423 | conn := MustNewDialer(server.URL)
424 | defer conn.Close()
425 |
426 | go conn.NextReader()
427 |
428 | <-done
429 | }
430 |
431 | func TestHandleClose(t *testing.T) {
432 | done := make(chan bool)
433 |
434 | ws := NewTestServer()
435 | ws.m.Config.PingPeriod = time.Millisecond
436 |
437 | ws.m.HandleClose(func(s *Session, code int, text string) error {
438 | close(done)
439 | return nil
440 | })
441 |
442 | server := httptest.NewServer(ws)
443 | defer server.Close()
444 |
445 | conn := MustNewDialer(server.URL)
446 |
447 | conn.WriteMessage(websocket.CloseMessage, nil)
448 |
449 | <-done
450 | }
451 |
452 | func TestHandleError(t *testing.T) {
453 | done := make(chan bool)
454 |
455 | ws := NewTestServer()
456 |
457 | ws.m.HandleError(func(s *Session, err error) {
458 | var closeError *websocket.CloseError
459 | assert.ErrorAs(t, err, &closeError)
460 | close(done)
461 | })
462 |
463 | server := httptest.NewServer(ws)
464 | defer server.Close()
465 |
466 | conn := MustNewDialer(server.URL)
467 |
468 | conn.Close()
469 |
470 | <-done
471 | }
472 |
473 | func TestHandleErrorWrite(t *testing.T) {
474 | writeError := make(chan struct{})
475 | disconnect := make(chan struct{})
476 |
477 | ws := NewTestServer()
478 | ws.m.Config.WriteWait = 0
479 |
480 | ws.m.HandleConnect(func(s *Session) {
481 | err := s.Write(TestMsg)
482 | assert.Nil(t, err)
483 | })
484 |
485 | ws.m.HandleError(func(s *Session, err error) {
486 | assert.NotNil(t, err)
487 |
488 | if os.IsTimeout(err) {
489 | close(writeError)
490 | }
491 | })
492 |
493 | ws.m.HandleDisconnect(func(s *Session) {
494 | close(disconnect)
495 | })
496 |
497 | server := httptest.NewServer(ws)
498 | defer server.Close()
499 |
500 | conn := MustNewDialer(server.URL)
501 | defer conn.Close()
502 |
503 | go conn.NextReader()
504 |
505 | <-writeError
506 | <-disconnect
507 | }
508 |
509 | func TestErrClosed(t *testing.T) {
510 | res := make(chan *Session)
511 |
512 | ws := NewTestServer()
513 |
514 | ws.m.HandleConnect(func(s *Session) {
515 | ws.m.CloseWithMsg(TestMsg)
516 | })
517 |
518 | ws.m.HandleDisconnect(func(s *Session) {
519 | res <- s
520 | })
521 |
522 | server := httptest.NewServer(ws)
523 | defer server.Close()
524 |
525 | conn := MustNewDialer(server.URL)
526 | defer conn.Close()
527 |
528 | go conn.ReadMessage()
529 |
530 | s := <-res
531 |
532 | assert.True(t, s.IsClosed())
533 | assert.True(t, ws.m.IsClosed())
534 | _, err := ws.m.Sessions()
535 | assert.ErrorIs(t, err, ErrClosed)
536 | assert.ErrorIs(t, ws.m.Close(), ErrClosed)
537 | assert.ErrorIs(t, ws.m.CloseWithMsg(TestMsg), ErrClosed)
538 |
539 | assert.ErrorIs(t, ws.m.Broadcast(TestMsg), ErrClosed)
540 | assert.ErrorIs(t, ws.m.BroadcastBinary(TestMsg), ErrClosed)
541 | assert.ErrorIs(t, ws.m.BroadcastFilter(TestMsg, func(s *Session) bool { return true }), ErrClosed)
542 | assert.ErrorIs(t, ws.m.BroadcastBinaryFilter(TestMsg, func(s *Session) bool { return true }), ErrClosed)
543 | assert.ErrorIs(t, ws.m.HandleRequest(nil, nil), ErrClosed)
544 | }
545 |
546 | func TestErrSessionClosed(t *testing.T) {
547 | res := make(chan *Session)
548 |
549 | ws := NewTestServer()
550 |
551 | ws.m.HandleConnect(func(s *Session) {
552 | s.CloseWithMsg(TestMsg)
553 | })
554 |
555 | ws.m.HandleDisconnect(func(s *Session) {
556 | res <- s
557 | })
558 |
559 | server := httptest.NewServer(ws)
560 | defer server.Close()
561 |
562 | conn := MustNewDialer(server.URL)
563 | defer conn.Close()
564 |
565 | go conn.ReadMessage()
566 |
567 | s := <-res
568 |
569 | assert.True(t, s.IsClosed())
570 | assert.ErrorIs(t, s.Write(TestMsg), ErrSessionClosed)
571 | assert.ErrorIs(t, s.WriteBinary(TestMsg), ErrSessionClosed)
572 | assert.ErrorIs(t, s.CloseWithMsg(TestMsg), ErrSessionClosed)
573 | assert.ErrorIs(t, s.Close(), ErrSessionClosed)
574 | assert.ErrorIs(t, ws.m.BroadcastMultiple(TestMsg, []*Session{s}), ErrSessionClosed)
575 |
576 | assert.ErrorIs(t, s.writeRaw(envelope{}), ErrWriteClosed)
577 | s.writeMessage(envelope{})
578 | }
579 |
580 | func TestErrMessageBufferFull(t *testing.T) {
581 | done := make(chan bool)
582 |
583 | ws := NewTestServerHandler(func(session *Session, msg []byte) {
584 | session.Write(msg)
585 | session.Write(msg)
586 | })
587 | ws.m.Config.MessageBufferSize = 0
588 | ws.m.HandleError(func(s *Session, err error) {
589 | if errors.Is(err, ErrMessageBufferFull) {
590 | close(done)
591 | }
592 | })
593 | server := httptest.NewServer(ws)
594 | defer server.Close()
595 |
596 | conn := MustNewDialer(server.URL)
597 | defer conn.Close()
598 |
599 | conn.WriteMessage(websocket.TextMessage, TestMsg)
600 |
601 | <-done
602 | }
603 |
604 | func TestSessionKeys(t *testing.T) {
605 | ws := NewTestServer()
606 |
607 | ws.m.HandleConnect(func(session *Session) {
608 | session.Set("stamp", time.Now().UnixNano())
609 | })
610 | ws.m.HandleMessage(func(session *Session, msg []byte) {
611 | stamp := session.MustGet("stamp").(int64)
612 | session.Write([]byte(strconv.Itoa(int(stamp))))
613 | })
614 | server := httptest.NewServer(ws)
615 | defer server.Close()
616 |
617 | fn := func(msg string) bool {
618 | conn := MustNewDialer(server.URL)
619 | defer conn.Close()
620 |
621 | conn.WriteMessage(websocket.TextMessage, []byte(msg))
622 |
623 | _, ret, err := conn.ReadMessage()
624 |
625 | assert.Nil(t, err)
626 |
627 | stamp, err := strconv.Atoi(string(ret))
628 |
629 | assert.Nil(t, err)
630 |
631 | diff := int(time.Now().UnixNano()) - stamp
632 |
633 | assert.Greater(t, diff, 0)
634 |
635 | return true
636 | }
637 |
638 | assert.Nil(t, quick.Check(fn, nil))
639 | }
640 |
641 | func TestSessionKeysConcurrent(t *testing.T) {
642 | ss := make(chan *Session)
643 |
644 | ws := NewTestServer()
645 |
646 | ws.m.HandleConnect(func(s *Session) {
647 | ss <- s
648 | })
649 |
650 | server := httptest.NewServer(ws)
651 | defer server.Close()
652 |
653 | conn := MustNewDialer(server.URL)
654 | defer conn.Close()
655 |
656 | s := <-ss
657 |
658 | var wg sync.WaitGroup
659 |
660 | for i := 0; i < 100; i++ {
661 | wg.Add(1)
662 |
663 | go func() {
664 | s.Set("test", TestMsg)
665 |
666 | v1, exists := s.Get("test")
667 |
668 | assert.True(t, exists)
669 | assert.Equal(t, v1, TestMsg)
670 |
671 | v2 := s.MustGet("test")
672 |
673 | assert.Equal(t, v1, v2)
674 |
675 | wg.Done()
676 | }()
677 | }
678 |
679 | wg.Wait()
680 |
681 | for i := 0; i < 100; i++ {
682 | wg.Add(1)
683 |
684 | go func() {
685 | s.UnSet("test")
686 |
687 | _, exists := s.Get("test")
688 |
689 | assert.False(t, exists)
690 |
691 | wg.Done()
692 | }()
693 | }
694 |
695 | wg.Wait()
696 | }
697 |
698 | func TestMisc(t *testing.T) {
699 | res := make(chan *Session)
700 |
701 | ws := NewTestServer()
702 |
703 | ws.m.HandleConnect(func(s *Session) {
704 | res <- s
705 | })
706 |
707 | server := httptest.NewServer(ws)
708 | defer server.Close()
709 |
710 | conn := MustNewDialer(server.URL)
711 | defer conn.Close()
712 |
713 | go conn.ReadMessage()
714 |
715 | s := <-res
716 |
717 | assert.Contains(t, s.LocalAddr().String(), "127.0.0.1")
718 | assert.Contains(t, s.RemoteAddr().String(), "127.0.0.1")
719 | assert.Equal(t, FormatCloseMessage(websocket.CloseMessage, "test"), websocket.FormatCloseMessage(websocket.CloseMessage, "test"))
720 | assert.Panics(t, func() {
721 | s.MustGet("test")
722 | })
723 | }
724 |
725 | func TestHandleSentMessage(t *testing.T) {
726 | test := func(h func(*TestServer, chan bool), w func(*websocket.Conn)) {
727 | done := make(chan bool)
728 |
729 | ws := NewTestServer()
730 | server := httptest.NewServer(ws)
731 | defer server.Close()
732 |
733 | h(ws, done)
734 |
735 | conn := MustNewDialer(server.URL)
736 | defer conn.Close()
737 |
738 | w(conn)
739 |
740 | <-done
741 | }
742 |
743 | test(func(ws *TestServer, done chan bool) {
744 | ws.m.HandleMessage(func(s *Session, msg []byte) {
745 | s.Write(msg)
746 | })
747 |
748 | ws.m.HandleSentMessage(func(s *Session, msg []byte) {
749 | assert.Equal(t, TestMsg, msg)
750 | close(done)
751 | })
752 | }, func(conn *websocket.Conn) {
753 | conn.WriteMessage(websocket.TextMessage, TestMsg)
754 | })
755 |
756 | test(func(ws *TestServer, done chan bool) {
757 | ws.m.HandleMessageBinary(func(s *Session, msg []byte) {
758 | s.WriteBinary(msg)
759 | })
760 |
761 | ws.m.HandleSentMessageBinary(func(s *Session, msg []byte) {
762 | assert.Equal(t, TestMsg, msg)
763 | close(done)
764 | })
765 | }, func(conn *websocket.Conn) {
766 | conn.WriteMessage(websocket.BinaryMessage, TestMsg)
767 | })
768 | }
769 |
770 | func TestConcurrentMessageHandling(t *testing.T) {
771 | testTimeout := func(cmh bool, msgType int) bool {
772 | base := time.Millisecond * 100
773 | done := make(chan struct{})
774 |
775 | handler := func(s *Session, msg []byte) {
776 | if len(msg) == 0 {
777 | done <- struct{}{}
778 | return
779 | }
780 |
781 | time.Sleep(base * 2)
782 | }
783 |
784 | ws := NewTestServerHandler(func(session *Session, msg []byte) {})
785 | if msgType == websocket.TextMessage {
786 | ws.m.HandleMessage(handler)
787 | } else {
788 | ws.m.HandleMessageBinary(handler)
789 | }
790 |
791 | ws.m.Config.ConcurrentMessageHandling = cmh
792 | ws.m.Config.PongWait = base
793 |
794 | var errorSet atomic.Bool
795 | ws.m.HandleError(func(s *Session, err error) {
796 | errorSet.Store(true)
797 | done <- struct{}{}
798 | })
799 |
800 | server := httptest.NewServer(ws)
801 | defer server.Close()
802 |
803 | conn := MustNewDialer(server.URL)
804 | defer conn.Close()
805 |
806 | conn.WriteMessage(msgType, TestMsg)
807 | conn.WriteMessage(msgType, TestMsg)
808 |
809 | time.Sleep(base / 4)
810 |
811 | conn.WriteMessage(msgType, nil)
812 |
813 | <-done
814 |
815 | return errorSet.Load()
816 | }
817 |
818 | t.Run("text should error", func(t *testing.T) {
819 | errorSet := testTimeout(false, websocket.TextMessage)
820 |
821 | if !errorSet {
822 | t.FailNow()
823 | }
824 | })
825 |
826 | t.Run("text should not error", func(t *testing.T) {
827 | errorSet := testTimeout(true, websocket.TextMessage)
828 |
829 | if errorSet {
830 | t.FailNow()
831 | }
832 | })
833 |
834 | t.Run("binary should error", func(t *testing.T) {
835 | errorSet := testTimeout(false, websocket.BinaryMessage)
836 |
837 | if !errorSet {
838 | t.FailNow()
839 | }
840 | })
841 |
842 | t.Run("binary should not error", func(t *testing.T) {
843 | errorSet := testTimeout(true, websocket.BinaryMessage)
844 |
845 | if errorSet {
846 | t.FailNow()
847 | }
848 | })
849 | }
850 |
--------------------------------------------------------------------------------
/session.go:
--------------------------------------------------------------------------------
1 | package melody
2 |
3 | import (
4 | "net"
5 | "net/http"
6 | "sync"
7 | "time"
8 |
9 | "github.com/gorilla/websocket"
10 | )
11 |
12 | // Session wrapper around websocket connections.
13 | type Session struct {
14 | Request *http.Request
15 | Keys map[string]any
16 | conn *websocket.Conn
17 | output chan envelope
18 | outputDone chan struct{}
19 | melody *Melody
20 | open bool
21 | rwmutex *sync.RWMutex
22 | }
23 |
24 | func (s *Session) writeMessage(message envelope) {
25 | if s.closed() {
26 | s.melody.errorHandler(s, ErrWriteClosed)
27 | return
28 | }
29 |
30 | select {
31 | case s.output <- message:
32 | default:
33 | s.melody.errorHandler(s, ErrMessageBufferFull)
34 | }
35 | }
36 |
37 | func (s *Session) writeRaw(message envelope) error {
38 | if s.closed() {
39 | return ErrWriteClosed
40 | }
41 |
42 | s.conn.SetWriteDeadline(time.Now().Add(s.melody.Config.WriteWait))
43 | err := s.conn.WriteMessage(message.t, message.msg)
44 |
45 | if err != nil {
46 | return err
47 | }
48 |
49 | return nil
50 | }
51 |
52 | func (s *Session) closed() bool {
53 | s.rwmutex.RLock()
54 | defer s.rwmutex.RUnlock()
55 |
56 | return !s.open
57 | }
58 |
59 | func (s *Session) close() {
60 | s.rwmutex.Lock()
61 | open := s.open
62 | s.open = false
63 | s.rwmutex.Unlock()
64 | if open {
65 | s.conn.Close()
66 | close(s.outputDone)
67 | }
68 | }
69 |
70 | func (s *Session) ping() {
71 | s.writeRaw(envelope{t: websocket.PingMessage, msg: []byte{}})
72 | }
73 |
74 | func (s *Session) writePump() {
75 | ticker := time.NewTicker(s.melody.Config.PingPeriod)
76 | defer ticker.Stop()
77 |
78 | loop:
79 | for {
80 | select {
81 | case msg := <-s.output:
82 | err := s.writeRaw(msg)
83 |
84 | if err != nil {
85 | s.melody.errorHandler(s, err)
86 | break loop
87 | }
88 |
89 | if msg.t == websocket.CloseMessage {
90 | break loop
91 | }
92 |
93 | if msg.t == websocket.TextMessage {
94 | s.melody.messageSentHandler(s, msg.msg)
95 | }
96 |
97 | if msg.t == websocket.BinaryMessage {
98 | s.melody.messageSentHandlerBinary(s, msg.msg)
99 | }
100 | case <-ticker.C:
101 | s.ping()
102 | case _, ok := <-s.outputDone:
103 | if !ok {
104 | break loop
105 | }
106 | }
107 | }
108 |
109 | s.close()
110 | }
111 |
112 | func (s *Session) readPump() {
113 | s.conn.SetReadLimit(s.melody.Config.MaxMessageSize)
114 | s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
115 |
116 | s.conn.SetPongHandler(func(string) error {
117 | s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
118 | s.melody.pongHandler(s)
119 | return nil
120 | })
121 |
122 | if s.melody.closeHandler != nil {
123 | s.conn.SetCloseHandler(func(code int, text string) error {
124 | return s.melody.closeHandler(s, code, text)
125 | })
126 | }
127 |
128 | for {
129 | t, message, err := s.conn.ReadMessage()
130 |
131 | if err != nil {
132 | s.melody.errorHandler(s, err)
133 | break
134 | }
135 |
136 | if s.melody.Config.ConcurrentMessageHandling {
137 | go s.handleMessage(t, message)
138 | } else {
139 | s.handleMessage(t, message)
140 | }
141 | }
142 | }
143 |
144 | func (s *Session) handleMessage(t int, message []byte) {
145 | switch t {
146 | case websocket.TextMessage:
147 | s.melody.messageHandler(s, message)
148 | case websocket.BinaryMessage:
149 | s.melody.messageHandlerBinary(s, message)
150 | }
151 | }
152 |
153 | // Write writes message to session.
154 | func (s *Session) Write(msg []byte) error {
155 | if s.closed() {
156 | return ErrSessionClosed
157 | }
158 |
159 | s.writeMessage(envelope{t: websocket.TextMessage, msg: msg})
160 |
161 | return nil
162 | }
163 |
164 | // WriteBinary writes a binary message to session.
165 | func (s *Session) WriteBinary(msg []byte) error {
166 | if s.closed() {
167 | return ErrSessionClosed
168 | }
169 |
170 | s.writeMessage(envelope{t: websocket.BinaryMessage, msg: msg})
171 |
172 | return nil
173 | }
174 |
175 | // Close closes session.
176 | func (s *Session) Close() error {
177 | if s.closed() {
178 | return ErrSessionClosed
179 | }
180 |
181 | s.writeMessage(envelope{t: websocket.CloseMessage, msg: []byte{}})
182 |
183 | return nil
184 | }
185 |
186 | // CloseWithMsg closes the session with the provided payload.
187 | // Use the FormatCloseMessage function to format a proper close message payload.
188 | func (s *Session) CloseWithMsg(msg []byte) error {
189 | if s.closed() {
190 | return ErrSessionClosed
191 | }
192 |
193 | s.writeMessage(envelope{t: websocket.CloseMessage, msg: msg})
194 |
195 | return nil
196 | }
197 |
198 | // Set is used to store a new key/value pair exclusively for this session.
199 | // It also lazy initializes s.Keys if it was not used previously.
200 | func (s *Session) Set(key string, value any) {
201 | s.rwmutex.Lock()
202 | defer s.rwmutex.Unlock()
203 |
204 | if s.Keys == nil {
205 | s.Keys = make(map[string]any)
206 | }
207 |
208 | s.Keys[key] = value
209 | }
210 |
211 | // Get returns the value for the given key, ie: (value, true).
212 | // If the value does not exists it returns (nil, false)
213 | func (s *Session) Get(key string) (value any, exists bool) {
214 | s.rwmutex.RLock()
215 | defer s.rwmutex.RUnlock()
216 |
217 | if s.Keys != nil {
218 | value, exists = s.Keys[key]
219 | }
220 |
221 | return
222 | }
223 |
224 | // MustGet returns the value for the given key if it exists, otherwise it panics.
225 | func (s *Session) MustGet(key string) any {
226 | if value, exists := s.Get(key); exists {
227 | return value
228 | }
229 |
230 | panic("Key \"" + key + "\" does not exist")
231 | }
232 |
233 | // UnSet will delete the key and has no return value
234 | func (s *Session) UnSet(key string) {
235 | s.rwmutex.Lock()
236 | defer s.rwmutex.Unlock()
237 | if s.Keys != nil {
238 | delete(s.Keys, key)
239 | }
240 | }
241 |
242 | // IsClosed returns the status of the connection.
243 | func (s *Session) IsClosed() bool {
244 | return s.closed()
245 | }
246 |
247 | // LocalAddr returns the local addr of the connection.
248 | func (s *Session) LocalAddr() net.Addr {
249 | return s.conn.LocalAddr()
250 | }
251 |
252 | // RemoteAddr returns the remote addr of the connection.
253 | func (s *Session) RemoteAddr() net.Addr {
254 | return s.conn.RemoteAddr()
255 | }
256 |
257 | // WebsocketConnection returns the underlying websocket connection.
258 | // This can be used to e.g. set/read additional websocket options or to write sychronous messages.
259 | func (s *Session) WebsocketConnection() *websocket.Conn {
260 | return s.conn
261 | }
262 |
--------------------------------------------------------------------------------