// +build !no_ws package main import "fmt" import "sync" import "time" import "bytes" import "strconv" import "runtime" import "net/http" import "github.com/gorilla/websocket" import "github.com/shirou/gopsutil/cpu" import "github.com/shirou/gopsutil/mem" type WS_User struct { conn *websocket.Conn User *User } type WS_Hub struct { online_users map[int]*WS_User online_guests map[*WS_User]bool guests sync.RWMutex users sync.RWMutex } var ws_hub WS_Hub var ws_upgrader = websocket.Upgrader{ReadBufferSize:1024,WriteBufferSize:1024} func init() { enable_websockets = true admin_stats_watchers = make(map[*WS_User]bool) ws_hub = WS_Hub{ online_users: make(map[int]*WS_User), online_guests: make(map[*WS_User]bool), } } func (hub *WS_Hub) GuestCount() int { defer hub.guests.RUnlock() hub.guests.RLock() return len(hub.online_guests) } func (hub *WS_Hub) UserCount() int { defer hub.users.RUnlock() hub.users.RLock() return len(hub.online_users) } func route_websockets(w http.ResponseWriter, r *http.Request) { user, ok := SimpleSessionCheck(w,r) if !ok { return } conn, err := ws_upgrader.Upgrade(w,r,nil) if err != nil { return } userptr, err := users.CascadeGet(user.ID) if err != nil && err != ErrStoreCapacityOverflow { return } ws_user := &WS_User{conn,userptr} if user.ID == 0 { ws_hub.guests.Lock() ws_hub.online_guests[ws_user] = true ws_hub.guests.Unlock() } else { ws_hub.users.Lock() ws_hub.online_users[user.ID] = ws_user ws_hub.users.Unlock() } //conn.SetReadLimit(/* put the max request size from earlier here? */) //conn.SetReadDeadline(time.Now().Add(60 * time.Second)) var current_page []byte for { _, message, err := conn.ReadMessage() if err != nil { if user.ID == 0 { ws_hub.guests.Lock() delete(ws_hub.online_guests,ws_user) ws_hub.guests.Unlock() } else { ws_hub.users.Lock() delete(ws_hub.online_users,user.ID) ws_hub.users.Unlock() } break } //fmt.Println("Message",message) //fmt.Println("Message",string(message)) messages := bytes.Split(message,[]byte("\r")) for _, msg := range messages { //fmt.Println("Submessage",msg) //fmt.Println("Submessage",string(msg)) if bytes.HasPrefix(msg,[]byte("page ")) { msgblocks := bytes.SplitN(msg,[]byte(" "),2) if len(msgblocks) < 2 { continue } if !bytes.Equal(msgblocks[1],current_page) { ws_leave_page(ws_user, current_page) current_page = msgblocks[1] //fmt.Println("Current Page: ",current_page) //fmt.Println("Current Page: ",string(current_page)) ws_page_responses(ws_user, current_page) } } /*if bytes.Equal(message,[]byte(`start-view`)) { } else if bytes.Equal(message,[]byte(`end-view`)) { }*/ } } conn.Close() } func ws_page_responses(ws_user *WS_User, page []byte) { switch(string(page)) { case "/panel/": //fmt.Println("/panel/ WS Route") w, err := ws_user.conn.NextWriter(websocket.TextMessage) if err != nil { //fmt.Println(err.Error()) return } fmt.Println(ws_hub.online_users) uonline := ws_hub.UserCount() gonline := ws_hub.GuestCount() totonline := uonline + gonline w.Write([]byte("set #dash-totonline " + strconv.Itoa(totonline) + " online\r")) w.Write([]byte("set #dash-gonline " + strconv.Itoa(gonline) + " guests online\r")) w.Write([]byte("set #dash-uonline " + strconv.Itoa(uonline) + " users online\r")) w.Close() // Listen for changes and inform the admins... admin_stats_mutex.Lock() watchers := len(admin_stats_watchers) admin_stats_watchers[ws_user] = true if watchers == 0 { go admin_stats_ticker() } admin_stats_mutex.Unlock() } } func ws_leave_page(ws_user *WS_User, page []byte) { switch(string(page)) { case "/panel/": admin_stats_mutex.Lock() delete(admin_stats_watchers,ws_user) admin_stats_mutex.Unlock() } } var admin_stats_watchers map[*WS_User]bool var admin_stats_mutex sync.RWMutex func admin_stats_ticker() { time.Sleep(time.Second) var last_uonline int = -1 var last_gonline int = -1 var last_totonline int = -1 var last_cpu_perc int = -1 var last_available_ram int64 = -1 var no_stat_updates bool = false var onlineColour, onlineGuestsColour, onlineUsersColour, cpustr, cpuColour, ramstr, ramColour string var cpuerr, ramerr error var memres *mem.VirtualMemoryStat var cpu_perc []float64 AdminStatLoop: for { //fmt.Println("tick tock") admin_stats_mutex.RLock() watch_count := len(admin_stats_watchers) admin_stats_mutex.RUnlock() if watch_count == 0 { break AdminStatLoop } cpu_perc, cpuerr = cpu.Percent(time.Duration(time.Second),true) memres, ramerr = mem.VirtualMemory() uonline := ws_hub.UserCount() gonline := ws_hub.GuestCount() totonline := uonline + gonline // It's far more likely that the CPU Usage will change than the other stats, so we'll optimise them seperately... no_stat_updates = (uonline == last_uonline && gonline == last_gonline && totonline == last_totonline) if no_stat_updates && int(cpu_perc[0]) == last_cpu_perc && last_available_ram == int64(memres.Available) { time.Sleep(time.Second) continue } if !no_stat_updates { if totonline > 10 { onlineColour = "stat_green" } else if totonline > 3 { onlineColour = "stat_orange" } else { onlineColour = "stat_red" } if gonline > 10 { onlineGuestsColour = "stat_green" } else if gonline > 1 { onlineGuestsColour = "stat_orange" } else { onlineGuestsColour = "stat_red" } if uonline > 5 { onlineUsersColour = "stat_green" } else if uonline > 1 { onlineUsersColour = "stat_orange" } else { onlineUsersColour = "stat_red" } } if cpuerr != nil { cpustr = "Unknown" } else { calcperc := int(cpu_perc[0]) / runtime.NumCPU() cpustr = strconv.Itoa(calcperc) if calcperc < 30 { cpuColour = "stat_green" } else if calcperc < 75 { cpuColour = "stat_orange" } else { cpuColour = "stat_red" } } if ramerr != nil { ramstr = "Unknown" } else { total_count, total_unit := convert_byte_unit(float64(memres.Total)) used_count := convert_byte_in_unit(float64(memres.Total - memres.Available),total_unit) // Round totals with .9s up, it's how most people see it anyway. Floats are notoriously imprecise, so do it off 0.85 var totstr string if (total_count - float64(int(total_count))) > 0.85 { used_count += 1.0 - (total_count - float64(int(total_count))) totstr = strconv.Itoa(int(total_count) + 1) } else { totstr = fmt.Sprintf("%.1f",total_count) } if used_count > total_count { used_count = total_count } ramstr = fmt.Sprintf("%.1f",used_count) + " / " + totstr + total_unit ramperc := ((memres.Total - memres.Available) * 100) / memres.Total if ramperc < 50 { ramColour = "stat_green" } else if ramperc < 75 { ramColour = "stat_orange" } else { ramColour = "stat_red" } } admin_stats_mutex.RLock() watchers := admin_stats_watchers admin_stats_mutex.RUnlock() for watcher, _ := range watchers { w, err := watcher.conn.NextWriter(websocket.TextMessage) if err != nil { fmt.Println(err.Error()) admin_stats_mutex.Lock() delete(admin_stats_watchers,watcher) admin_stats_mutex.Unlock() continue } if !no_stat_updates { w.Write([]byte("set #dash-totonline " + strconv.Itoa(totonline) + " online\r")) w.Write([]byte("set #dash-gonline " + strconv.Itoa(gonline) + " guests online\r")) w.Write([]byte("set #dash-uonline " + strconv.Itoa(uonline) + " users online\r")) w.Write([]byte("set-class #dash-totonline grid_stat " + onlineColour + "\r")) w.Write([]byte("set-class #dash-gonline grid_stat " + onlineGuestsColour + "\r")) w.Write([]byte("set-class #dash-uonline grid_stat " + onlineUsersColour + "\r")) } w.Write([]byte("set #dash-cpu CPU: " + cpustr + "%\r")) w.Write([]byte("set-class #dash-cpu grid_istat " + cpuColour + "\r")) w.Write([]byte("set #dash-ram RAM: " + ramstr + "\r")) w.Write([]byte("set-class #dash-ram grid_istat " + ramColour + "\r")) w.Close() } last_uonline = uonline last_gonline = gonline last_totonline = totonline last_cpu_perc = int(cpu_perc[0]) last_available_ram = int64(memres.Available) //time.Sleep(time.Second) } }