(2.12) Atomic batch: reject unsupported headers

Signed-off-by: Maurice van Veen <[email protected]>

Author: MauriceVanVeen

server/errors.json

@@ -1750,10 +1750,10 @@
     "deprecates": ""
   },
   {
-    "constant": "JSAtomicPublishDuplicateErr",
+    "constant": "JSAtomicPublishUnsupportedHeaderBatchErr",
     "code": 400,
     "error_code": 10177,
-    "description": "atomic publish duplicates not allowed",
+    "description": "atomic publish unsupported header used: {header}",
     "comment": "",
     "help": "",
     "url": "",

server/jetstream_batching_test.go

@@ -194,7 +194,7 @@ func TestJetStreamAtomicBatchPublishDedupeNotAllowed(t *testing.T) {
 		require_NoError(t, err)
 		require_NoError(t, json.Unmarshal(rmsg.Data, &pubAck))
 		require_NotNil(t, pubAck.Error)
-		require_Error(t, pubAck.Error, NewJSAtomicPublishDuplicateError())
+		require_Error(t, pubAck.Error, NewJSAtomicPublishUnsupportedHeaderBatchError("Nats-Msg-Id"))
 	}
 
 	for _, storage := range []StorageType{FileStorage, MemoryStorage} {
@@ -399,3 +399,42 @@ func TestJetStreamAtomicBatchPublishCleanup(t *testing.T) {
 	t.Run("StepDown", func(t *testing.T) { test(t, StepDown) })
 	t.Run("Delete", func(t *testing.T) { test(t, Delete) })
 }
+
+func TestJetStreamAtomicBatchPublishDenyHeaders(t *testing.T) {
+	test := func(t *testing.T, replicas int) {
+		c := createJetStreamClusterExplicit(t, "R3S", 3)
+		defer c.shutdown()
+
+		nc, js := jsClientConnect(t, c.randomServer())
+		defer nc.Close()
+
+		cfg := &StreamConfig{
+			Name:               "TEST",
+			Subjects:           []string{"foo"},
+			Storage:            FileStorage,
+			AllowAtomicPublish: true,
+			Replicas:           replicas,
+		}
+		_, err := jsStreamCreate(t, nc, cfg)
+		require_NoError(t, err)
+
+		// We might support these headers later on, but for now error.
+		for key, value := range map[string]string{
+			"Nats-Msg-Id":                 "msgId",
+			"Nats-Expected-Last-Sequence": "0",
+			"Nats-Expected-Last-Msg-Id":   "msgId",
+		} {
+			t.Run(key, func(t *testing.T) {
+				m := nats.NewMsg("foo")
+				m.Header.Set("Nats-Batch-Id", "uuid")
+				m.Header.Set("Nats-Batch-Sequence", "1")
+				m.Header.Set("Nats-Batch-Commit", "1")
+				m.Header.Set(key, value)
+				_, err = js.PublishMsg(m)
+				require_Error(t, err, NewJSAtomicPublishUnsupportedHeaderBatchError(key))
+			})
+		}
+	}
+
+	t.Run("R3", func(t *testing.T) { test(t, 3) })
+}

server/jetstream_cluster.go

@@ -8233,6 +8233,21 @@ func (mset *stream) processClusteredInboundMsg(subject, reply string, hdr, msg [
 			sm    *StoreMsg
 			sz    int
 		)
+
+		errorOnUnsupported := func(seq uint64, header string) *ApiError {
+			apiErr := NewJSAtomicPublishUnsupportedHeaderBatchError(header)
+			// TODO(mvv): reset in-memory expected header maps
+			mset.clseq -= seq - 1
+			mset.clMu.Unlock()
+			cleanup()
+			batches.mu.Unlock()
+			if canRespond {
+				buf, _ := json.Marshal(&JSPubAckResponse{PubAck: &PubAck{Stream: name}, Error: apiErr})
+				outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, nil, buf, nil, 0))
+			}
+			return apiErr
+		}
+
 		for seq := uint64(1); seq <= batchSeq; seq++ {
 			if seq == batchSeq {
 				bsubj, bhdr, bmsg = subject, hdr, msg
@@ -8251,19 +8266,15 @@ func (mset *stream) processClusteredInboundMsg(subject, reply string, hdr, msg [
 				return apiErr
 			}
 
-			// Duplicates are fully rejected when using batch.
+			// Reject unsupported headers.
 			if msgId := getMsgId(bhdr); msgId != _EMPTY_ {
-				// TODO(mvv): reset in-memory expected header maps
-				mset.clseq -= seq - 1
-				mset.clMu.Unlock()
-				cleanup()
-				batches.mu.Unlock()
-				apiErr := NewJSAtomicPublishDuplicateError()
-				if canRespond {
-					buf, _ := json.Marshal(&JSPubAckResponse{PubAck: &PubAck{Stream: name}, Error: apiErr})
-					outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, nil, buf, nil, 0))
-				}
-				return apiErr
+				return errorOnUnsupported(seq, JSMsgId)
+			}
+			if _, ok = getExpectedLastSeq(hdr); ok {
+				return errorOnUnsupported(seq, JSExpectedLastSeq)
+			}
+			if getExpectedLastMsgId(hdr) != _EMPTY_ {
+				return errorOnUnsupported(seq, JSExpectedLastMsgId)
 			}
 
 			var apiErr *ApiError

server/jetstream_errors_generated.go

@@ -11,15 +11,15 @@ const (
 	// JSAtomicPublishDisabledErr atomic publish is disabled
 	JSAtomicPublishDisabledErr ErrorIdentifier = 10174
 
-	// JSAtomicPublishDuplicateErr atomic publish duplicates not allowed
-	JSAtomicPublishDuplicateErr ErrorIdentifier = 10177
-
 	// JSAtomicPublishIncompleteBatchErr atomic publish batch is incomplete
 	JSAtomicPublishIncompleteBatchErr ErrorIdentifier = 10176
 
 	// JSAtomicPublishMissingSeqErr atomic publish sequence is missing
 	JSAtomicPublishMissingSeqErr ErrorIdentifier = 10175
 
+	// JSAtomicPublishUnsupportedHeaderBatchErr atomic publish unsupported header used: {header}
+	JSAtomicPublishUnsupportedHeaderBatchErr ErrorIdentifier = 10177
+
 	// JSBadRequestErr bad request
 	JSBadRequestErr ErrorIdentifier = 10003
 
@@ -541,9 +541,9 @@ var (
 	ApiErrors = map[ErrorIdentifier]*ApiError{
 		JSAccountResourcesExceededErr:              {Code: 400, ErrCode: 10002, Description: "resource limits exceeded for account"},
 		JSAtomicPublishDisabledErr:                 {Code: 400, ErrCode: 10174, Description: "atomic publish is disabled"},
-		JSAtomicPublishDuplicateErr:                {Code: 400, ErrCode: 10177, Description: "atomic publish duplicates not allowed"},
 		JSAtomicPublishIncompleteBatchErr:          {Code: 400, ErrCode: 10176, Description: "atomic publish batch is incomplete"},
 		JSAtomicPublishMissingSeqErr:               {Code: 400, ErrCode: 10175, Description: "atomic publish sequence is missing"},
+		JSAtomicPublishUnsupportedHeaderBatchErr:   {Code: 400, ErrCode: 10177, Description: "atomic publish unsupported header used: {header}"},
 		JSBadRequestErr:                            {Code: 400, ErrCode: 10003, Description: "bad request"},
 		JSClusterIncompleteErr:                     {Code: 503, ErrCode: 10004, Description: "incomplete results"},
 		JSClusterNoPeersErrF:                       {Code: 400, ErrCode: 10005, Description: "{err}"},
@@ -761,34 +761,40 @@ func NewJSAtomicPublishDisabledError(opts ...ErrorOption) *ApiError {
 	return ApiErrors[JSAtomicPublishDisabledErr]
 }
 
-// NewJSAtomicPublishDuplicateError creates a new JSAtomicPublishDuplicateErr error: "atomic publish duplicates not allowed"
-func NewJSAtomicPublishDuplicateError(opts ...ErrorOption) *ApiError {
+// NewJSAtomicPublishIncompleteBatchError creates a new JSAtomicPublishIncompleteBatchErr error: "atomic publish batch is incomplete"
+func NewJSAtomicPublishIncompleteBatchError(opts ...ErrorOption) *ApiError {
 	eopts := parseOpts(opts)
 	if ae, ok := eopts.err.(*ApiError); ok {
 		return ae
 	}
 
-	return ApiErrors[JSAtomicPublishDuplicateErr]
+	return ApiErrors[JSAtomicPublishIncompleteBatchErr]
 }
 
-// NewJSAtomicPublishIncompleteBatchError creates a new JSAtomicPublishIncompleteBatchErr error: "atomic publish batch is incomplete"
-func NewJSAtomicPublishIncompleteBatchError(opts ...ErrorOption) *ApiError {
+// NewJSAtomicPublishMissingSeqError creates a new JSAtomicPublishMissingSeqErr error: "atomic publish sequence is missing"
+func NewJSAtomicPublishMissingSeqError(opts ...ErrorOption) *ApiError {
 	eopts := parseOpts(opts)
 	if ae, ok := eopts.err.(*ApiError); ok {
 		return ae
 	}
 
-	return ApiErrors[JSAtomicPublishIncompleteBatchErr]
+	return ApiErrors[JSAtomicPublishMissingSeqErr]
 }
 
-// NewJSAtomicPublishMissingSeqError creates a new JSAtomicPublishMissingSeqErr error: "atomic publish sequence is missing"
-func NewJSAtomicPublishMissingSeqError(opts ...ErrorOption) *ApiError {
+// NewJSAtomicPublishUnsupportedHeaderBatchError creates a new JSAtomicPublishUnsupportedHeaderBatchErr error: "atomic publish unsupported header used: {header}"
+func NewJSAtomicPublishUnsupportedHeaderBatchError(header interface{}, opts ...ErrorOption) *ApiError {
 	eopts := parseOpts(opts)
 	if ae, ok := eopts.err.(*ApiError); ok {
 		return ae
 	}
 
-	return ApiErrors[JSAtomicPublishMissingSeqErr]
+	e := ApiErrors[JSAtomicPublishUnsupportedHeaderBatchErr]
+	args := e.toReplacerArgs([]interface{}{"{header}", header})
+	return &ApiError{
+		Code:        e.Code,
+		ErrCode:     e.ErrCode,
+		Description: strings.NewReplacer(args...).Replace(e.Description),
+	}
 }
 
 // NewJSBadRequestError creates a new JSBadRequestErr error: "bad request"