From 9797d0f3d7d61609f94fa67ed5002dc954ef3f97 Mon Sep 17 00:00:00 2001
From: Rob Findley <rfindley@google.com>
Date: Mon, 23 Jun 2025 16:49:18 +0000
Subject: [PATCH] internal/mcp: make parseEventID robust, and add a test

Address comments on CL 682555 about using sscanf for parseEventID,
adding a test for parseEventID and formatEventID.

Change-Id: Ie31c5431c33829eaf244a5927da9e07bf5c300b5
Reviewed-on: https://go-review.googlesource.com/c/tools/+/683336
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
Auto-Submit: Robert Findley <rfindley@google.com>
---
 mcp/streamable.go      | 21 +++++++++++++------
 mcp/streamable_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/mcp/streamable.go b/mcp/streamable.go
index 37b0b40..63211c2 100644
--- a/mcp/streamable.go
+++ b/mcp/streamable.go
@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"strconv"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -449,20 +450,28 @@ stream:
 // streamID and message index idx.
 //
 // See also [parseEventID].
-func formatEventID(id streamID, idx int) string {
-	return fmt.Sprintf("%d_%d", id, idx)
+func formatEventID(sid streamID, idx int) string {
+	return fmt.Sprintf("%d_%d", sid, idx)
 }
 
 // parseEventID parses a Last-Event-ID value into a logical stream id and
 // index.
 //
 // See also [formatEventID].
-func parseEventID(eventID string) (conn streamID, idx int, ok bool) {
-	_, err := fmt.Sscanf(eventID, "%d_%d", &conn, &idx)
-	if err != nil || conn < 0 || idx < 0 {
+func parseEventID(eventID string) (sid streamID, idx int, ok bool) {
+	parts := strings.Split(eventID, "_")
+	if len(parts) != 2 {
 		return 0, 0, false
 	}
-	return conn, idx, true
+	stream, err := strconv.ParseInt(parts[0], 10, 64)
+	if err != nil || stream < 0 {
+		return 0, 0, false
+	}
+	idx, err = strconv.Atoi(parts[1])
+	if err != nil || idx < 0 {
+		return 0, 0, false
+	}
+	return streamID(stream), idx, true
 }
 
 // Read implements the [Connection] interface.
diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go
index 1137d50..26ca928 100644
--- a/mcp/streamable_test.go
+++ b/mcp/streamable_test.go
@@ -538,3 +538,49 @@ func mustMarshal(t *testing.T, v any) json.RawMessage {
 	}
 	return data
 }
+
+func TestEventID(t *testing.T) {
+	tests := []struct {
+		sid streamID
+		idx int
+	}{
+		{0, 0},
+		{0, 1},
+		{1, 0},
+		{1, 1},
+		{1234, 5678},
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("%d_%d", test.sid, test.idx), func(t *testing.T) {
+			eventID := formatEventID(test.sid, test.idx)
+			gotSID, gotIdx, ok := parseEventID(eventID)
+			if !ok {
+				t.Fatalf("parseEventID(%q) failed, want ok", eventID)
+			}
+			if gotSID != test.sid || gotIdx != test.idx {
+				t.Errorf("parseEventID(%q) = %d, %d, want %d, %d", eventID, gotSID, gotIdx, test.sid, test.idx)
+			}
+		})
+	}
+
+	invalid := []string{
+		"",
+		"_",
+		"1_",
+		"_1",
+		"a_1",
+		"1_a",
+		"-1_1",
+		"1_-1",
+	}
+
+	for _, eventID := range invalid {
+		t.Run(fmt.Sprintf("invalid_%q", eventID), func(t *testing.T) {
+			if _, _, ok := parseEventID(eventID); ok {
+				t.Errorf("parseEventID(%q) succeeded, want failure", eventID)
+			}
+		})
+	}
+}
+