From 291e18099050ff9e19b8ee25c2ffad68d9baafef Mon Sep 17 00:00:00 2001
From: kim <89579420+NyaaaWhatsUpDoc@users.noreply.github.com>
Date: Tue, 20 Feb 2024 18:07:49 +0000
Subject: [bugfix] fix possible mutex lockup during streaming code (#2633)

* rewrite Stream{} to use much less mutex locking, update related code

* use new context for the stream context

* ensure stream gets closed on return of writeTo / readFrom WSConn()

* ensure stream write timeout gets cancelled

* remove embedded context type from Stream{}, reformat log messages for consistency

* use c.Request.Context() for context passed into Stream().Open()

* only return 1 boolean, fix tests to expect multiple stream types in messages

* changes to ping logic

* further improved ping logic

* don't export unused function types, update message sending to only include relevant stream type

* ensure stream gets closed :facepalm:

* update to error log on failed json marshal (instead of panic)

* inverse websocket read error checking to _ignore_ expected close errors
---
 internal/processing/workers/fromfediapi_test.go | 53 +++++++++----------------
 1 file changed, 19 insertions(+), 34 deletions(-)

(limited to 'internal/processing/workers/fromfediapi_test.go')

diff --git a/internal/processing/workers/fromfediapi_test.go b/internal/processing/workers/fromfediapi_test.go
index 799eaf2dc..446355628 100644
--- a/internal/processing/workers/fromfediapi_test.go
+++ b/internal/processing/workers/fromfediapi_test.go
@@ -130,14 +130,9 @@ func (suite *FromFediAPITestSuite) TestProcessReplyMention() {
 	suite.Equal(replyingStatus.ID, notif.StatusID)
 	suite.False(*notif.Read)
 
-	// the notification should be streamed
-	var msg *stream.Message
-	select {
-	case msg = <-wssStream.Messages:
-		// fine
-	case <-time.After(5 * time.Second):
-		suite.FailNow("no message from wssStream")
-	}
+	ctx, _ := context.WithTimeout(context.Background(), time.Second*5)
+	msg, ok := wssStream.Recv(ctx)
+	suite.True(ok)
 
 	suite.Equal(stream.EventTypeNotification, msg.Event)
 	suite.NotEmpty(msg.Payload)
@@ -203,14 +198,10 @@ func (suite *FromFediAPITestSuite) TestProcessFave() {
 	suite.Equal(fave.StatusID, notif.StatusID)
 	suite.False(*notif.Read)
 
-	// 2. a notification should be streamed
-	var msg *stream.Message
-	select {
-	case msg = <-wssStream.Messages:
-		// fine
-	case <-time.After(5 * time.Second):
-		suite.FailNow("no message from wssStream")
-	}
+	ctx, _ := context.WithTimeout(context.Background(), time.Second*5)
+	msg, ok := wssStream.Recv(ctx)
+	suite.True(ok)
+
 	suite.Equal(stream.EventTypeNotification, msg.Event)
 	suite.NotEmpty(msg.Payload)
 	suite.EqualValues([]string{stream.TimelineNotifications}, msg.Stream)
@@ -277,7 +268,9 @@ func (suite *FromFediAPITestSuite) TestProcessFaveWithDifferentReceivingAccount(
 	suite.False(*notif.Read)
 
 	// 2. no notification should be streamed to the account that received the fave message, because they weren't the target
-	suite.Empty(wssStream.Messages)
+	ctx, _ := context.WithTimeout(context.Background(), time.Second*5)
+	_, ok := wssStream.Recv(ctx)
+	suite.False(ok)
 }
 
 func (suite *FromFediAPITestSuite) TestProcessAccountDelete() {
@@ -405,14 +398,10 @@ func (suite *FromFediAPITestSuite) TestProcessFollowRequestLocked() {
 	})
 	suite.NoError(err)
 
-	// a notification should be streamed
-	var msg *stream.Message
-	select {
-	case msg = <-wssStream.Messages:
-		// fine
-	case <-time.After(5 * time.Second):
-		suite.FailNow("no message from wssStream")
-	}
+	ctx, _ = context.WithTimeout(ctx, time.Second*5)
+	msg, ok := wssStream.Recv(context.Background())
+	suite.True(ok)
+
 	suite.Equal(stream.EventTypeNotification, msg.Event)
 	suite.NotEmpty(msg.Payload)
 	suite.EqualValues([]string{stream.TimelineHome}, msg.Stream)
@@ -423,7 +412,7 @@ func (suite *FromFediAPITestSuite) TestProcessFollowRequestLocked() {
 	suite.Equal(originAccount.ID, notif.Account.ID)
 
 	// no messages should have been sent out, since we didn't need to federate an accept
-	suite.Empty(suite.httpClient.SentMessages)
+	suite.Empty(&suite.httpClient.SentMessages)
 }
 
 func (suite *FromFediAPITestSuite) TestProcessFollowRequestUnlocked() {
@@ -503,14 +492,10 @@ func (suite *FromFediAPITestSuite) TestProcessFollowRequestUnlocked() {
 	suite.Equal(originAccount.URI, accept.To)
 	suite.Equal("Accept", accept.Type)
 
-	// a notification should be streamed
-	var msg *stream.Message
-	select {
-	case msg = <-wssStream.Messages:
-		// fine
-	case <-time.After(5 * time.Second):
-		suite.FailNow("no message from wssStream")
-	}
+	ctx, _ = context.WithTimeout(ctx, time.Second*5)
+	msg, ok := wssStream.Recv(context.Background())
+	suite.True(ok)
+
 	suite.Equal(stream.EventTypeNotification, msg.Event)
 	suite.NotEmpty(msg.Payload)
 	suite.EqualValues([]string{stream.TimelineHome}, msg.Stream)
-- 
cgit v1.2.3