add allowCustomTags

algorand · Sep 8, 2022 · 06f5c56 · 06f5c56
1 parent 9df1d0c
commit 06f5c56
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/network/wsNetwork_test.go b/network/wsNetwork_test.go
@@ -51,6 +51,10 @@ import (
 
 const sendBufferLength = 1000
 
+func init() {
+	allowCustomTags = true
+}
+
 func TestMain(m *testing.M) {
 	logging.Base().SetLevel(logging.Debug)
 	os.Exit(m.Run())

diff --git a/network/wsPeer.go b/network/wsPeer.go
@@ -49,6 +49,9 @@ const msgsInReadBufferPerPeer = 10
 
 var tagStringList []string
 
+// allowCustomTags is set by tests to allow non-protocol-defined message tags.
+var allowCustomTags bool
+
 func init() {
 	tagStringList = make([]string, len(protocol.TagList))
 	for i, t := range protocol.TagList {
@@ -500,6 +503,9 @@ func (wp *wsPeer) readLoop() {
 		case protocol.VoteBundleTag:
 		default: // unrecognized tag
 			wp.OutOfProtocol = true
+			if !allowCustomTags {
+				continue // drop message, skip adding it to queue
+			}
 		}
 		if len(msg.Data) > 0 && wp.incomingMsgFilter != nil && dedupSafeTag(msg.Tag) {
 			if wp.incomingMsgFilter.CheckIncomingMessage(msg.Tag, msg.Data, true, true) {