123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290 |
- package main
- import (
- "bytes"
- "fmt"
- native "log"
- "net/http"
- "os"
- "time"
- ws "github.com/gorilla/websocket"
- "github.com/kelseyhightower/envconfig"
- )
- var (
- version string
- date string
- logger = native.New(os.Stdout, "", 0)
- errors = native.New(os.Stderr, "", 0)
- )
- var upgrader = ws.Upgrader{
- ReadBufferSize: 1024,
- WriteBufferSize: 1024,
- CheckOrigin: func(r *http.Request) bool {
- return true
- },
- }
- type config struct {
- Port int64 `default:"8080"`
- Prefix string `default:"/"`
- }
- func main() {
- var c config
- if err := envconfig.Process("ws", &c); err != nil {
- errors.Fatal(err)
- }
- mngr := manager{
- clients: make(map[*client]bool),
- add: make(chan *client),
- remove: make(chan *client),
- broadcast: make(chan []byte),
- }
- go func() {
- for {
- select {
- case client := <-mngr.add:
- mngr.clients[client] = true
- case client := <-mngr.remove:
- if _, ok := mngr.clients[client]; ok {
- delete(mngr.clients, client)
- close(client.send)
- }
- case message := <-mngr.broadcast:
- for client := range mngr.clients {
- select {
- case client.send <- message:
- default:
- close(client.send)
- delete(mngr.clients, client)
- }
- }
- }
- }
- }()
- http.HandleFunc(fmt.Sprintf("%sws", c.Prefix), func(w http.ResponseWriter, r *http.Request) {
- conn, err := upgrader.Upgrade(w, r, nil)
- if err != nil {
- logger.Println(err)
- return
- }
- client := &client{
- conn: conn,
- send: make(chan []byte, 256),
- mngr: mngr,
- }
- mngr.add <- client
- go client.write()
- go client.read()
- })
- http.HandleFunc(c.Prefix, index(c.Prefix))
- if c.Prefix != "/" {
- http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
- w.WriteHeader(http.StatusForbidden)
- })
- }
- http.HandleFunc(fmt.Sprintf("%sversion", c.Prefix), func(w http.ResponseWriter, r *http.Request) {
- if len(version) > 0 && len(date) > 0 {
- fmt.Fprintf(w, "version: %s (built at %s)", version, date)
- } else {
- w.WriteHeader(http.StatusOK)
- }
- })
- http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
- w.WriteHeader(http.StatusOK)
- })
- logger.Printf("[service] listening on port %d", c.Port)
- if err := http.ListenAndServe(fmt.Sprintf(":%d", c.Port), nil); err != nil {
- errors.Fatal(err)
- }
- }
- type manager struct {
- clients map[*client]bool
- add chan *client
- remove chan *client
- broadcast chan []byte
- }
- type client struct {
- conn *ws.Conn
- send chan []byte
- mngr manager
- }
- const (
- maxMessageSize = 512
- readDeadline = 60 * time.Second
- writeDeadline = 10 * time.Second
- pingPeriod = (writeDeadline * 9) / 10
- )
- var (
- newline = []byte{'\n'}
- space = []byte{' '}
- )
- func (c *client) read() {
- defer func() {
- c.mngr.remove <- c
- c.conn.Close()
- }()
- c.conn.SetReadLimit(maxMessageSize)
- c.conn.SetReadDeadline(time.Now().Add(readDeadline))
- c.conn.SetPongHandler(func(string) error {
- c.conn.SetReadDeadline(time.Now().Add(readDeadline))
- return nil
- })
- for {
- _, message, err := c.conn.ReadMessage()
- if err != nil {
- if ws.IsUnexpectedCloseError(err, ws.CloseGoingAway, ws.CloseAbnormalClosure) {
- logger.Printf("error: %v", err)
- }
- break
- }
- message = bytes.TrimSpace(bytes.Replace(message, newline, space, -1))
- c.mngr.broadcast <- message
- }
- }
- func (c *client) write() {
- ticker := time.NewTicker(pingPeriod)
- defer func() {
- ticker.Stop()
- c.conn.Close()
- }()
- for {
- select {
- case message, ok := <-c.send:
- c.conn.SetWriteDeadline(time.Now().Add(writeDeadline))
- if !ok {
- c.conn.WriteMessage(ws.CloseMessage, []byte{})
- return
- }
- w, err := c.conn.NextWriter(ws.TextMessage)
- if err != nil {
- return
- }
- w.Write(message)
- n := len(c.send)
- for i := 0; i < n; i++ {
- w.Write(newline)
- w.Write(<-c.send)
- }
- if err := w.Close(); err != nil {
- return
- }
- case <-ticker.C:
- c.conn.SetWriteDeadline(time.Now().Add(writeDeadline))
- if err := c.conn.WriteMessage(ws.PingMessage, nil); err != nil {
- return
- }
- }
- }
- }
- var html = `<!DOCTYPE html>
- <html lang="en">
- <head>
- <meta charset="UTF-8">
- <title>Chat Example</title>
- <script type="text/javascript">
- window.onload = function () {
- var conn;
- var msg = document.getElementById("msg");
- var log = document.getElementById("log");
- function appendLog(item) {
- var doScroll = log.scrollTop > log.scrollHeight - log.clientHeight - 1;
- log.appendChild(item);
- if (doScroll) {
- log.scrollTop = log.scrollHeight - log.clientHeight;
- }
- }
- document.getElementById("form").onsubmit = function () {
- if (!conn) {
- return false;
- }
- if (!msg.value) {
- return false;
- }
- conn.send(msg.value);
- msg.value = "";
- return false;
- };
- if (window["WebSocket"]) {
- conn = new WebSocket("ws://" + document.location.host + "%sws");
- conn.onclose = function (evt) {
- var item = document.createElement("div");
- item.innerHTML = "<b>Connection closed.</b>";
- appendLog(item);
- };
- conn.onmessage = function (evt) {
- var messages = evt.data.split('\n');
- for (var i = 0; i < messages.length; i++) {
- var item = document.createElement("div");
- item.innerText = messages[i];
- appendLog(item);
- }
- };
- } else {
- var item = document.createElement("div");
- item.innerHTML = "<b>Your browser does not support WebSockets.</b>";
- appendLog(item);
- }
- };
- </script>
- <style type="text/css">
- html {
- overflow: hidden;
- }
- body {
- overflow: hidden;
- padding: 0;
- margin: 0;
- width: 100%%;
- height: 100%%;
- background: maroon;
- }
- #log {
- background: white;
- margin: 0;
- padding: 0.5em 0.5em 0.5em 0.5em;
- position: absolute;
- top: 0.5em;
- left: 0.5em;
- right: 0.5em;
- bottom: 3em;
- overflow: auto;
- }
- #form {
- padding: 0 0.5em 0 0.5em;
- margin: 0;
- position: absolute;
- bottom: 1em;
- left: 0px;
- width: 100%;
- overflow: hidden;
- }
- </style>
- <link rel="shortcut icon" href="data:image/x-icon;," type="image/x-icon">
- </head>
- <body>
- <div id="log"></div>
- <form id="form">
- <input type="submit" value="Send" />
- <input type="text" id="msg" size="64" />
- </form>
- </body>
- </html>`
- func index(prefix string) func(http.ResponseWriter, *http.Request) {
- return func(w http.ResponseWriter, r *http.Request) {
- fmt.Fprint(w, fmt.Sprintf(html, prefix))
- }
- }
|