package ws import ( "encoding/json" "fmt" "github.com/gorilla/websocket" "goseg/auth" "goseg/broadcast" "goseg/config" "goseg/structs" "net/http" "time" ) var ( upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, } ) // func handleConnection(c *websocket.Conn) { // // Read the first message from the client which should be the token // messageType, p, err := c.ReadMessage() // if err != nil { // config.Logger.Error(fmt.Errorf("%v",err)) // return // } // token := string(p) // // Verify the token // isValid, _, err := CheckToken(token, c, false) // 'false' assumes it's not a setup // if !isValid || err != nil { // config.Logger.Info("Invalid token provided by client.") // c.Close() // return // } // // rest of logic // } // switch on ws event cases func WsHandler(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { config.Logger.Error(fmt.Sprintf("Couldn't upgrade websocket connection: %v", err)) return } // manage broadcasts and clients thru the broadcast package broadcast.RegisterClient(conn) defer broadcast.UnregisterClient(conn) // keepalive for ws conn.SetPongHandler(func(string) error { conn.SetReadDeadline(time.Now().Add(60 * time.Second)) return nil }) pingInterval := 15 * time.Second go func() { ticker := time.NewTicker(pingInterval) defer ticker.Stop() for { select { case <-ticker.C: if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } }() for { _, msg, err := conn.ReadMessage() if err != nil { return } var prelim structs.WsType var payload structs.WsPayload if err := json.Unmarshal(msg, &prelim); err != nil { fmt.Println("Error unmarshalling message:", err) continue } switch prelim.Payload.Type { case "login": if err = loginHandler(msg, payload); err != nil { config.Logger.Error(fmt.Sprintf("%v", err)) } case "verify": if err = verifyHandler(msg, payload, r, conn); err != nil { config.Logger.Error(fmt.Sprintf("%v", err)) } case "setup": config.Logger.Info("Setup") // setup.Setup(payload) case "new_ship": config.Logger.Info("New ship") case "pier_upload": config.Logger.Info("Pier upload") case "password": config.Logger.Info("Password") case "system": config.Logger.Info("System") case "startram": config.Logger.Info("StarTram") case "urbit": config.Logger.Info("Urbit") case "support": config.Logger.Info("Support") case "broadcast": if err := broadcast.BroadcastToClients(); err != nil { errmsg := fmt.Sprintf("Unable to broadcast to peer(s): %v", err) config.Logger.Error(errmsg) } default: errmsg := fmt.Sprintf("Unknown request type: %s", prelim.Payload.Type) config.Logger.Warn(errmsg) } } } // validate password and add to auth session func loginHandler(msg []byte, payload structs.WsPayload) error { config.Logger.Info("Login") now := time.Now().Format("2006-01-02_15:04:05") payload.Payload = structs.WsLoginPayload{} if err := json.Unmarshal(msg, &payload); err != nil { return fmt.Errorf("Error unmarshalling message: %v", err) } loginPayload, ok := payload.Payload.(structs.WsLoginPayload) if !ok { return fmt.Errorf("Error casting to LoginPayload") } isAuthenticated := auth.AuthenticateLogin(loginPayload.Password) if isAuthenticated { if err := auth.AddSession(payload.Token.ID, payload.Token.Token, now, true); err != nil { return fmt.Errorf("Unable to process login: %v", err) } } else { config.Logger.Info("Login failed") } return nil } func verifyHandler(msg []byte, payload structs.WsPayload, r *http.Request, conn *websocket.Conn) error { config.Logger.Info("Verify") payload.Payload = structs.WsLoginPayload{} // if we can't unmarshal, assume no token if err := json.Unmarshal(msg, &payload); err != nil { resp, err := auth.CreateToken(conn, r, false) if err != nil { fmt.Errorf("Couldn't create token: %v", err) } respJson, err := json.Marshal(resp) if err != nil { return fmt.Errorf("Error marshalling token: %v", err) } if err := conn.WriteMessage(websocket.TextMessage, respJson); err != nil { return fmt.Errorf("Error writing response: %v", err) } } return nil }