fix: [SyslogRelay] adjust the logic of stop

electricbubble · Apr 24, 2021 · b6b9bb3 · b6b9bb3
1 parent c40aa39
commit b6b9bb3
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/lockdown_test.go b/lockdown_test.go
@@ -53,6 +53,7 @@ func Test_lockdown_SyslogRelayService(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	syslogRelaySrv.Stop()
 
 	lines := syslogRelaySrv.Lines()
 

diff --git a/syslogrelay.go b/syslogrelay.go
@@ -12,8 +12,9 @@ var _ SyslogRelay = (*syslogRelay)(nil)
 
 func newSyslogRelay(client *libimobiledevice.SyslogRelayClient) *syslogRelay {
 	r := &syslogRelay{
-		client: client,
-		stop:   make(chan bool),
+		client:    client,
+		stop:      make(chan bool),
+		isReading: false,
 	}
 	r.reader = bufio.NewReader(r.client.InnerConn().RawConn())
 	return r
@@ -22,18 +23,23 @@ func newSyslogRelay(client *libimobiledevice.SyslogRelayClient) *syslogRelay {
 type syslogRelay struct {
 	client *libimobiledevice.SyslogRelayClient
 
-	reader *bufio.Reader
-	stop   chan bool
+	reader    *bufio.Reader
+	stop      chan bool
+	isReading bool
 }
 
 func (r *syslogRelay) Lines() <-chan string {
 	out := make(chan string)
+	r.isReading = true
 
 	go func() {
+		defer func() {
+			close(out)
+			r.isReading = false
+		}()
 		for {
 			select {
 			case <-r.stop:
-				close(out)
 				return
 			default:
 				bs, err := r.readLine()
@@ -54,7 +60,9 @@ func (r *syslogRelay) Lines() <-chan string {
 }
 
 func (r *syslogRelay) Stop() {
-	r.stop <- true
+	if r.isReading {
+		r.stop <- true
+	}
 }
 
 func (r *syslogRelay) readLine() ([]byte, error) {