Fix some data races

[rbdr/mobius] / hotline / server.go

diff --git a/hotline/server.go b/hotline/server.go
index aebff6ed8af71e73b4534b631fb127ea0d27eb00..cf022fc9f9e3a2d3a6649253f1256e371cb079a3 100644 (file)
--- 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:
@@ -688,9 +702,11 @@ func (s *Server) handleFileTransfer(conn io.ReadWriteCloser) error {
 
                s.Logger.Infow("File download started", "filePath", fullFilePath, "transactionRef", fileTransfer.ReferenceNumber)
 
-               // Start by sending flat file object to client
-               if _, err := conn.Write(ffo.BinaryMarshal()); err != nil {
-                       return err
+               if fileTransfer.options == nil {
+                       // Start by sending flat file object to client
+                       if _, err := conn.Write(ffo.BinaryMarshal()); err != nil {
+                               return err
+                       }
                }
 
                file, err := FS.Open(fullFilePath)