Fix race condition on clients map

Author: g3force

internal/app/rcon/server.go

@@ -27,6 +27,7 @@ type Server struct {
 	listener          net.Listener
 	running           bool
 	mutex             sync.Mutex
+	clientsMutex      sync.Mutex
 }
 
 type Client struct {
@@ -76,12 +77,35 @@ func (s *Server) Stop() {
 	if err := s.listener.Close(); err != nil {
 		log.Printf("Could not close listener: %v", err)
 	}
-	for id := range s.clients {
+	for _, id := range s.GetClientIds() {
 		s.CloseConnection(id)
 	}
 	s.listener = nil
 }
 
+func (s *Server) GetClient(id string) (*Client, bool) {
+	s.clientsMutex.Lock()
+	defer s.clientsMutex.Unlock()
+	client, ok := s.clients[id]
+	return client, ok
+}
+
+func (s *Server) PutClient(id string, client *Client) {
+	s.clientsMutex.Lock()
+	defer s.clientsMutex.Unlock()
+	s.clients[id] = client
+}
+
+func (s *Server) GetClientIds() []string {
+	s.clientsMutex.Lock()
+	defer s.clientsMutex.Unlock()
+	var clients []string
+	for k := range s.clients {
+		clients = append(clients, k)
+	}
+	return clients
+}
+
 func (s *Server) listen() {
 	log.Print("Listening on ", s.address)
 
@@ -132,6 +156,8 @@ func (s *Server) loadTlsConfig() (*tls.Config, error) {
 }
 
 func (s *Server) CloseConnection(id string) {
+	s.clientsMutex.Lock()
+	defer s.clientsMutex.Unlock()
 	delete(s.clients, id)
 	log.Printf("Connection to %v closed", id)
 }

internal/app/rcon/server_autoref.go

@@ -39,12 +39,8 @@ func (c *AutoRefClient) receiveRegistration(reader *bufio.Reader, server *AutoRe
 		return errors.New("No identifier specified")
 	}
 	c.id = *registration.Identifier
-	if _, exists := server.clients[c.id]; exists {
-		var clients []string
-		for k := range server.clients {
-			clients = append(clients, k)
-		}
-		return errors.Errorf("AutoRef Client with given identifier already registered: %v", clients)
+	if _, exists := server.GetClient(c.id); exists {
+		return errors.Errorf("AutoRef Client with given identifier already registered: %v", server.GetClientIds())
 	}
 	c.pubKey = server.trustedKeys[c.id]
 	if c.pubKey != nil {
@@ -78,7 +74,7 @@ func (s *AutoRefServer) handleClientConnection(conn net.Conn) {
 		return
 	}
 
-	s.clients[client.id] = client.Client
+	s.PutClient(client.id, client.Client)
 	defer func() {
 		s.gcEngine.UpdateGcState(func(gcState *engine.GcState) {
 			delete(gcState.AutoRefState, client.id)

internal/app/rcon/server_remotecontrol.go

@@ -48,7 +48,7 @@ func (c *RemoteControlClient) receiveRegistration(reader *bufio.Reader, server *
 	}
 	c.team = registration.Team
 	c.id = registration.Team.String()
-	if _, exists := server.clients[c.id]; exists {
+	if _, exists := server.GetClient(c.id); exists {
 		return errors.New("Team already registered: " + c.id)
 	}
 	c.pubKey = server.trustedKeys[c.id]
@@ -88,7 +88,7 @@ func (s *RemoteControlServer) handleClientConnection(conn net.Conn) {
 		return
 	}
 
-	s.clients[client.id] = client.Client
+	s.PutClient(client.id, client.Client)
 	defer func() {
 		s.updateConnectionState(client, false)
 		s.CloseConnection(client.id)
@@ -130,7 +130,7 @@ func (s *RemoteControlServer) updateConnectionState(client RemoteControlClient,
 }
 
 func (s *RemoteControlServer) SendRequest(teamName string, request *ControllerToRemoteControl) error {
-	if client, ok := s.clients[teamName]; ok {
+	if client, ok := s.GetClient(teamName); ok {
 		return sslconn.SendMessage(client.conn, request)
 	}
 	return errors.Errorf("Remote control client '%v' not connected", teamName)

internal/app/rcon/server_team.go

@@ -64,7 +64,7 @@ func (c *TeamClient) receiveRegistration(reader *bufio.Reader, server *TeamServe
 	c.id = team.String() + "-" + *registration.TeamName
 	c.team = team
 	c.teamName = *registration.TeamName
-	if _, exists := server.clients[c.id]; exists {
+	if _, exists := server.GetClient(c.id); exists {
 		return errors.New("Team with given name already registered: " + c.id)
 	}
 	c.pubKey = server.trustedKeys[*registration.TeamName]
@@ -110,7 +110,7 @@ func (s *TeamServer) handleClientConnection(conn net.Conn) {
 		return
 	}
 
-	s.clients[client.id] = client.Client
+	s.PutClient(client.id, client.Client)
 	defer func() {
 		s.gcEngine.UpdateGcState(func(gcState *engine.GcState) {
 			if teamState, ok := gcState.TeamState[client.team.String()]; ok {
@@ -157,7 +157,7 @@ func (s *TeamServer) handleClientConnection(conn net.Conn) {
 }
 
 func (s *TeamServer) SendRequest(teamName string, request *ControllerToTeam) error {
-	if client, ok := s.clients[teamName]; ok {
+	if client, ok := s.GetClient(teamName); ok {
 		return sslconn.SendMessage(client.conn, request)
 	}
 	return errors.Errorf("Team Client '%v' not connected", teamName)