├── .gitignore ├── docs ├── http-vs-ws.png ├── server-flowchart.png ├── http-vs-ws.xml └── server-flowchart.xml ├── Makefile ├── go.mod ├── cmd └── main │ └── main.go ├── go.sum ├── internal └── websocket │ ├── model.go │ ├── handler.go │ └── server.go └── README.md /.gitignore: -------------------------------------------------------------------------------- 1 | .vscode 2 | main -------------------------------------------------------------------------------- /docs/http-vs-ws.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/madeindra/golang-websocket/HEAD/docs/http-vs-ws.png -------------------------------------------------------------------------------- /docs/server-flowchart.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/madeindra/golang-websocket/HEAD/docs/server-flowchart.png -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: 2 | 3 | setup: 4 | go mod tidy 5 | 6 | build: .PHONY 7 | go build -o main ./cmd/main 8 | 9 | run: .PHONY build 10 | ./main -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/madeindra/golang-websocket 2 | 3 | go 1.15 4 | 5 | require ( 6 | github.com/google/uuid v1.2.0 7 | github.com/gorilla/websocket v1.4.2 8 | ) 9 | -------------------------------------------------------------------------------- /cmd/main/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "net/http" 5 | 6 | "github.com/madeindra/golang-websocket/internal/websocket" 7 | ) 8 | 9 | func main() { 10 | http.HandleFunc("/socket", websocket.HandleWS) 11 | 12 | if err := http.ListenAndServe(":8080", nil); err != nil { 13 | panic(err) 14 | } 15 | } 16 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/google/uuid v1.2.0 h1:qJYtXnJRWmpe7m/3XlyhrsLrEURqHRM2kxzoxXqyUDs= 2 | github.com/google/uuid v1.2.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= 3 | github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc= 4 | github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= 5 | -------------------------------------------------------------------------------- /internal/websocket/model.go: -------------------------------------------------------------------------------- 1 | package websocket 2 | 3 | import "github.com/gorilla/websocket" 4 | 5 | // Subscription is a type for each string of topic and the clients that subscribe to it 6 | type Subscription map[string]Client 7 | 8 | // Client is a type that describe the clients' ID and their connection 9 | type Client map[string]*websocket.Conn 10 | 11 | // Message is a struct for message to be sent by the client 12 | type Message struct { 13 | Action string `json:"action"` 14 | Topic string `json:"topic"` 15 | Message string `json:"message"` 16 | } 17 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Websocket Server in Go 2 | 3 | ## Running the server 4 | 1. Clone this repository 5 | 6 | 2. Mount the repository & run this command to install dependencies 7 | ``` 8 | make setup 9 | ``` 10 | 11 | 3. Run the websocket server 12 | ``` 13 | make run 14 | ``` 15 | 16 | 4. Websocket server will be running on `localhost:8080` 17 | 18 | ## Using this server with client 19 | 1. After running the server, open your Websocket client. If you don't have any, try `Websocket King` extension for chrome. 20 | 21 | 2. Connect to `ws://localhost:8080/socket`, you will be greeted by the server. 22 | ``` 23 | Server: Welcome! Your ID is f0ab664a-5af3-4f8d-8afe-eb93085267e4 24 | ``` 25 | 26 | 3. To subscribe to a topic, send this payload (*topic can be anything*) 27 | ``` 28 | { 29 | "action": "subscribe", 30 | "topic": "world" 31 | } 32 | ``` 33 | 34 | 4. To send a message to the topic's subscribers, send payload in this format 35 | ``` 36 | { 37 | "action": "publish", 38 | "topic": "world", 39 | "message": "Hello world!" 40 | } 41 | ``` 42 | 43 | 5. To unsubscribe from the topic, send this payload (*topic can be anything*) 44 | ``` 45 | { 46 | "action": "unsubscribe", 47 | "topic": "world" 48 | } 49 | ``` 50 | 51 | ## HTTP vs WebSocket 52 | 53 | You might be asking "why should I use Websocket instead of REST API"? 54 | 55 | REST API uses HTTP which can only send response once per request. 56 | 57 | Meanwhile, WebSocket can be used for persistent bidirectional communication without the need of reestablishing connection everytime. 58 | 59 | This can be useful in some scenario like chatting or pub-sub. 60 | 61 | Here is the diagram to visualize the difference between HTTP and WebSocket. 62 | 63 | ![HTTP-vs-WebSocket](./docs/http-vs-ws.png) 64 | 65 | ## Flowchart 66 | This flowchart describes how this server works. 67 | ![Server-Flowchart](./docs/server-flowchart.png) 68 | 69 | ## Project Structure 70 | ``` 71 | cmd 72 | └── main 73 | └── main.go 74 | internal 75 | └── websocket 76 | └── handler.go 77 | └── model.go 78 | └── server.go 79 | ``` 80 | ### Main files 81 | **main.go**: the main file to be executed. 82 | 83 | ### Handler 84 | 85 | **handler.go**: handles open/close connection & pass the message to the server. 86 | 87 | ### Model 88 | 89 | **model.go**: stores the models used by the server. 90 | 91 | ### Server 92 | **server.go**: runs specific action according to the client message, also containes functions that needed by the server to work properly as a websocket server. 93 | 94 | ## Further Work 95 | 96 | This repository is far from ideal. It's just a proof-of-concept. 97 | 98 | While this repository is close to a pub-sub, it can still be used for a chat server. 99 | 100 | For example, we can add a function on socket connected so that client will be automatically subscribes to their own ID as a topic. Other clients then will use those user's ID as a topic to publish a message. 101 | 102 | I have tried building such solution combined with Authorization to prevent other user from subscribing to other's ID and it does work. 103 | 104 | ## Credit 105 | This repository is inspired by [Golang-PubSub by @tabvn](https://github.com/tabvn/golang-pubsub-youtube) -------------------------------------------------------------------------------- /internal/websocket/handler.go: -------------------------------------------------------------------------------- 1 | package websocket 2 | 3 | import ( 4 | "fmt" 5 | "net/http" 6 | "time" 7 | 8 | "github.com/google/uuid" 9 | "github.com/gorilla/websocket" 10 | ) 11 | 12 | const ( 13 | // time to read the next client's pong message 14 | pongWait = 60 * time.Second 15 | // time period to send pings to client 16 | pingPeriod = (pongWait * 9) / 10 17 | // time allowed to write a message to client 18 | writeWait = 10 * time.Second 19 | // max message size allowed 20 | maxMessageSize = 512 21 | // I/O read buffer size 22 | readBufferSize = 1024 23 | // I/O write buffer size 24 | writeBufferSize = 1024 25 | ) 26 | 27 | // Initialize server with empty subscription 28 | var server = &Server{Subscriptions: make(Subscription)} 29 | 30 | // http to websocket upgrader 31 | var upgrader = websocket.Upgrader{ 32 | ReadBufferSize: readBufferSize, 33 | WriteBufferSize: writeBufferSize, 34 | CheckOrigin: func(r *http.Request) bool { 35 | // allow all origin 36 | return true 37 | }, 38 | } 39 | 40 | func HandleWS(w http.ResponseWriter, r *http.Request) { 41 | // upgrades connection to websocket 42 | conn, err := upgrader.Upgrade(w, r, nil) 43 | if err != nil { 44 | w.WriteHeader(http.StatusInternalServerError) 45 | w.Write([]byte("failed upgrading connection")) 46 | return 47 | } 48 | defer conn.Close() 49 | 50 | // create new client id 51 | clientID := uuid.New().String() 52 | 53 | // greet the new client 54 | server.Send(conn, fmt.Sprintf("Server: Welcome! Your ID is %s", clientID)) 55 | 56 | // create channel to signal client health 57 | done := make(chan struct{}) 58 | 59 | go writePump(conn, clientID, done) 60 | readPump(conn, clientID, done) 61 | } 62 | 63 | // readPump process incoming messages and set the settings 64 | func readPump(conn *websocket.Conn, clientID string, done chan<- struct{}) { 65 | // set limit, deadline to read & pong handler 66 | conn.SetReadLimit(maxMessageSize) 67 | conn.SetReadDeadline(time.Now().Add(pongWait)) 68 | conn.SetPongHandler(func(string) error { 69 | conn.SetReadDeadline(time.Now().Add(pongWait)) 70 | return nil 71 | }) 72 | 73 | // message handling 74 | for { 75 | // read incoming message 76 | _, msg, err := conn.ReadMessage() 77 | // if error occured 78 | if err != nil { 79 | // remove from the client 80 | server.RemoveClient(clientID) 81 | // set health status to unhealthy by closing channel 82 | close(done) 83 | // stop process 84 | break 85 | } 86 | 87 | // if no error, process incoming message 88 | server.ProcessMessage(conn, clientID, msg) 89 | } 90 | } 91 | 92 | // writePump sends ping to the client 93 | func writePump(conn *websocket.Conn, clientID string, done <-chan struct{}) { 94 | // create ping ticker 95 | ticker := time.NewTicker(pingPeriod) 96 | defer ticker.Stop() 97 | 98 | for { 99 | select { 100 | case <-ticker.C: 101 | // send ping message 102 | err := conn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait)) 103 | if err != nil { 104 | // if error sending ping, remove this client from the server 105 | server.RemoveClient(clientID) 106 | // stop sending ping 107 | return 108 | } 109 | case <-done: 110 | // if process is done, stop sending ping 111 | return 112 | } 113 | } 114 | } 115 | -------------------------------------------------------------------------------- /internal/websocket/server.go: -------------------------------------------------------------------------------- 1 | package websocket 2 | 3 | import ( 4 | "encoding/json" 5 | "strings" 6 | "sync" 7 | 8 | "github.com/gorilla/websocket" 9 | ) 10 | 11 | // constants for action type 12 | const ( 13 | publish = "publish" 14 | subscribe = "subscribe" 15 | unsubscribe = "unsubscribe" 16 | ) 17 | 18 | // constants for server message 19 | const ( 20 | errInvalidMessage = "Server: Invalid msg" 21 | errActionUnrecognizable = "Server: Action unrecognized" 22 | ) 23 | 24 | // Server is the struct to handle the Server functions & manage the Subscriptions 25 | type Server struct { 26 | Subscriptions Subscription 27 | } 28 | 29 | // Send simply sends message to the websocket client 30 | func (s *Server) Send(conn *websocket.Conn, message string) { 31 | // send simple message 32 | conn.WriteMessage(websocket.TextMessage, []byte(message)) 33 | } 34 | 35 | // SendWithWait sends message to the websocket client using wait group, allowing usage with goroutines 36 | func (s *Server) SendWithWait(conn *websocket.Conn, message string, wg *sync.WaitGroup) { 37 | // send simple message 38 | conn.WriteMessage(websocket.TextMessage, []byte(message)) 39 | 40 | // set the task as done 41 | wg.Done() 42 | } 43 | 44 | // RemoveClient removes the clients from the server subscription map 45 | func (s *Server) RemoveClient(clientID string) { 46 | // loop all topics 47 | for _, client := range s.Subscriptions { 48 | // delete the client from all the topic's client map 49 | delete(client, clientID) 50 | } 51 | } 52 | 53 | // ProcessMessage handle message according to the action type 54 | func (s *Server) ProcessMessage(conn *websocket.Conn, clientID string, msg []byte) *Server { 55 | // parse message 56 | m := Message{} 57 | if err := json.Unmarshal(msg, &m); err != nil { 58 | s.Send(conn, errInvalidMessage) 59 | } 60 | 61 | // convert all action to lowercase and remove whitespace 62 | action := strings.TrimSpace(strings.ToLower(m.Action)) 63 | 64 | switch action { 65 | case publish: 66 | s.Publish(m.Topic, []byte(m.Message)) 67 | 68 | case subscribe: 69 | s.Subscribe(conn, clientID, m.Topic) 70 | 71 | case unsubscribe: 72 | s.Unsubscribe(clientID, m.Topic) 73 | 74 | default: 75 | s.Send(conn, errActionUnrecognizable) 76 | } 77 | 78 | return s 79 | } 80 | 81 | // Publish sends a message to all subscribing clients of a topic 82 | func (s *Server) Publish(topic string, message []byte) { 83 | // if topic does not exist, stop the process 84 | if _, exist := s.Subscriptions[topic]; !exist { 85 | return 86 | } 87 | 88 | // if topic exist 89 | client := s.Subscriptions[topic] 90 | 91 | // send the message to the clients 92 | var wg sync.WaitGroup 93 | for _, conn := range client { 94 | // add 1 job to wait group 95 | wg.Add(1) 96 | 97 | // send with goroutines 98 | go s.SendWithWait(conn, string(message), &wg) 99 | } 100 | 101 | // wait until all goroutines jobs done 102 | wg.Wait() 103 | } 104 | 105 | // Subscribe adds a client to a topic's client map 106 | func (s *Server) Subscribe(conn *websocket.Conn, clientID string, topic string) { 107 | // if topic exist, check the client map 108 | if _, exist := s.Subscriptions[topic]; exist { 109 | client := s.Subscriptions[topic] 110 | 111 | // if client already subbed, stop the process 112 | if _, subbed := client[clientID]; subbed { 113 | return 114 | } 115 | 116 | // if not subbed, add to client map 117 | client[clientID] = conn 118 | return 119 | } 120 | 121 | // if topic does not exist, create a new topic 122 | newClient := make(Client) 123 | s.Subscriptions[topic] = newClient 124 | 125 | // add the client to the topic 126 | s.Subscriptions[topic][clientID] = conn 127 | } 128 | 129 | // Unsubscribe removes a clients from a topic's client map 130 | func (s *Server) Unsubscribe(clientID string, topic string) { 131 | // if topic exist, check the client map 132 | if _, exist := s.Subscriptions[topic]; exist { 133 | client := s.Subscriptions[topic] 134 | 135 | // remove the client from the topic's client map 136 | delete(client, clientID) 137 | } 138 | } 139 | -------------------------------------------------------------------------------- /docs/http-vs-ws.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | -------------------------------------------------------------------------------- /docs/server-flowchart.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | 103 | 104 | 105 | 106 | 107 | 108 | 109 | 110 | 111 | 112 | 113 | 114 | 115 | 116 | 117 | 118 | 119 | 120 | 121 | 122 | 123 | 124 | 125 | 126 | 127 | 128 | 129 | 130 | 131 | 132 | 133 | 134 | 135 | 136 | 137 | 138 | 139 | 140 | 141 | 142 | 143 | 144 | 145 | 146 | 147 | 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | 160 | 161 | 162 | 163 | 164 | 165 | 166 | 167 | 168 | 169 | 170 | 171 | 172 | 173 | 174 | 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | 183 | 184 | 185 | 186 | 187 | 188 | 189 | 190 | 191 | 192 | 193 | 194 | 195 | 196 | 197 | 198 | 199 | 200 | 201 | 202 | 203 | 204 | 205 | 206 | 207 | 208 | 209 | 210 | 211 | 212 | 213 | 214 | 215 | 216 | 217 | 218 | 219 | 220 | 221 | 222 | 223 | 224 | 225 | 226 | 227 | 228 | 229 | 230 | 231 | 232 | 233 | 234 | 235 | 236 | 237 | 238 | 239 | 240 | 241 | 242 | 243 | 244 | 245 | --------------------------------------------------------------------------------