From 3c08964b38faad09668cedfa10413d7067d1fc1b Mon Sep 17 00:00:00 2001
From: Ben Grewell <BGrewell@gmail.com>
Date: Mon, 14 Jun 2021 14:00:32 -0700
Subject: [PATCH] add locks to map access

---
 .idea/workspace.xml |  2 +-
 controller.go       | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/.idea/workspace.xml b/.idea/workspace.xml
index b07a49a..baffb5a 100644
--- a/.idea/workspace.xml
+++ b/.idea/workspace.xml
@@ -3,7 +3,7 @@
   <component name="ChangeListManager">
     <list default="true" id="fc2840de-29dc-4fca-8e0e-a283562f60ca" name="Default Changelist" comment="">
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/client.go" beforeDir="false" afterPath="$PROJECT_DIR$/client.go" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/controller.go" beforeDir="false" afterPath="$PROJECT_DIR$/controller.go" afterDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
diff --git a/controller.go b/controller.go
index 22f67d2..73dd202 100644
--- a/controller.go
+++ b/controller.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"net"
+	"sync"
 	"time"
 
 	api "github.com/BGrewell/go-iperf/api/go"
@@ -16,6 +17,8 @@ func NewController(port int) (controller *Controller, err error) {
 		Port:    port,
 		clients: make(map[string]*Client),
 		servers: make(map[string]*Server),
+		clientLock: sync.Mutex{},
+		serverLock: sync.Mutex{},
 	}
 	err = c.startListener()
 
@@ -36,6 +39,8 @@ type Controller struct {
 	api.UnimplementedCommandServer
 	Port int
 	cmdClient api.CommandClient
+	clientLock sync.Mutex
+	serverLock sync.Mutex
 	clients map[string]*Client
 	servers map[string]*Server
 }
@@ -52,7 +57,9 @@ func (c *Controller) GrpcRequestServer(context.Context, *api.StartServerRequest)
 		return nil, err
 	}
 
+	c.serverLock.Lock()
 	c.servers[srv.Id] = srv
+	c.serverLock.Unlock()
 	
 	reply := &api.StartServerResponse{
 		Id:         srv.Id,
@@ -88,13 +95,17 @@ func (c *Controller) NewServer() (server *Server, err error) {
 	freePort, err := GetUnusedTcpPort()
 	s := NewServer()
 	s.SetPort(freePort)
+	c.serverLock.Lock()
 	c.servers[s.Id] = s
+	c.serverLock.Unlock()
 	return s, nil
 }
 
 // StopServer shuts down an iperf server and frees any actively used resources
 func (c *Controller) StopServer(id string) (err error) {
+	c.serverLock.Lock()
 	delete(c.servers, id)
+	c.serverLock.Unlock()
 	return nil
 }
 
@@ -114,14 +125,18 @@ func (c *Controller) NewClient(serverAddr string) (client *Client, err error) {
 	
 	cli := NewClient(serverAddr)
 	cli.SetPort(srvPort)
+	c.clientLock.Lock()
 	c.clients[cli.Id] = cli
+	c.clientLock.Unlock()
 
 	return cli, nil
 }
 
 // StopClient will clean up the server side connection and shut down any actively used resources
 func (c *Controller) StopClient(id string) (err error) {
+	c.clientLock.Lock()
 	delete(c.clients, id)
+	c.clientLock.Unlock()
 	return nil
 }