From 9cf2949ef8a6d04ecbc8b4b9b9c428e732b98185 Mon Sep 17 00:00:00 2001
From: Xavrax <matedahlke@gmail.com>
Date: Mon, 9 Dec 2024 22:33:17 +0100
Subject: [PATCH] validation of message type

---
 publish_request.go      | 25 ++++++++++++++++++++++++-
 publish_request_test.go | 14 ++++++++++++++
 pubnub.go               |  2 ++
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/publish_request.go b/publish_request.go
index 4b4544a1..0bed78af 100644
--- a/publish_request.go
+++ b/publish_request.go
@@ -170,7 +170,8 @@ func (b *publishBuilder) QueryParam(queryParam map[string]string) *publishBuilde
 	return b
 }
 
-// CustomMessageType sets the Custom Message Type for the Publish request.
+// CustomMessageType sets the User-specified message type string - limited by 3-50 case-sensitive alphanumeric characters
+// with only `-` and `_` special characters allowed.
 func (b *publishBuilder) CustomMessageType(messageType string) *publishBuilder {
     b.opts.CustomMessageType = messageType
 
@@ -187,6 +188,24 @@ func (b *publishBuilder) Execute() (*PublishResponse, StatusResponse, error) {
 	return newPublishResponse(rawJSON, status)
 }
 
+func (o *publishOpts) isCustomMessageTypeCorrect() bool {
+    if len(o.CustomMessageType) == 0 {
+        return true
+    }
+
+    if len(o.CustomMessageType) < 3 || len(o.CustomMessageType) > 50 {
+        return false
+    }
+
+    for _, c := range o.CustomMessageType {
+        if !('a' <= c && 'z' >= c) && !('A' <= c && 'Z' >= c) && c != '-' && c != '_' {
+            return false
+        }
+    }
+
+    return true
+}
+
 func (o *publishOpts) validate() error {
 	if o.config().PublishKey == "" {
 		return newValidationError(o, StrMissingPubKey)
@@ -204,6 +223,10 @@ func (o *publishOpts) validate() error {
 		return newValidationError(o, StrMissingMessage)
 	}
 
+    if !o.isCustomMessageTypeCorrect() {
+        return newValidationError(o, StrInvalidCustomMessageType)
+    }
+
 	return nil
 }
 
diff --git a/publish_request_test.go b/publish_request_test.go
index 41661540..c4cebc6c 100644
--- a/publish_request_test.go
+++ b/publish_request_test.go
@@ -80,6 +80,7 @@ func AssertSuccessPublishGet2(t *testing.T, expectedString string, message inter
 	o.TTL(10)
 	o.ShouldStore(false)
 	o.DoNotReplicate(true)
+    o.CustomMessageType("custom")
 
 	path, err := o.opts.buildPath()
 	assert.Nil(err)
@@ -99,6 +100,7 @@ func AssertSuccessPublishGet2(t *testing.T, expectedString string, message inter
 	expected.Set("pnsdk", Version)
 	expected.Set("norep", "true")
 	expected.Set("store", "0")
+    expected.Set("custom_message_type", "custom")
 
 	h.AssertQueriesEqual(t, expected, query,
 		[]string{"seqn", "pnsdk", "uuid", "store"}, []string{})
@@ -532,3 +534,15 @@ func TestPublishValidateSubscribeKey(t *testing.T) {
 
 	assert.Equal("pubnub/validation: pubnub: Publish: Missing Subscribe Key", opts.validate().Error())
 }
+
+func TestPublishValidateCustomMessageType(t *testing.T) {
+    assert := assert.New(t)
+    pn := NewPubNub(NewDemoConfig())
+    opts := newPublishOpts(pn, pn.ctx)
+    opts.CustomMessageType = "custom-message_type"
+    assert.True(opts.isCustomMessageTypeCorrect())
+    opts.CustomMessageType = "a"
+    assert.False(opts.isCustomMessageTypeCorrect())
+    opts.CustomMessageType = "!@#$%^&*("
+    assert.False(opts.isCustomMessageTypeCorrect())
+}
diff --git a/pubnub.go b/pubnub.go
index 93c2a585..5ea96106 100644
--- a/pubnub.go
+++ b/pubnub.go
@@ -55,6 +55,8 @@ const (
 	StrMissingFileName = "Missing File Name"
 	// StrMissingToken shows `Missing PAMv3 token` message
 	StrMissingToken = "Missing PAMv3 token"
+    // StrInvalidCustomMessageType shows `Invalid CustomMessageType` message
+    StrInvalidCustomMessageType = "Invalid CustomMessageType: size different than 3-50 or contains invalid characters"
 )
 
 // PubNub No server connection will be established when you create a new PubNub object.