fix consumer group rewrite

proxy/rewrite_groups.go

@@ -88,6 +88,7 @@ func (c *conn) handleJoinGroup(hdr protocol.RequestHeader, body []byte) error {
 	if err := c.roundTripTyped(req, resp, hdr.ClientID); err != nil {
 		return errors.Wrap(err, "handleJoinGroup")
 	}
+	rewrite.JoinGroupResponseOut(c.tenant.TopicPrefix, resp)
 	return c.writeTypedResponse(hdr, resp)
 }
 
@@ -103,6 +104,7 @@ func (c *conn) handleSyncGroup(hdr protocol.RequestHeader, body []byte) error {
 	if err := c.roundTripTyped(req, resp, hdr.ClientID); err != nil {
 		return errors.Wrap(err, "handleSyncGroup")
 	}
+	rewrite.SyncGroupResponseOut(c.tenant.TopicPrefix, resp)
 	return c.writeTypedResponse(hdr, resp)
 }
 

rewrite/groups.go

@@ -102,14 +102,91 @@ func OffsetDeleteResponseOut(prefix string, resp *kmsg.OffsetDeleteResponse) {
 	}
 }
 
-// JoinGroupRequestIn prefixes the group ID.
+// consumerProtocolType is the only ProtocolType whose JoinGroup/SyncGroup
+// member blobs we decode; others (e.g. Kafka Connect) pass through untouched.
+const consumerProtocolType = "consumer"
+
+// rewriteConsumerSubscription applies fn to every topic in a ConsumerMemberMetadata
+// blob (JoinGroup protocol metadata). A blob that doesn't decode is returned as-is.
+func rewriteConsumerSubscription(meta []byte, fn func(string) string) []byte {
+	if len(meta) == 0 {
+		return meta
+	}
+	var m kmsg.ConsumerMemberMetadata
+	if err := m.ReadFrom(meta); err != nil {
+		return meta
+	}
+	for i := range m.Topics {
+		m.Topics[i] = fn(m.Topics[i])
+	}
+	for i := range m.OwnedPartitions {
+		m.OwnedPartitions[i].Topic = fn(m.OwnedPartitions[i].Topic)
+	}
+	return m.AppendTo(nil)
+}
+
+// rewriteConsumerAssignment applies fn to every topic in a ConsumerMemberAssignment
+// blob (SyncGroup member assignment). A blob that doesn't decode is returned as-is.
+func rewriteConsumerAssignment(asg []byte, fn func(string) string) []byte {
+	if len(asg) == 0 {
+		return asg
+	}
+	var a kmsg.ConsumerMemberAssignment
+	if err := a.ReadFrom(asg); err != nil {
+		return asg
+	}
+	for i := range a.Topics {
+		a.Topics[i].Topic = fn(a.Topics[i].Topic)
+	}
+	return a.AppendTo(nil)
+}
+
+// JoinGroupRequestIn prefixes the group ID and the subscribed topics in each
+// member's subscription blob.
 func JoinGroupRequestIn(prefix string, req *kmsg.JoinGroupRequest) {
 	req.Group = PrefixIn(prefix, req.Group)
+	if prefix == "" || req.ProtocolType != consumerProtocolType {
+		return
+	}
+	for i := range req.Protocols {
+		req.Protocols[i].Metadata = rewriteConsumerSubscription(req.Protocols[i].Metadata,
+			func(t string) string { return PrefixIn(prefix, t) })
+	}
+}
+
+// JoinGroupResponseOut strips the prefix from the subscription topics returned
+// to the leader, so it assigns in client space (matching what it sees via Metadata).
+func JoinGroupResponseOut(prefix string, resp *kmsg.JoinGroupResponse) {
+	if prefix == "" || (resp.ProtocolType != nil && *resp.ProtocolType != consumerProtocolType) {
+		return
+	}
+	for i := range resp.Members {
+		resp.Members[i].ProtocolMetadata = rewriteConsumerSubscription(resp.Members[i].ProtocolMetadata,
+			func(t string) string { return StripOut(prefix, t) })
+	}
 }
 
-// SyncGroupRequestIn prefixes the group ID.
+// SyncGroupRequestIn prefixes the group ID and the topics in each assignment
+// blob the leader submits.
 func SyncGroupRequestIn(prefix string, req *kmsg.SyncGroupRequest) {
 	req.Group = PrefixIn(prefix, req.Group)
+	if prefix == "" || (req.ProtocolType != nil && *req.ProtocolType != consumerProtocolType) {
+		return
+	}
+	for i := range req.GroupAssignment {
+		req.GroupAssignment[i].MemberAssignment = rewriteConsumerAssignment(req.GroupAssignment[i].MemberAssignment,
+			func(t string) string { return PrefixIn(prefix, t) })
+	}
+}
+
+// SyncGroupResponseOut strips the prefix from the topics in the member's
+// assignment, so it fetches client-space names (re-prefixed on Fetch).
+func SyncGroupResponseOut(prefix string, resp *kmsg.SyncGroupResponse) {
+	if prefix == "" || (resp.ProtocolType != nil && *resp.ProtocolType != consumerProtocolType) {
+		return
+	}
+	resp.MemberAssignment = rewriteConsumerAssignment(resp.MemberAssignment,
+		func(t string) string { return StripOut(prefix, t) })
 }
 
 // HeartbeatRequestIn prefixes the group ID.
@@ -129,11 +206,21 @@ func DescribeGroupsRequestIn(prefix string, req *kmsg.DescribeGroupsRequest) {
 	}
 }
 
-// DescribeGroupsResponseOut strips the tenant prefix from each group ID in
-// the response.
+// DescribeGroupsResponseOut strips the tenant prefix from each group ID and
+// from the topics in every member's subscription and assignment blob.
 func DescribeGroupsResponseOut(prefix string, resp *kmsg.DescribeGroupsResponse) {
+	strip := func(t string) string { return StripOut(prefix, t) }
 	for i := range resp.Groups {
-		resp.Groups[i].Group = StripOut(prefix, resp.Groups[i].Group)
+		g := &resp.Groups[i]
+		g.Group = StripOut(prefix, g.Group)
+		if prefix == "" || (g.ProtocolType != "" && g.ProtocolType != consumerProtocolType) {
+			continue
+		}
+		for j := range g.Members {
+			m := &g.Members[j]
+			m.ProtocolMetadata = rewriteConsumerSubscription(m.ProtocolMetadata, strip)
+			m.MemberAssignment = rewriteConsumerAssignment(m.MemberAssignment, strip)
+		}
 	}
 }
 

rewrite/groups_test.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/bubunyo/kroxy/rewrite"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"github.com/twmb/franz-go/pkg/kmsg"
 )
 
@@ -155,6 +156,27 @@ func TestDescribeGroups_RoundTrip(t *testing.T) {
 	assert.Equal(t, "b", resp.Groups[1].Group)
 }
 
+func TestDescribeGroups_StripsMemberBlobs(t *testing.T) {
+	t.Parallel()
+	resp := &kmsg.DescribeGroupsResponse{
+		Groups: []kmsg.DescribeGroupsResponseGroup{{
+			Group:        "tA.consumer-1",
+			ProtocolType: "consumer",
+			Members: []kmsg.DescribeGroupsResponseGroupMember{{
+				MemberID:         "m1",
+				ProtocolMetadata: encSub([]string{"tA.orders", "tA.events"}, nil),
+				MemberAssignment: encAsg("tA.orders", []int32{0, 1}),
+			}},
+		}},
+	}
+	rewrite.DescribeGroupsResponseOut(tp, resp)
+
+	assert.Equal(t, "consumer-1", resp.Groups[0].Group)
+	m := resp.Groups[0].Members[0]
+	assert.Equal(t, []string{"orders", "events"}, decSub(t, m.ProtocolMetadata).Topics)
+	assert.Equal(t, "orders", decAsg(t, m.MemberAssignment).Topics[0].Topic)
+}
+
 func TestListGroups_FiltersOtherTenants(t *testing.T) {
 	t.Parallel()
 	resp := &kmsg.ListGroupsResponse{
@@ -245,3 +267,125 @@ func TestTransactions_PrefixesTxnAndGroup(t *testing.T) {
 	rewrite.TxnOffsetCommitResponseOut(tp, tcr)
 	assert.Equal(t, "orders", tcr.Topics[0].Topic)
 }
+
+func encSub(topics []string, owned map[string][]int32) []byte {
+	m := kmsg.NewConsumerMemberMetadata()
+	m.Version = 1
+	m.Topics = topics
+	m.UserData = []byte("ud")
+	for tname, parts := range owned {
+		op := kmsg.NewConsumerMemberMetadataOwnedPartition()
+		op.Topic = tname
+		op.Partitions = parts
+		m.OwnedPartitions = append(m.OwnedPartitions, op)
+	}
+	return m.AppendTo(nil)
+}
+
+func decSub(t *testing.T, b []byte) kmsg.ConsumerMemberMetadata {
+	t.Helper()
+	var m kmsg.ConsumerMemberMetadata
+	require.NoError(t, m.ReadFrom(b))
+	return m
+}
+
+func encAsg(topic string, parts []int32) []byte {
+	a := kmsg.NewConsumerMemberAssignment()
+	a.Version = 1
+	at := kmsg.NewConsumerMemberAssignmentTopic()
+	at.Topic = topic
+	at.Partitions = parts
+	a.Topics = append(a.Topics, at)
+	a.UserData = []byte("ud")
+	return a.AppendTo(nil)
+}
+
+func decAsg(t *testing.T, b []byte) kmsg.ConsumerMemberAssignment {
+	t.Helper()
+	var a kmsg.ConsumerMemberAssignment
+	require.NoError(t, a.ReadFrom(b))
+	return a
+}
+
+func TestJoinGroup_RewritesSubscriptionTopics(t *testing.T) {
+	t.Parallel()
+
+	req := &kmsg.JoinGroupRequest{
+		Group:        "consumer-1",
+		ProtocolType: "consumer",
+		Protocols: []kmsg.JoinGroupRequestProtocol{
+			{Name: "range", Metadata: encSub([]string{"orders", "events"}, map[string][]int32{"orders": {0, 1}})},
+		},
+	}
+	rewrite.JoinGroupRequestIn(tp, req)
+	assert.Equal(t, "tA.consumer-1", req.Group)
+
+	got := decSub(t, req.Protocols[0].Metadata)
+	assert.Equal(t, []string{"tA.orders", "tA.events"}, got.Topics)
+	require.Len(t, got.OwnedPartitions, 1)
+	assert.Equal(t, "tA.orders", got.OwnedPartitions[0].Topic)
+	assert.Equal(t, []int32{0, 1}, got.OwnedPartitions[0].Partitions)
+	assert.Equal(t, []byte("ud"), got.UserData)
+
+	pt := "consumer"
+	resp := &kmsg.JoinGroupResponse{
+		ProtocolType: &pt,
+		Members: []kmsg.JoinGroupResponseMember{
+			{MemberID: "m1", ProtocolMetadata: encSub([]string{"tA.orders", "tA.events"}, nil)},
+		},
+	}
+	rewrite.JoinGroupResponseOut(tp, resp)
+	gotResp := decSub(t, resp.Members[0].ProtocolMetadata)
+	assert.Equal(t, []string{"orders", "events"}, gotResp.Topics)
+}
+
+func TestSyncGroup_RewritesAssignmentTopics(t *testing.T) {
+	t.Parallel()
+
+	req := &kmsg.SyncGroupRequest{
+		Group: "consumer-1",
+		GroupAssignment: []kmsg.SyncGroupRequestGroupAssignment{
+			{MemberID: "m1", MemberAssignment: encAsg("orders", []int32{0, 1})},
+		},
+	}
+	rewrite.SyncGroupRequestIn(tp, req)
+	assert.Equal(t, "tA.consumer-1", req.Group)
+	gotReq := decAsg(t, req.GroupAssignment[0].MemberAssignment)
+	require.Len(t, gotReq.Topics, 1)
+	assert.Equal(t, "tA.orders", gotReq.Topics[0].Topic)
+	assert.Equal(t, []int32{0, 1}, gotReq.Topics[0].Partitions)
+
+	resp := &kmsg.SyncGroupResponse{MemberAssignment: encAsg("tA.orders", []int32{0, 1})}
+	rewrite.SyncGroupResponseOut(tp, resp)
+	gotResp := decAsg(t, resp.MemberAssignment)
+	assert.Equal(t, "orders", gotResp.Topics[0].Topic)
+	assert.Equal(t, []byte("ud"), gotResp.UserData)
+}
+
+func TestJoinGroup_NonConsumerProtocolUntouched(t *testing.T) {
+	t.Parallel()
+
+	blob := []byte("opaque-connect-payload")
+	req := &kmsg.JoinGroupRequest{
+		Group:        "connect-cluster",
+		ProtocolType: "connect",
+		Protocols:    []kmsg.JoinGroupRequestProtocol{{Name: "default", Metadata: blob}},
+	}
+	rewrite.JoinGroupRequestIn(tp, req)
+	assert.Equal(t, "tA.connect-cluster", req.Group)
+	assert.Equal(t, blob, req.Protocols[0].Metadata, "non-consumer protocol metadata must pass through unchanged")
+}
+
+func TestJoinGroup_EmptyPrefixLeavesBlobsByteIdentical(t *testing.T) {
+	t.Parallel()
+
+	orig := encSub([]string{"orders"}, nil)
+	cp := append([]byte(nil), orig...)
+	req := &kmsg.JoinGroupRequest{
+		Group:        "consumer-1",
+		ProtocolType: "consumer",
+		Protocols:    []kmsg.JoinGroupRequestProtocol{{Name: "range", Metadata: cp}},
+	}
+	rewrite.JoinGroupRequestIn("", req)
+	assert.Equal(t, orig, req.Protocols[0].Metadata)
+}