Merge pull request #897 from MykolaBalakin/feat/jsonrpc2-manual-receive-mode-upstream

feat(jsonrpc2): add manual receive mode to prevent message loss

README.md

@@ -160,3 +160,4 @@ There are a bunch of environmental variables that can be set inside the docker c
 * `JSON_RPC_IGNORE_AVATARS`: When set to `true`, avatars are not automatically downloaded in json-rpc mode (default: `false`)
 * `JSON_RPC_IGNORE_STICKERS`: When set to `true`, sticker packs are not automatically downloaded in json-rpc mode (default: `false`)
 * `JSON_RPC_TRUST_NEW_IDENTITIES`: Choose how to trust new identities in json-rpc mode. Supported values: `on-first-use`, `always`, `never`. (default: `on-first-use`)
+* `JSON_RPC_RECEIVE_MODE`: Controls when signal-cli pulls inbound messages from the Signal servers in json-rpc mode. Supported values: `on-start` (default), `manual`. In the default `on-start` mode, signal-cli auto-receives messages and pushes them as JSON-RPC notifications regardless of whether any websocket subscriber is currently attached to `/v1/receive/{number}` — messages arriving while no subscriber is attached (e.g. during a brief subscriber redeploy) are silently dropped. In `manual` mode signal-cli only fetches messages while at least one websocket subscriber is attached; while no subscriber is attached, messages remain queued server-side under Signal's normal retention rules and are delivered on the next subscriber re-attach. Recommended for any deployment where the websocket consumer is restarted or scaled.

src/api/api.go

@@ -573,7 +573,7 @@ func (a *Api) SendV2(c *gin.Context) {
 }
 
 func (a *Api) handleSignalReceive(ws *websocket.Conn, number string, stop chan struct{}) {
-	receiveChannel, channelUuid, err := a.signalClient.GetReceiveChannel()
+	receiveChannel, channelUuid, err := a.signalClient.GetReceiveChannel(number)
 	if err != nil {
 		log.Error("Couldn't get receive channel: ", err.Error())
 		return

src/client/client.go

@@ -53,6 +53,13 @@ const (
 	JsonRpc
 )
 
+type JsonRpcReceiveMode int
+
+const (
+	OnStart JsonRpcReceiveMode = iota + 1
+	Manual
+)
+
 type GroupLinkState int
 
 const (
@@ -406,10 +413,11 @@ type SignalClient struct {
 	signalCliApiConfig       *utils.SignalCliApiConfig
 	cliClient                *CliClient
 	receiveWebhookUrl        string
+	jsonRpcReceiveMode       JsonRpcReceiveMode
 }
 
 func NewSignalClient(signalCliConfig string, attachmentTmpDir string, avatarTmpDir string, signalCliMode SignalCliMode,
-	jsonRpc2ClientConfigPath string, signalCliApiConfigPath string, receiveWebhookUrl string) *SignalClient {
+	jsonRpc2ClientConfigPath string, signalCliApiConfigPath string, receiveWebhookUrl string, jsonRpcReceiveMode JsonRpcReceiveMode) *SignalClient {
 	return &SignalClient{
 		signalCliConfig:          signalCliConfig,
 		attachmentTmpDir:         attachmentTmpDir,
@@ -419,6 +427,7 @@ func NewSignalClient(signalCliConfig string, attachmentTmpDir string, avatarTmpD
 		jsonRpc2Clients:          make(map[string]*JsonRpc2Client),
 		signalCliApiConfigPath:   signalCliApiConfigPath,
 		receiveWebhookUrl:        receiveWebhookUrl,
+		jsonRpcReceiveMode:       jsonRpcReceiveMode,
 	}
 }
 
@@ -442,7 +451,7 @@ func (s *SignalClient) Init(maxRetries int) error {
 
 		tcpPortsNumberMapping := s.jsonRpc2ClientConfig.GetTcpPortsForNumbers()
 		for number, tcpPort := range tcpPortsNumberMapping {
-			s.jsonRpc2Clients[number] = NewJsonRpc2Client(s.signalCliApiConfig, number)
+			s.jsonRpc2Clients[number] = NewJsonRpc2Client(s.signalCliApiConfig, number, s.jsonRpcReceiveMode)
 			err := s.jsonRpc2Clients[number].Dial("127.0.0.1:"+strconv.FormatInt(tcpPort, 10), maxRetries)
 			if err != nil {
 				return err
@@ -1084,12 +1093,12 @@ func (s *SignalClient) Receive(number string, timeout int64, ignoreAttachments b
 	}
 }
 
-func (s *SignalClient) GetReceiveChannel() (chan JsonRpc2ReceivedMessage, string, error) {
+func (s *SignalClient) GetReceiveChannel(number string) (chan JsonRpc2ReceivedMessage, string, error) {
 	jsonRpc2Client, err := s.getJsonRpc2Client()
 	if err != nil {
 		return nil, "", err
 	}
-	return jsonRpc2Client.GetReceiveChannel()
+	return jsonRpc2Client.GetReceiveChannel(number)
 }
 
 func (s *SignalClient) RemoveReceiveChannel(channelUuid string) {

src/client/jsonrpc2.go

@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"net"
 	"net/http"
 	"strconv"
@@ -56,23 +57,38 @@ func (r *RateLimitErrorType) Error() string {
 	return r.Err.Error()
 }
 
-type JsonRpc2Client struct {
-	conn                     net.Conn
-	receivedResponsesById    map[string]chan JsonRpc2MessageResponse
-	receivedMessagesChannels map[string]chan JsonRpc2ReceivedMessage
-	signalCliApiConfig       *utils.SignalCliApiConfig
-	number                   string
-	receivedMessagesMutex    sync.Mutex
-	receivedResponsesMutex   sync.Mutex
-	address                  string
+// receiveSubscription tracks the state of a manual-mode signal-cli
+// subscribeReceive call: the subscription id assigned by signal-cli and
+// a refcount of websocket subscribers attached to that account.
+type receiveSubscription struct {
+	id       int64
+	refcount int
 }
 
-func NewJsonRpc2Client(signalCliApiConfig *utils.SignalCliApiConfig, number string) *JsonRpc2Client {
+type JsonRpc2Client struct {
+	conn                      net.Conn
+	receivedResponsesById     map[string]chan JsonRpc2MessageResponse
+	receivedMessagesChannels  map[string]chan JsonRpc2ReceivedMessage
+	receiveSubscriptions      map[string]*receiveSubscription
+	accountsByChannelUuid     map[string]string
+	signalCliApiConfig        *utils.SignalCliApiConfig
+	number                    string
+	jsonRpcReceiveMode        JsonRpcReceiveMode
+	receivedMessagesMutex     sync.Mutex
+	receivedResponsesMutex    sync.Mutex
+	receiveSubscriptionsMutex sync.Mutex
+	address                   string
+}
+
+func NewJsonRpc2Client(signalCliApiConfig *utils.SignalCliApiConfig, number string, jsonRpcReceiveMode JsonRpcReceiveMode) *JsonRpc2Client {
 	return &JsonRpc2Client{
 		signalCliApiConfig:       signalCliApiConfig,
 		number:                   number,
+		jsonRpcReceiveMode:       jsonRpcReceiveMode,
 		receivedResponsesById:    make(map[string]chan JsonRpc2MessageResponse),
 		receivedMessagesChannels: make(map[string]chan JsonRpc2ReceivedMessage),
+		receiveSubscriptions:     make(map[string]*receiveSubscription),
+		accountsByChannelUuid:    make(map[string]string),
 	}
 }
 
@@ -238,6 +254,21 @@ func (r *JsonRpc2Client) ReceiveData(number string, receiveWebhookUrl string) {
 		var resp1 JsonRpc2ReceivedMessage
 		json.Unmarshal([]byte(str), &resp1)
 		if resp1.Method == "receive" {
+			// In manual receive-mode signal-cli wraps the envelope in
+			// {"subscription":N,"result":{...}}; in auto mode it sends
+			// the envelope directly. Unwrap so the broadcast format is
+			// the same in both modes and downstream consumers (e.g. the
+			// websocket handler) don't have to know which mode is in use.
+			if r.jsonRpcReceiveMode == Manual {
+				var manualWrapper struct {
+					Subscription int64           `json:"subscription"`
+					Result       json.RawMessage `json:"result"`
+				}
+				if err := json.Unmarshal(resp1.Params, &manualWrapper); err == nil && len(manualWrapper.Result) > 0 {
+					resp1.Params = manualWrapper.Result
+				}
+			}
+
 			r.receivedMessagesMutex.Lock()
 			for _, c := range r.receivedMessagesChannels {
 				select {
@@ -246,7 +277,6 @@ func (r *JsonRpc2Client) ReceiveData(number string, receiveWebhookUrl string) {
 				default:
 					log.Debug("Couldn't send message to golang channel, as there's no receiver")
 				}
-				continue
 			}
 			r.receivedMessagesMutex.Unlock()
 
@@ -275,16 +305,92 @@ func (r *JsonRpc2Client) ReceiveData(number string, receiveWebhookUrl string) {
 	}
 }
 
-func (r *JsonRpc2Client) GetReceiveChannel() (chan JsonRpc2ReceivedMessage, string, error) {
-	c := make(chan JsonRpc2ReceivedMessage)
+func (r *JsonRpc2Client) subscribeReceive(account string) (int64, error) {
+	type subscribeReceiveArgs struct{}
+	resultStr, err := r.getRaw("subscribeReceive", &account, subscribeReceiveArgs{})
+	if err != nil {
+		return 0, err
+	}
+	var subscriptionId int64
+	if err := json.Unmarshal([]byte(resultStr), &subscriptionId); err != nil {
+		return 0, err
+	}
+	return subscriptionId, nil
+}
+
+func (r *JsonRpc2Client) unsubscribeReceive(account string, subscriptionId int64) error {
+	type unsubscribeReceiveArgs struct {
+		Subscription int64 `json:"subscription"`
+	}
+	_, err := r.getRaw("unsubscribeReceive", &account, unsubscribeReceiveArgs{Subscription: subscriptionId})
+	return err
+}
+
+func (r *JsonRpc2Client) acquireReceiveSubscription(account string) error {
+	r.receiveSubscriptionsMutex.Lock()
+	defer r.receiveSubscriptionsMutex.Unlock()
+
+	if sub, ok := r.receiveSubscriptions[account]; ok {
+		sub.refcount++
+		return nil
+	}
+
+	id, err := r.subscribeReceive(account)
+	if err != nil {
+		return err
+	}
+
+	r.receiveSubscriptions[account] = &receiveSubscription{id: id, refcount: 1}
+	log.Debugf("Subscribed to receive notifications for account %s (subscription=%d)", account, id)
+	return nil
+}
+
+func (r *JsonRpc2Client) releaseReceiveSubscription(account string) error {
+	r.receiveSubscriptionsMutex.Lock()
+
+	sub, ok := r.receiveSubscriptions[account]
+	if !ok {
+		r.receiveSubscriptionsMutex.Unlock()
+		return nil
+	}
+
+	sub.refcount--
+	if sub.refcount > 0 {
+		r.receiveSubscriptionsMutex.Unlock()
+		return nil
+	}
+
+	delete(r.receiveSubscriptions, account)
+	err := r.unsubscribeReceive(account, sub.id)
+
+	r.receiveSubscriptionsMutex.Unlock()
+
+	if err != nil {
+		log.Warnf("Couldn't unsubscribed from receive notifications for account %s (subscription=%d): %s", account, sub.id, err.Error())
+	} else {
+		log.Debugf("Unsubscribed from receive notifications for account %s (subscription=%d)", account, sub.id)
+	}
+
+	return err
+}
+
+func (r *JsonRpc2Client) GetReceiveChannel(account string) (chan JsonRpc2ReceivedMessage, string, error) {
+	c := make(chan JsonRpc2ReceivedMessage, 64)
 
 	channelUuid, err := uuid.NewV4()
 	if err != nil {
 		return c, "", err
 	}
 
+	if r.jsonRpcReceiveMode == Manual {
+		if err := r.acquireReceiveSubscription(account); err != nil {
+			return c, "", fmt.Errorf("Couldn't create subscription for account %s: %w", account, err)
+		}
+	}
+
 	r.receivedMessagesMutex.Lock()
 	r.receivedMessagesChannels[channelUuid.String()] = c
+	r.accountsByChannelUuid[channelUuid.String()] = account
 	r.receivedMessagesMutex.Unlock()
 
 	return c, channelUuid.String(), nil
@@ -293,5 +399,13 @@ func (r *JsonRpc2Client) GetReceiveChannel() (chan JsonRpc2ReceivedMessage, stri
 func (r *JsonRpc2Client) RemoveReceiveChannel(channelUuid string) {
 	r.receivedMessagesMutex.Lock()
 	delete(r.receivedMessagesChannels, channelUuid)
+	account := r.accountsByChannelUuid[channelUuid]
+	delete(r.accountsByChannelUuid, channelUuid)
 	r.receivedMessagesMutex.Unlock()
+
+	if r.jsonRpcReceiveMode == Manual {
+		if err := r.releaseReceiveSubscription(account); err != nil {
+			log.Errorf("Couldn't release subscription for account %s: %s", account, err.Error())
+		}
+	}
 }

src/main.go

@@ -156,6 +156,18 @@ func main() {
 		}
 	}
 
+	jsonRpcReceiveMode := client.OnStart
+	jsonRpcReceiveModeValue := utils.GetEnv("JSON_RPC_RECEIVE_MODE", "")
+	if jsonRpcReceiveModeValue != "" && signalCliMode != client.JsonRpc {
+		log.Fatal("Env variable JSON_RPC_RECEIVE_MODE can't be used with normal or native mode")
+	} else if jsonRpcReceiveModeValue == "on-start" {
+		jsonRpcReceiveMode = client.OnStart
+	} else if jsonRpcReceiveModeValue == "manual" {
+		jsonRpcReceiveMode = client.Manual
+	} else if jsonRpcReceiveModeValue != "" {
+		log.Fatal("Invalid JSON_RPC_RECEIVE_MODE environment variable set!")
+	}
+
 	webhookUrl := utils.GetEnv("RECEIVE_WEBHOOK_URL", "")
 	if webhookUrl != "" && signalCliMode != client.JsonRpc {
 		log.Fatal("Env variable RECEIVE_WEBHOOK_URL can only be used with mode json-rpc!")
@@ -163,7 +175,7 @@ func main() {
 
 	jsonRpc2ClientConfigPathPath := *signalCliConfig + "/jsonrpc2.yml"
 	signalCliApiConfigPath := *signalCliConfig + "/api-config.yml"
-	signalClient := client.NewSignalClient(*signalCliConfig, *attachmentTmpDir, *avatarTmpDir, signalCliMode, jsonRpc2ClientConfigPathPath, signalCliApiConfigPath, webhookUrl)
+	signalClient := client.NewSignalClient(*signalCliConfig, *attachmentTmpDir, *avatarTmpDir, signalCliMode, jsonRpc2ClientConfigPathPath, signalCliApiConfigPath, webhookUrl, jsonRpcReceiveMode)
 	err = signalClient.Init(60)
 	if err != nil {
 		log.Fatal("Couldn't init Signal Client: ", err.Error())

src/scripts/jsonrpc2-helper.go

@@ -68,6 +68,15 @@ func main() {
 		args = append(args, "--ignore-attachments")
 	}
 
+	receiveMode := utils.GetEnv("JSON_RPC_RECEIVE_MODE", "")
+	if receiveMode == "on-start" {
+		args = append(args, "--receive-mode=on-start")
+	} else if receiveMode == "manual" {
+		args = append(args, "--receive-mode=manual")
+	} else if receiveMode != "" {
+		log.Fatal("Invalid JSON_RPC_RECEIVE_MODE environment variable set!")
+	}
+
 	args = append(args, []string{"--tcp", "127.0.0.1:" + strconv.FormatInt(tcpPort, 10)}...)
 
 	// write jsonrpc.yml config file