gorilla-wsserver_20201129160513.go 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183
  1. package main
  2. import (
  3. "encoding/json"
  4. "flag"
  5. "fmt"
  6. "log"
  7. "net/http"
  8. "github.com/google/uuid"
  9. "github.com/gorilla/websocket"
  10. )
  11. var wsaddr = flag.String("wsaddr", "localhost:12345", "websocket service address")
  12. type WSMsg struct {
  13. Mtype string `json:"mtype"`
  14. ID int `json:"id,omitempty"`
  15. Payload string `json:"payload,string"`
  16. }
  17. type Handler struct {
  18. MType string
  19. Callback func(WSMsg, *WSPeer) error
  20. }
  21. type WSServer struct {
  22. upgrader *websocket.Upgrader
  23. handler map[string]func(WSMsg, *WSPeer) error
  24. peers map[uuid.UUID]*WSPeer
  25. }
  26. func NewWSServer() *WSServer {
  27. upgrader := websocket.Upgrader{} // use default options
  28. upgrader.CheckOrigin = func(r *http.Request) bool { return true } // cors-origin '*'
  29. wss := WSServer{upgrader: &upgrader, peers: map[uuid.UUID]*WSPeer{}}
  30. return &wss
  31. }
  32. func (wss *WSServer) Write(name string, sendMsg WSMsg) {
  33. peers := wss.GetPeers(name)
  34. if len(peers) == 0 {
  35. log.Println("WSServer:Write - no sendChans for", name)
  36. return
  37. }
  38. for _, peer := range peers {
  39. peer.Answer(sendMsg)
  40. }
  41. }
  42. func (wss *WSServer) GetPeers(searchname string) []*WSPeer {
  43. cons := []*WSPeer{}
  44. for key, con := range wss.peers {
  45. username := key.String()
  46. if username == searchname {
  47. cons = append(cons, con)
  48. }
  49. }
  50. return cons
  51. }
  52. func (wss *WSServer) CleanPeer(name uuid.UUID) {
  53. delete(wss.peers, name)
  54. }
  55. func (wss *WSServer) RegisterHandler(handlers []Handler) {
  56. if wss.handler == nil {
  57. wss.handler = map[string]func(WSMsg, *WSPeer) error{}
  58. }
  59. for _, handler := range handlers {
  60. log.Println("Registrering handler for MType: ", handler.MType)
  61. wss.handler[handler.MType] = handler.Callback
  62. }
  63. }
  64. type WSPeer struct {
  65. uuid uuid.UUID
  66. conn *websocket.Conn
  67. answers *chan WSMsg
  68. }
  69. func (peer *WSPeer) WorkAnswers() {
  70. for {
  71. msg := WSMsg{}
  72. msg = <-*peer.answers
  73. fmt.Println("sending on: ", peer.uuid)
  74. peer._answer(msg)
  75. }
  76. }
  77. func (peer *WSPeer) Answer(msg WSMsg) {
  78. if peer.answers == nil {
  79. fmt.Println("empty chan for :", peer.uuid)
  80. }
  81. *peer.answers <- msg
  82. }
  83. func (peer *WSPeer) _answer(msg WSMsg) error {
  84. log.Println("sending message: ", msg)
  85. jsonMsg, err := json.Marshal(msg)
  86. if err != nil {
  87. log.Printf("write: err marshalling \"%s\"\n%s", msg, err.Error())
  88. return err
  89. }
  90. err = peer.conn.WriteMessage(websocket.BinaryMessage, jsonMsg)
  91. if err != nil {
  92. log.Println("write:", err)
  93. return err
  94. }
  95. return nil
  96. }
  97. func (wss WSServer) WebsocketHandler(w http.ResponseWriter, r *http.Request) {
  98. log.Println("request on /ws from", r.RemoteAddr)
  99. c, err := wss.upgrader.Upgrade(w, r, nil)
  100. if err != nil {
  101. log.Print("upgrade:", err)
  102. return
  103. }
  104. defer c.Close()
  105. peer := WSPeer{uuid: uuid.New(), conn: c}
  106. bla := make(chan WSMsg, 10)
  107. peer.answers = &bla
  108. go peer.WorkAnswers()
  109. wss.peers[peer.uuid] = &peer
  110. for {
  111. _, message, err := c.ReadMessage()
  112. if err != nil {
  113. log.Println("read:", err)
  114. break
  115. }
  116. var rawmsg map[string]interface{}
  117. err = json.Unmarshal(message, &rawmsg)
  118. msg := WSMsg{}
  119. for k, v := range rawmsg {
  120. switch k {
  121. case "id":
  122. id := v.(float64)
  123. msg.ID = int(id)
  124. case "mtype":
  125. msg.Mtype = v.(string)
  126. case "payload":
  127. payload, ok := v.(string)
  128. if !ok {
  129. payloadmap := v.(map[string]interface{})
  130. var payloadbyte []byte
  131. payloadbyte, err = json.Marshal(payloadmap)
  132. payload = string(payloadbyte)
  133. }
  134. msg.Payload = payload
  135. }
  136. }
  137. if err != nil {
  138. log.Println("error unmarshalling message:", err)
  139. log.Printf("recv: %s", message)
  140. return
  141. }
  142. handler, ok := wss.handler[msg.Mtype]
  143. if !ok {
  144. log.Println("unknown MType", msg.Mtype)
  145. keys := ""
  146. for k := range wss.handler {
  147. keys = fmt.Sprintf("%s\n%s", keys, k)
  148. }
  149. log.Println("available MTypes:\n", keys, "----")
  150. remsg := WSMsg{Mtype: "UNKNOWN", Payload: ""}
  151. peer.Answer(remsg)
  152. return
  153. }
  154. err = handler(msg, &peer)
  155. if err != nil {
  156. log.Println("handler failed", err.Error())
  157. //writeChan <- msg
  158. }
  159. log.Println("message handled")
  160. }
  161. }