feat: support quoted message headers

2026-09-23 17:46:51 +08:00
parent aa9f562e1e
commit 3e92021254
5 changed files with 235 additions and 1 deletions
@@ -6312,7 +6312,8 @@ func TestServeIMMessageSearchContextAndMedia(t *testing.T) {
 		store.Messages = oldMessages
 	}()
 	subs.EXPECT().Get(types.TenantID(7), "grpAlpha", types.Uid(42), false).Return(&types.Subscription{TenantID: 7, User: types.Uid(42).String(), Topic: "grpAlpha", ModeWant: types.ModeCPublic, ModeGiven: types.ModeCPublic}, nil).Times(3)
-	textMsg := types.Message{TenantID: 7, SeqId: 3, Topic: "grpAlpha", From: types.Uid(42).String(), Head: types.KVMap{"type": "text"}, Content: map[string]any{"text": "hello world"}, ObjHeader: types.ObjHeader{CreatedAt: time.Date(2026, 8, 13, 9, 0, 0, 0, time.UTC)}}
+	replyHead := map[string]any{"seq": float64(2), "type": "text", "text": "quoted hello", "from": types.Uid(43).UserId()}
+	textMsg := types.Message{TenantID: 7, SeqId: 3, Topic: "grpAlpha", From: types.Uid(42).String(), Head: types.KVMap{"type": "text", "x-im-reply": replyHead}, Content: map[string]any{"text": "hello world"}, ObjHeader: types.ObjHeader{CreatedAt: time.Date(2026, 8, 13, 9, 0, 0, 0, time.UTC)}}
 	fileMsg := types.Message{TenantID: 7, SeqId: 4, Topic: "grpAlpha", From: types.Uid(42).String(), Head: types.KVMap{"type": "image"}, Content: map[string]any{"file_id": "file-1", "mime": "image/png"}, ObjHeader: types.ObjHeader{CreatedAt: time.Date(2026, 8, 13, 9, 1, 0, 0, time.UTC)}}
 	messages.EXPECT().Search(types.TenantID(7), "grpAlpha", types.Uid(42), "hello", "text", "", 50).Return([]types.Message{textMsg}, "", nil)
 	messages.EXPECT().GetAll(types.TenantID(7), "grpAlpha", types.Uid(42), gomock.Any()).Return([]types.Message{textMsg}, nil)
@@ -6334,6 +6335,10 @@ func TestServeIMMessageSearchContextAndMedia(t *testing.T) {
 	if searchSnapshot["from"] != types.Uid(42).UserId() {
 		t.Fatalf("search snapshot from must use full user_id: %#v", searchSnapshot)
 	}
+	searchHead := searchSnapshot["head"].(map[string]any)
+	if got := searchHead["x-im-reply"].(map[string]any); got["seq"] != replyHead["seq"] || got["type"] != replyHead["type"] || got["text"] != replyHead["text"] || got["from"] != replyHead["from"] {
+		t.Fatalf("search snapshot must preserve x-im-reply: %#v", got)
+	}

 	contextReq := httptest.NewRequest(http.MethodGet, "/v1/im/topics/grpAlpha/messages/3/context", nil)
 	contextReq.Header.Set("Authorization", "Bearer token-1")
@@ -6347,6 +6352,10 @@ func TestServeIMMessageSearchContextAndMedia(t *testing.T) {
 	if len(contextItems) != 1 || contextItems[0].(map[string]any)["from"] != types.Uid(42).UserId() {
 		t.Fatalf("context from must use full user_id: %#v", contextData)
 	}
+	contextHead := contextItems[0].(map[string]any)["head"].(map[string]any)
+	if got := contextHead["x-im-reply"].(map[string]any); got["seq"] != replyHead["seq"] || got["type"] != replyHead["type"] || got["text"] != replyHead["text"] || got["from"] != replyHead["from"] {
+		t.Fatalf("context must preserve x-im-reply: %#v", got)
+	}

 	mediaReq := httptest.NewRequest(http.MethodGet, "/v1/im/topics/grpAlpha/media?type=image", nil)
 	mediaReq.Header.Set("Authorization", "Bearer token-1")
@@ -3727,6 +3727,13 @@ func validateMessageHeaders(head types.KVMap) error {
 			return types.ErrMalformed
 		}
 	}
+	if rawReply, ok := head["x-im-reply"]; ok {
+		reply, err := normalizeIMReplyHeader(rawReply)
+		if err != nil {
+			return err
+		}
+		head["x-im-reply"] = reply
+	}
 	return nil
 }

@@ -3779,6 +3786,74 @@ func messageHeaderString(head types.KVMap, key string) (string, bool) {
 	return text, text != ""
 }

+func normalizeIMReplyHeader(value any) (map[string]any, error) {
+	var reply map[string]any
+	switch typed := value.(type) {
+	case map[string]any:
+		reply = typed
+	case types.KVMap:
+		reply = map[string]any(typed)
+	case string:
+		if strings.TrimSpace(typed) == "" {
+			return nil, types.ErrMalformed
+		}
+		if err := json.Unmarshal([]byte(typed), &reply); err != nil {
+			return nil, types.ErrMalformed
+		}
+	default:
+		return nil, types.ErrMalformed
+	}
+	if reply == nil || !validIMReplySeq(reply["seq"]) {
+		return nil, types.ErrMalformed
+	}
+	if text, ok := reply["type"].(string); !ok || strings.TrimSpace(text) == "" {
+		return nil, types.ErrMalformed
+	}
+	if text, ok := reply["text"].(string); !ok || strings.TrimSpace(text) == "" {
+		return nil, types.ErrMalformed
+	}
+	if from, ok := reply["from"]; ok {
+		if text, ok := from.(string); !ok || strings.TrimSpace(text) == "" {
+			return nil, types.ErrMalformed
+		}
+	}
+	return reply, nil
+}
+
+func validIMReplySeq(value any) bool {
+	switch typed := value.(type) {
+	case json.Number:
+		_, err := typed.Int64()
+		return err == nil
+	case float64:
+		return typed > 0 && typed == float64(int64(typed))
+	case float32:
+		return typed > 0 && typed == float32(int64(typed))
+	case int:
+		return typed > 0
+	case int8:
+		return typed > 0
+	case int16:
+		return typed > 0
+	case int32:
+		return typed > 0
+	case int64:
+		return typed > 0
+	case uint:
+		return typed > 0
+	case uint8:
+		return typed > 0
+	case uint16:
+		return typed > 0
+	case uint32:
+		return typed > 0
+	case uint64:
+		return typed > 0
+	default:
+		return false
+	}
+}
+
 func validIMMessageHeaderToken(value string, max int) bool {
 	value = strings.TrimSpace(value)
 	if value == "" || len(value) > max {
@@ -23,6 +23,8 @@ func TestValidateMessageHeadersProductContract(t *testing.T) {
 		{"x-im-schema": 1, "type": "audio", "x-im-type": "voice"},
 		{"x-im-schema": 1, "x-im-type": "system_notice"},
 		{"x-im-schema": 1, "x-im-type": "system_action"},
+		{"x-im-schema": 1, "x-im-type": "text", "x-im-reply": map[string]any{"seq": float64(128), "type": "text", "text": "hello", "from": "usr1"}},
+		{"x-im-schema": 1, "x-im-type": "text", "x-im-reply": `{"seq":128,"type":"text","text":"hello","from":"usr1"}`},
 		{"type": "system"},
 	}
 	for _, head := range valid {
@@ -36,6 +38,10 @@ func TestValidateMessageHeadersProductContract(t *testing.T) {
 		{"x-im-schema": true, "x-im-type": "text"},
 		{"x-im-schema": float64(1), "x-im-type": "unknown"},
 		{"type": "text", "x-im-type": "image"},
+		{"x-im-schema": 1, "x-im-type": "text", "x-im-reply": map[string]any{"seq": 0, "type": "text", "text": "hello"}},
+		{"x-im-schema": 1, "x-im-type": "text", "x-im-reply": map[string]any{"seq": 1.5, "type": "text", "text": "hello"}},
+		{"x-im-schema": 1, "x-im-type": "text", "x-im-reply": map[string]any{"seq": 1, "type": "text"}},
+		{"x-im-schema": 1, "x-im-type": "text", "x-im-reply": "not-json"},
 	}
 	for _, head := range invalid {
 		if err := validateMessageHeaders(head); err != types.ErrMalformed {
@@ -145,3 +151,21 @@ func TestActionTokenCachePreservesInternalFields(t *testing.T) {
 		t.Fatalf("cache round trip lost fields: got %#v", got)
 	}
 }
+
+func TestValidateMessageHeadersNormalizesReplyJSONString(t *testing.T) {
+	head := types.KVMap{
+		"x-im-schema": 1,
+		"x-im-type":   "text",
+		"x-im-reply":  `{"seq":128,"type":"text","text":"hello","from":"usr1"}`,
+	}
+	if err := validateMessageHeaders(head); err != nil {
+		t.Fatalf("validateMessageHeaders() = %v, want nil", err)
+	}
+	reply, ok := head["x-im-reply"].(map[string]any)
+	if !ok {
+		t.Fatalf("x-im-reply was not normalized to object: %#v", head["x-im-reply"])
+	}
+	if reply["seq"].(float64) != 128 || reply["type"] != "text" || reply["text"] != "hello" || reply["from"] != "usr1" {
+		t.Fatalf("unexpected normalized reply: %#v", reply)
+	}
+}
@@ -302,6 +302,70 @@ func TestHandleBroadcastDataP2P(t *testing.T) {
 	}
 }

+func TestHandleBroadcastDataPreservesReplyHead(t *testing.T) {
+	numUsers := 2
+	helper := TopicTestHelper{}
+	helper.setUp(t, numUsers, types.TopicCatP2P, "p2p-test" /*attach=*/, true)
+	defer helper.tearDown()
+
+	reply := map[string]any{
+		"seq":  int64(128),
+		"type": "text",
+		"text": "今天晚上一起吃饭吗？",
+		"from": helper.uids[1].UserId(),
+	}
+	helper.mm.EXPECT().Save(types.TenantID(1), gomock.Any(), gomock.Any(), gomock.Any()).
+		DoAndReturn(func(_ types.TenantID, msg *types.Message, _ []string, _ bool) (*types.MessageSaveResult, error) {
+			got, ok := msg.Head["x-im-reply"].(map[string]any)
+			if !ok {
+				t.Fatalf("stored x-im-reply: got %#v", msg.Head["x-im-reply"])
+			}
+			if got["seq"] != reply["seq"] || got["type"] != reply["type"] || got["text"] != reply["text"] || got["from"] != reply["from"] {
+				t.Fatalf("stored x-im-reply changed: got %#v want %#v", got, reply)
+			}
+			return &types.MessageSaveResult{SeqId: 1, MarkedReadBySender: true}, nil
+		})
+
+	from := helper.uids[0].UserId()
+	msg := &ClientComMessage{
+		AsUser:   from,
+		Original: from,
+		Pub: &MsgClientPub{
+			Topic: "p2p",
+			Head: map[string]any{
+				"mime":            "application/json",
+				"x-im-schema":     1,
+				"x-im-type":       "text",
+				"x-im-client-mid": "mid-reply-1",
+				"x-im-reply":      reply,
+			},
+			Content: map[string]any{"text": "好啊"},
+			NoEcho:  true,
+		},
+		sess: helper.sessions[0],
+	}
+	helper.topic.handleClientMsg(msg)
+	helper.finish()
+
+	if errorMsgs, hasError := helper.hubMessages["__ERROR__"]; hasError {
+		t.Fatal(errorMsgs[0].Ctrl.Text)
+	}
+	if len(helper.results[1].messages) != 1 {
+		t.Fatalf("peer: expected 1 message, got %d", len(helper.results[1].messages))
+	}
+	data := helper.results[1].messages[0].(*ServerComMessage).Data
+	if data == nil {
+		t.Fatalf("peer response must have data")
+	}
+	got, ok := data.Head["x-im-reply"].(map[string]any)
+	if !ok {
+		t.Fatalf("broadcast x-im-reply: got %#v", data.Head["x-im-reply"])
+	}
+	if got["seq"] != reply["seq"] || got["type"] != reply["type"] || got["text"] != reply["text"] || got["from"] != reply["from"] {
+		t.Fatalf("broadcast x-im-reply changed: got %#v want %#v", got, reply)
+	}
+}
+
 func TestHandleBroadcastDataP2PRequiresRelationship(t *testing.T) {
 	numUsers := 2
 	helper := TopicTestHelper{}
@@ -1130,6 +1194,65 @@ func TestSaveMessageMaybeKafkaPublishesWithoutDBPersist(t *testing.T) {
 	}
 }

+func TestSaveMessageMaybeKafkaPreservesReplyHead(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	oldKafkaBus := kafkaBus
+	oldStore := store.Store
+	oldMessages := store.Messages
+	oldPublisher := publishKafkaMessagePersist
+	defer func() {
+		kafkaBus = oldKafkaBus
+		store.Store = oldStore
+		store.Messages = oldMessages
+		publishKafkaMessagePersist = oldPublisher
+	}()
+
+	kafkaBus = &kafkaRuntime{}
+	mockStore := mock_store.NewMockPersistentStorageInterface(ctrl)
+	mockStore.EXPECT().GetUid().Return(types.Uid(102))
+	store.Store = mockStore
+	store.Messages = mock_store.NewMockMessagesPersistenceInterface(ctrl)
+	var publishedPayload kafkaMessagePersistPayload
+	publishKafkaMessagePersist = func(_ context.Context, _, _ string, envelope kafkaEnvelope) error {
+		if err := json.Unmarshal(envelope.Payload, &publishedPayload); err != nil {
+			t.Fatalf("failed to unmarshal kafka payload: %v", err)
+		}
+		return nil
+	}
+
+	reply := map[string]any{
+		"seq":  float64(12),
+		"type": "text",
+		"text": "quoted hello",
+		"from": types.Uid(2).UserId(),
+	}
+	msg := &types.Message{
+		TenantID: 1,
+		Topic:    "grp-test",
+		SeqId:    7,
+		From:     types.Uid(1).String(),
+		Head: types.KVMap{
+			"x-im-client-mid": "client-reply-1",
+			"x-im-type":       "text",
+			"x-im-reply":      reply,
+		},
+		Content: map[string]any{"text": "reply"},
+	}
+
+	if _, err := saveMessageMaybeKafka(context.Background(), 1, msg, nil, true); err != nil {
+		t.Fatalf("saveMessageMaybeKafka returned error: %v", err)
+	}
+	got, ok := publishedPayload.Message.Head["x-im-reply"].(map[string]any)
+	if !ok {
+		t.Fatalf("kafka payload x-im-reply: got %#v", publishedPayload.Message.Head["x-im-reply"])
+	}
+	if got["seq"] != reply["seq"] || got["type"] != reply["type"] || got["text"] != reply["text"] || got["from"] != reply["from"] {
+		t.Fatalf("kafka payload x-im-reply changed: got %#v want %#v", got, reply)
+	}
+}
+
 func TestShouldTryHotMessageFallback(t *testing.T) {
 	if !shouldTryHotMessageFallback(types.TopicCatGrp, 159, &types.QueryOpt{Since: 159, Limit: 1}) {
 		t.Fatal("expected fallback for latest group seq query")