Fix some data races

4 files changed, 28 insertions, 12 deletions

diff --git a/hotline/client_conn.go b/hotline/client_conn.go
index 4a471b3..65e18ee 100644
--- a/hotline/client_conn.go
+++ b/hotline/client_conn.go
@@ -151,7 +151,7 @@ func (cc *ClientConn) Authorize(access int) bool {
 }
 
 // Disconnect notifies other clients that a client has disconnected
-func (cc ClientConn) Disconnect() {
+func (cc *ClientConn) Disconnect() {
 	cc.Server.mux.Lock()
 	defer cc.Server.mux.Unlock()
 
@@ -165,7 +165,7 @@ func (cc ClientConn) Disconnect() {
 }
 
 // notifyOthers sends transaction t to other clients connected to the server
-func (cc ClientConn) notifyOthers(t Transaction) {
+func (cc *ClientConn) notifyOthers(t Transaction) {
 	for _, c := range sortedClients(cc.Server.Clients) {
 		if c.ID != cc.ID && c.Agreed {
 			t.clientID = c.ID
diff --git a/hotline/server.go b/hotline/server.go
index 5ff8842..cf022fc 100644
--- a/hotline/server.go
+++ b/hotline/server.go
@@ -646,28 +646,42 @@ const dlFldrActionNextFile = 3
 // handleFileTransfer receives a client net.Conn from the file transfer server, performs the requested transfer type, then closes the connection
 func (s *Server) handleFileTransfer(conn io.ReadWriteCloser) error {
 	defer func() {
+
 		if err := conn.Close(); err != nil {
 			s.Logger.Errorw("error closing connection", "error", err)
 		}
 	}()
 
+	defer func() {
+		if r := recover(); r != nil {
+			fmt.Println("stacktrace from panic: \n" + string(debug.Stack()))
+			s.Logger.Errorw("PANIC", "err", r, "trace", string(debug.Stack()))
+		}
+	}()
+
 	txBuf := make([]byte, 16)
-	_, err := conn.Read(txBuf)
-	if err != nil {
+	if _, err := io.ReadFull(conn, txBuf); err != nil {
 		return err
 	}
 
 	var t transfer
-	_, err = t.Write(txBuf)
-	if err != nil {
+	if _, err := t.Write(txBuf); err != nil {
 		return err
 	}
 
 	transferRefNum := binary.BigEndian.Uint32(t.ReferenceNumber[:])
-	fileTransfer := s.FileTransfers[transferRefNum]
+	defer func() {
+		s.mux.Lock()
+		delete(s.FileTransfers, transferRefNum)
+		s.mux.Unlock()
+	}()
 
-	// delete single use transferRefNum
-	delete(s.FileTransfers, transferRefNum)
+	s.mux.Lock()
+	fileTransfer, ok := s.FileTransfers[transferRefNum]
+	s.mux.Unlock()
+	if !ok {
+		return errors.New("invalid transaction ID")
+	}
 
 	switch fileTransfer.Type {
 	case FileDownload:
diff --git a/hotline/transaction.go b/hotline/transaction.go
index 04fb4b6..e7bfa58 100644
--- a/hotline/transaction.go
+++ b/hotline/transaction.go
@@ -214,7 +214,7 @@ func ReadFields(paramCount []byte, buf []byte) ([]Field, error) {
 	return fields, nil
 }
 
-func (t Transaction) MarshalBinary() (data []byte, err error) {
+func (t *Transaction) MarshalBinary() (data []byte, err error) {
 	payloadSize := t.Size()
 
 	fieldCount := make([]byte, 2)
@@ -238,7 +238,7 @@ func (t Transaction) MarshalBinary() (data []byte, err error) {
 }
 
 // Size returns the total size of the transaction payload
-func (t Transaction) Size() []byte {
+func (t *Transaction) Size() []byte {
 	bs := make([]byte, 4)
 
 	fieldSize := 0
@@ -251,7 +251,7 @@ func (t Transaction) Size() []byte {
 	return bs
 }
 
-func (t Transaction) GetField(id int) Field {
+func (t *Transaction) GetField(id int) Field {
 	for _, field := range t.Fields {
 		if id == int(binary.BigEndian.Uint16(field.ID)) {
 			return field
diff --git a/hotline/transaction_handlers.go b/hotline/transaction_handlers.go
index 4a86831..abe2c54 100644
--- a/hotline/transaction_handlers.go
+++ b/hotline/transaction_handlers.go
@@ -1521,12 +1521,14 @@ func HandleUploadFile(cc *ClientConn, t *Transaction) (res []Transaction, err er
 	transactionRef := cc.Server.NewTransactionRef()
 	data := binary.BigEndian.Uint32(transactionRef)
 
+	cc.Server.mux.Lock()
 	cc.Server.FileTransfers[data] = &FileTransfer{
 		FileName:        fileName,
 		FilePath:        filePath,
 		ReferenceNumber: transactionRef,
 		Type:            FileUpload,
 	}
+	cc.Server.mux.Unlock()
 
 	replyT := cc.NewReply(t, NewField(fieldRefNum, transactionRef))