func onAccept(s *gotalk.Sock) { // Keep track of connected sockets socksmu.Lock() defer socksmu.Unlock() Clients[s] = 1 s.CloseHandler = func(s *gotalk.Sock, _ int) { socksmu.Lock() defer socksmu.Unlock() delete(Clients, s) } }
// When a new node connects add said node to the Server struct Nodes // and process the node checkin. func (server *Server) onAccept(s *gotalk.Sock) { server.mu.Lock() defer server.mu.Unlock() go func() { var resp Message err := s.Request("checkin", nil, &resp) if err != nil { Log.Fatalf("%s", err) } Log.Infof("New node connected. (%s / %s)", resp.Data["name"], resp.Data["id"]) if _, ok := resp.Data["envdb-version"]; !ok { s.BufferNotify("die", []byte("This version of Envdb is out of date. Please upgrade.")) return } if !VersionCheck(Version, resp.Data["envdb-version"].(string)) { s.BufferNotify("die", []byte("Envdb version mismatch")) return } node := &NodeData{ Id: resp.Data["id"].(string), Name: resp.Data["name"].(string), EnvdbVersion: resp.Data["envdb-version"].(string), Online: true, Socket: s, OsQuery: resp.Data["osquery"].(bool), OsQueryVersion: resp.Data["osquery-version"].(string), OsQueryConfigPath: resp.Data["osquery-config-path"].(string), Ip: resp.Data["ip"].(string), Hostname: resp.Data["hostname"].(string), PendingDelete: false, Os: resp.Data["os"].(string), } server.Nodes[s] = node if _, err := NodeUpdateOrCreate(node); err != nil { Log.Error("unable to create or update node record") Log.Error("Error: ", err) } WebSocketSend("node-update", node) s.CloseHandler = func(s *gotalk.Sock, _ int) { server.mu.Lock() defer server.mu.Unlock() node := server.Nodes[s] node.Online = false Log.Infof("Node disconnected. (%s / %s)", node.Name, node.Id) dbNode, err := NodeUpdateOrCreate(node) if err != nil { Log.Error("unable to update node record") Log.Error("Error: ", err) } if dbNode.PendingDelete { dbNode.Delete() } else { WebSocketSend("node-update", node) } delete(server.Nodes, s) } }() }