gr792_kkm
/
kkm792


			
				
					
						
						
							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 // disable CORS
	},
}

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: blue;
    }
    #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))
	}
}