frontend/wsserver/gorilla-wsserver.go

package wsserver

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"strings"

	gamelangpb "git.alfi.li/gamelang/protobuf/gamelang"
	um "git.alfi.li/gamelang/systems/usermanager"

	"github.com/gorilla/websocket"
)

var wsaddr = flag.String("wsaddr", "localhost:12345", "websocket service address")

type WSMsg struct {
	Mtype   string `json:"mtype"`
	Payload string `json:"payload"`
}

type Handler struct {
	MType    string
	Callback func(WSMsg, gamelangpb.User) (WSMsg, error)
}

type WSServer struct {
	usermanager *um.UserManager
	userclient  gamelangpb.UserServiceClient

	upgrader   *websocket.Upgrader
	handler    map[string]func(WSMsg, gamelangpb.User) (WSMsg, error)
	writeChans map[string]chan WSMsg
}

func NewWSServer(umBroker []string, userclient gamelangpb.UserServiceClient, usermanager *um.UserManager) *WSServer {
	upgrader := websocket.Upgrader{} // use default options
	wss := WSServer{upgrader: &upgrader, writeChans: map[string]chan WSMsg{}}
	if userclient != nil {
		log.Print("init userclient")
		wss.userclient = userclient
	} else if usermanager != nil {
		log.Print("using usermanager")
		wss.usermanager = usermanager
	} else {
		log.Print("init usermanager")
		usermanager := um.NewUserManager(umBroker)
		wss.usermanager = usermanager
	}
	return &wss
}

func (wss *WSServer) Write(name string, sendMsg WSMsg) {
	chans := wss.GetWriteChans(name)
	if len(chans) == 0 {
		log.Println("WSServer:Write - no sendChans for", name)
		return
	}
	for _, con := range chans {
		con <- sendMsg
	}
}

func (wss *WSServer) GetWriteChans(searchname string) []chan WSMsg {
	cons := []chan WSMsg{}
	for key, con := range wss.writeChans {
		username := strings.Split(key, "\uffff")[0]
		if username == searchname {
			cons = append(cons, con)
		}
	}
	return cons
}

func (wss *WSServer) CleanWriteChan(name string) {
	delete(wss.writeChans, name)
}

func (wss *WSServer) RegisterHandler(handlers []Handler) {
	if wss.handler == nil {
		wss.handler = map[string]func(WSMsg, gamelangpb.User) (WSMsg, error){}
	}
	for _, handler := range handlers {
		log.Println("Registrering handler for MType: ", handler.MType)
		wss.handler[handler.MType] = handler.Callback
	}
}

type rawMsg struct {
	mtype int
	p     []byte
}

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
	}

	user := "not logged in"
	defer wss.CleanWriteChan(fmt.Sprintf("%s\uffff%s", user, r.RemoteAddr))
	defer c.Close()
	var userObj *gamelangpb.User
	var writeChan chan WSMsg

	for {
		mt, message, err := c.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			break
		}
		msg := WSMsg{}
		err = json.Unmarshal(message, &msg)

		if err != nil {
			log.Println("error unmarshalling message:", err)
			log.Printf("recv: %s", message)
			return
		}

		if user == "not logged in" && msg.Mtype != "login" {
			msg := WSMsg{Mtype: "login-needed", Payload: ""}
			jsonMsg, _ := json.Marshal(msg)
			c.WriteMessage(mt, jsonMsg)
			continue
		}

		if msg.Mtype == "login" {
			log.Println("login request:", msg.Payload)
			if user == "not logged in" {
				creds := strings.Split(msg.Payload, "\uffff")
				if len(creds) != 2 {
					log.Printf("login failed for \"%s\"", msg.Payload)
					msg := WSMsg{Mtype: "login-failed", Payload: ""}
					wss.answer(mt, c, msg)
					return
				}
				username := creds[0]
				password := creds[1]
				var ok bool
				if wss.userclient != nil {
					log.Print("checking user via userclient")
					checkuser := gamelangpb.User{Name: username, Password: []byte(password)}
					var err error
					var tmpUserObj *gamelangpb.User
					tmpUserObj, err = wss.userclient.CheckUser(context.Background(), &checkuser)
					if err != nil {
						ok = false
						log.Print(err.Error())
					} else {
						userObj = tmpUserObj
						ok = true
					}
				} else {
					log.Print("checking user via usermanager")
					var tmpUserObj gamelangpb.User
					tmpUserObj, ok = wss.usermanager.CheckUser(username, password)
					if ok {
						userObj = &tmpUserObj
					}
				}
				if !ok {
					log.Printf("login failed for \"%s\" with \"%s\"", username, msg.Payload)
					msg := WSMsg{Mtype: "login-failed", Payload: ""}
					wss.answer(mt, c, msg)
					return
				}
				user = username
				writeChan = make(chan WSMsg, 20)
				wss.writeChans[fmt.Sprintf("%s\uffff%s", username, r.RemoteAddr)] = writeChan
				go func() {
					for {
						newMsg := <-writeChan
						err := wss.answer(websocket.BinaryMessage, c, newMsg)
						if err != nil {
							if strings.Contains(err.Error(), "use of closed network connection") || strings.Contains(err.Error(), "connection reset by peer") {
								delete(wss.writeChans, fmt.Sprintf("%s\uffff%s", username, r.RemoteAddr))
							}
							log.Println("websocket writer error:", err.Error())
						}
					}
				}()
				log.Printf("login succeeded for \"%s\"", username)
				msg := WSMsg{Mtype: "login-succeeded", Payload: username}
				writeChan <- msg
			}
		} else {
			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: ""}
				writeChan <- remsg
				return
			}
			remsg, err := handler(msg, *userObj)
			if err != nil {
				log.Println("handler failed", err.Error())
				writeChan <- msg
			}
			writeChan <- remsg
		}
	}
}

func (wss WSServer) answer(mt int, c *websocket.Conn, msg WSMsg) error {
	jsonMsg, err := json.Marshal(msg)
	if err != nil {
		log.Printf("write: err marshalling \"%s\"\n%s", msg, err.Error())
		return err
	}
	err = c.WriteMessage(mt, jsonMsg)
	if err != nil {
		log.Println("write:", err)
		return err
	}
	return nil
}