package main import ( "encoding/json" "flag" "fmt" "log" "net/http" "github.com/google/uuid" "github.com/gorilla/websocket" ) var wsaddr = flag.String("wsaddr", "localhost:12345", "websocket service address") type WSMsg struct { Mtype string `json:"mtype"` ID int `json:"id,omitempty"` Payload interface{} `json:"payload,string"` } type Handler struct { MType string Callback func(WSMsg, *WSPeer) error } type WSServer struct { upgrader *websocket.Upgrader handler map[string]func(WSMsg, *WSPeer) error peers map[uuid.UUID]*WSPeer removeChan *chan *WSPeer } func NewWSServer(removechan *chan *WSPeer) *WSServer { upgrader := websocket.Upgrader{} // use default options upgrader.CheckOrigin = func(r *http.Request) bool { return true } // cors-origin '*' wss := WSServer{upgrader: &upgrader, peers: map[uuid.UUID]*WSPeer{}, removeChan: removechan} return &wss } func (wss *WSServer) Write(name string, sendMsg WSMsg) { peers := wss.GetPeers(name) if len(peers) == 0 { log.Println("WSServer:Write - no sendChans for", name) return } for _, peer := range peers { peer.Answer(sendMsg) } } func (wss *WSServer) GetPeers(searchname string) []*WSPeer { cons := []*WSPeer{} for key, con := range wss.peers { username := key.String() if username == searchname { cons = append(cons, con) } } return cons } func (wss *WSServer) GetPeercount() int { return len(wss.peers) } func (wss *WSServer) CleanPeer(name uuid.UUID) { delete(wss.peers, name) } func (wss *WSServer) RegisterHandler(handlers []Handler) { if wss.handler == nil { wss.handler = map[string]func(WSMsg, *WSPeer) error{} } for _, handler := range handlers { log.Println("Registrering handler for MType: ", handler.MType) wss.handler[handler.MType] = handler.Callback } } type WSPeer struct { uuid uuid.UUID conn *websocket.Conn answers *chan WSMsg } func (peer *WSPeer) WorkAnswers() { for { msg := WSMsg{} msg = <-*peer.answers fmt.Println("sending on: ", peer.uuid) peer._answer(msg) } } func (peer *WSPeer) Answer(msg WSMsg) { if peer.answers == nil { fmt.Println("empty chan for :", peer.uuid) } *peer.answers <- msg } func (peer *WSPeer) _answer(msg WSMsg) error { //log.Println("sending message: ", msg) jsonMsg, err := json.Marshal(msg) if err != nil { log.Printf("write: err marshalling \"%s\"\n%s", msg, err.Error()) return err } err = peer.conn.WriteMessage(websocket.BinaryMessage, jsonMsg) if err != nil { log.Println("write:", err) return err } return nil } func (wss WSServer) removePeer(peer *WSPeer) { log.Println("removing WSPeer", peer.uuid) delete(wss.peers, peer.uuid) *wss.removeChan <- peer } func (wss WSServer) WebsocketHandler(w http.ResponseWriter, r *http.Request) { log.Println("request on /ws from", r.RemoteAddr) c, err := wss.upgrader.Upgrade(w, r, nil) if err != nil { log.Print("upgrade:", err) return } defer c.Close() peer := WSPeer{uuid: uuid.New(), conn: c} defer wss.removePeer(&peer) bla := make(chan WSMsg, 10) peer.answers = &bla go peer.WorkAnswers() wss.peers[peer.uuid] = &peer for { _, message, err := c.ReadMessage() if err != nil { log.Println("read:", err) break } var rawmsg map[string]interface{} err = json.Unmarshal(message, &rawmsg) msg := WSMsg{} for k, v := range rawmsg { switch k { case "id": id := v.(float64) msg.ID = int(id) case "mtype": msg.Mtype = v.(string) case "payload": msg.Payload = v /* payload, ok := v.(string) if !ok { payloadmap := v.(map[string]interface{}) var payloadbyte []byte payloadbyte, err = json.Marshal(payloadmap) payload = string(payloadbyte) } msg.Payload = payload */ } } if err != nil { log.Println("error unmarshalling message:", err) //log.Printf("recv: %s", message) return } handler, ok := wss.handler[msg.Mtype] if !ok { log.Println("unknown MType", msg.Mtype) keys := "" for k := range wss.handler { keys = fmt.Sprintf("%s\n%s", keys, k) } log.Println("available MTypes:\n", keys, "----") remsg := WSMsg{Mtype: "UNKNOWN", Payload: ""} peer.Answer(remsg) return } err = handler(msg, &peer) if err != nil { log.Println("handler failed", err.Error()) //writeChan <- msg } log.Println("message handled") } }