nnet/cpn/server_ws.go

package cpn

import (
	"net/http"

	"git.wenlab.co/joe/nnet"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
)

type CallbackWsPath func(http.ResponseWriter, *http.Request)

type WsServer struct {
	*Hub
	addr     string
	wspath   string
	svr      *http.Server
	upgrader *websocket.Upgrader
	routes   map[string]CallbackWsPath
}

func NewWsServer(cf *nnet.HubConfig, cb nnet.ISessionCallback, p nnet.IProtocol,
	addr string, wspath string, routes map[string]CallbackWsPath) nnet.IHub {
	return &WsServer{
		Hub:    newHub(cf, cb, p),
		addr:   addr,
		wspath: wspath,
		upgrader: &websocket.Upgrader{
			ReadBufferSize:  cf.ReadBufSize,
			WriteBufferSize: cf.WriteBufSize,
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		},
		routes: routes,
	}
}

func (self *WsServer) Start() error {
	router := mux.NewRouter()
	for k, v := range self.routes {
		router.HandleFunc(k, v)
	}
	router.HandleFunc("/", self.do_homepage)
	router.HandleFunc(self.wspath, func(w http.ResponseWriter, r *http.Request) {
		self.do_new_session(w, r)
	})
	self.svr = &http.Server{
		Addr:    self.addr,
		Handler: router,
	}
	err := self.svr.ListenAndServe()
	return err
}

func (self *WsServer) do_homepage(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	}
	if r.Method != "GET" {
		http.Error(w, "Method now allowed", http.StatusMethodNotAllowed)
		return
	}
	w.Write([]byte("welcome to nnet v" + nnet.VERSION))
}

func (self *WsServer) do_new_session(w http.ResponseWriter, r *http.Request) {
	conn, err := self.upgrader.Upgrade(w, r, nil)
	if nil != err {
		return
	}
	self.wg.Add(1)
	go func() {
		ses := newSession(NewWsConn(conn), self)
		ses.Do()
		self.wg.Done()
	}()
}

func (self *WsServer) Stop() error {
	self.svr.Close()
	return self.Hub.Stop()
}