go搭建websocket服务

之前整理go包管理工具的时候用的是自己做websocket服务的半成品代码,今天把做websocket服务的代码完善下

所用包

go官方是没有websocket的包的,这里笔者用的是第三方开发的一个包 gorilla/websocket,代码是笔者在官方的demo的基础上更改的

实现功能

笔者目前次项目简单的实现了前端绑定后,给服务端发消息,服务的把收到的消息直接再发送给前端, 这个功能是官方demo里就有的,笔者主要实现的是,添加了个接口后端接受GET请求后, 主动发消息给所有绑定的前端服务
接口:

send 接受msg参数把msg发给所有绑定的socket客户端
echo 默认绑定事件

home 页面入口

代码

package main

import (
	"flag"
	"html/template"
	"log"
	"net/http"
	"github.com/gorilla/websocket"
)

/*
type room struct{
    roomId int64
    users sync.Map
}*/

var addr = flag.String("addr", "localhost:8080", "http service address")

var upgrader = websocket.Upgrader{} // use default options

//var c *websocket.Conn
//定义socket 集合 所有socket共用一个集合 扩展下  定义组  每个组 socket消息互通
var room map[int]*websocket.Conn
var i int = 0

func send(w http.ResponseWriter, r *http.Request) {
    query := r.URL.Query()
    msg := query.Get("msg")
    for clt := range room{
        err := room[clt].WriteMessage(websocket.TextMessage, []byte(msg))
        if err != nil {
            delete(room, clt)
            log.Println("write:", err)
            log.Println("delete:", clt)
        }
    }
    w.Write([]byte("send"))
}
func echo(w http.ResponseWriter, r *http.Request) {
	//var err error
    c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
    room[i] = c
    i++
	defer c.Close()
	for {
		mt, message, err := c.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			break
		}
		log.Printf("recv: %s", message)
		err = c.WriteMessage(mt, message)
		if err != nil {
			log.Println("write:", err)
			break
		}
	}
}

func home(w http.ResponseWriter, r *http.Request) {
	homeTemplate.Execute(w, "ws://"+r.Host+"/echo")
}

func main() {
    room = make(map[int]*websocket.Conn)
	flag.Parse()
	log.SetFlags(0)
	http.HandleFunc("/echo", echo)
	http.HandleFunc("/", home)
    http.HandleFunc("/send", send)
	log.Fatal(http.ListenAndServe(*addr, nil))
}

var homeTemplate = template.Must(template.New("").Parse(`
<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<script>  
window.addEventListener("load", function(evt) {
    var output = document.getElementById("output");
    var input = document.getElementById("input");
    var ws;
    var print = function(message) {
        var d = document.createElement("div");
        d.textContent = message;
        output.appendChild(d);
        output.scroll(0, output.scrollHeight);
    };
    document.getElementById("open").onclick = function(evt) {
        if (ws) {
            return false;
        }
        ws = new WebSocket("{{.}}");
        ws.onopen = function(evt) {
            print("OPEN");
        }
        ws.onclose = function(evt) {
            print("CLOSE");
            ws = null;
        }
        ws.onmessage = function(evt) {
            print("RESPONSE: " + evt.data);
        }
        ws.onerror = function(evt) {
            print("ERROR: " + evt.data);
        }
        return false;
    };
    document.getElementById("send").onclick = function(evt) {
        if (!ws) {
            return false;
        }
        print("SEND: " + input.value);
        ws.send(input.value);
        return false;
    };
    document.getElementById("close").onclick = function(evt) {
        if (!ws) {
            return false;
        }
        ws.close();
        return false;
    };
});
</script>
</head>
<body>
<table>
<tr><td valign="top" width="50%">
<p>Click "Open" to create a connection to the server, 
"Send" to send a message to the server and "Close" to close the connection. 
You can change the message and send multiple times.
<p>
<form>
<button id="open">Open</button>
<button id="close">Close</button>
<p><input id="input" type="text" value="Hello world!">
<button id="send">Send</button>
</form>
</td><td valign="top" width="50%">
<div id="output" style="max-height: 70vh;overflow-y: scroll;"></div>
</td></tr></table>
</body>
</html>
`))