bot,api: fix stop channel race

Author: davseby

api.go

@@ -27,18 +27,21 @@ func (b *Bot) Raw(method string, payload interface{}) ([]byte, error) {
 		return nil, err
 	}
 
-	// Cancel the request immediately without waiting for the timeout  when bot is about to stop.
+	// Cancel the request immediately without waiting for the timeout
+	// when bot is about to stop.
 	// This may become important if doing long polling with long timeout.
-	exit := make(chan struct{})
-	defer close(exit)
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
 	go func() {
+		b.stopMu.RLock()
+		stopCh := b.stopClient
+		b.stopMu.RUnlock()
+
 		select {
-		case <-b.stopClient:
+		case <-stopCh:
 			cancel()
-		case <-exit:
+		case <-ctx.Done():
 		}
 	}()
 

bot.go

@@ -10,6 +10,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -81,7 +82,9 @@ type Bot struct {
 	parseMode   ParseMode
 	stop        chan chan struct{}
 	client      *http.Client
-	stopClient  chan struct{}
+
+	stopMu     sync.RWMutex
+	stopClient chan struct{}
 }
 
 // Settings represents a utility struct for passing certain
@@ -198,10 +201,14 @@ func (b *Bot) Start() {
 	}
 
 	// do nothing if called twice
+	b.stopMu.Lock()
 	if b.stopClient != nil {
+		b.stopMu.Unlock()
 		return
 	}
+
 	b.stopClient = make(chan struct{})
+	b.stopMu.Unlock()
 
 	stop := make(chan struct{})
 	stopConfirm := make(chan struct{})
@@ -221,17 +228,20 @@ func (b *Bot) Start() {
 			close(stop)
 			<-stopConfirm
 			close(confirm)
-			b.stopClient = nil
 			return
 		}
 	}
 }
 
 // Stop gracefully shuts the poller down.
 func (b *Bot) Stop() {
+	b.stopMu.Lock()
 	if b.stopClient != nil {
 		close(b.stopClient)
+		b.stopClient = nil
 	}
+	b.stopMu.Unlock()
+
 	confirm := make(chan struct{})
 	b.stop <- confirm
 	<-confirm