Avoid json Unmarshal on proto CancelJobRequest (#31178)

apache · May 3, 2024 · e64c359 · e64c359
1 parent f468d40
commit e64c359
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/sdks/go/pkg/beam/runners/prism/internal/web/web.go b/sdks/go/pkg/beam/runners/prism/internal/web/web.go
@@ -373,8 +373,12 @@ type jobCancelHandler struct {
 	Jobcli jobpb.JobServiceClient
 }
 
+type cancelJobRequest struct {
+	JobId string `json:"job_id"`
+}
+
 func (h *jobCancelHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	var cancelRequest *jobpb.CancelJobRequest
+	var cancelRequest *cancelJobRequest
 	if r.Method != http.MethodPost {
 		http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
 		return
@@ -395,7 +399,10 @@ func (h *jobCancelHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	resp, err := h.Jobcli.Cancel(r.Context(), cancelRequest)
+	// Forward JobId from POST body avoids direct json Unmarshall on composite types containing protobuf message types.
+	resp, err := h.Jobcli.Cancel(r.Context(), &jobpb.CancelJobRequest{
+		JobId: cancelRequest.JobId,
+	})
 	if err != nil {
 		statusCode := status.Code(err)
 		httpCode := http.StatusInternalServerError