Protect subscription completed flag with a mutex to avoid closing mul…

…tiple times the channel or writing to a closed channel.

conn.go

@@ -6,6 +6,7 @@ import (
 	"log"
 	"net"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/go-stomp/stomp/frame"
@@ -508,11 +509,12 @@ func (c *Conn) Subscribe(destination string, ack AckMode, opts ...func(*frame.Fr
 	}
 
 	sub := &Subscription{
-		id:          id,
-		destination: destination,
-		conn:        c,
-		ackMode:     ack,
-		C:           make(chan *Message, 16),
+		id:             id,
+		destination:    destination,
+		conn:           c,
+		ackMode:        ack,
+		C:              make(chan *Message, 16),
+		completedMutex: &sync.Mutex{},
 	}
 	go sub.readLoop(ch)
 

subscription.go

@@ -3,6 +3,7 @@ package stomp
 import (
 	"fmt"
 	"log"
+	"sync"
 
 	"github.com/go-stomp/stomp/frame"
 )
@@ -12,12 +13,13 @@ import (
 //
 // Once a client has subscribed, it can receive messages from the C channel.
 type Subscription struct {
-	C           chan *Message
-	id          string
-	destination string
-	conn        *Conn
-	ackMode     AckMode
-	completed   bool
+	C              chan *Message
+	id             string
+	destination    string
+	conn           *Conn
+	ackMode        AckMode
+	completed      bool
+	completedMutex *sync.Mutex
 }
 
 // BUG(jpj): If the client does not read messages from the Subscription.C
@@ -65,8 +67,12 @@ func (s *Subscription) Unsubscribe(opts ...func(*frame.Frame) error) error {
 	}
 
 	s.conn.sendFrame(f)
-	s.completed = true
-	close(s.C)
+	s.completedMutex.Lock()
+	if !s.completed {
+		s.completed = true
+		close(s.C)
+	}
+	s.completedMutex.Unlock()
 	return nil
 }
 
@@ -109,9 +115,11 @@ func (s *Subscription) readLoop(ch chan *frame.Frame) {
 				Header:       f.Header,
 				Body:         f.Body,
 			}
+			s.completedMutex.Lock()
 			if !s.completed {
 				s.C <- msg
 			}
+			s.completedMutex.Unlock()
 		} else if f.Command == frame.ERROR {
 			message, _ := f.Header.Contains(frame.Message)
 			text := fmt.Sprintf("Subscription %s: %s: ERROR message:%s",
@@ -131,11 +139,13 @@ func (s *Subscription) readLoop(ch chan *frame.Frame) {
 				Header:       f.Header,
 				Body:         f.Body,
 			}
+			s.completedMutex.Lock()
 			if !s.completed {
+				s.completed = true
 				s.C <- msg
+				close(s.C)
 			}
-			s.completed = true
-			close(s.C)
+			s.completedMutex.Unlock()
 			return
 		}
 	}