fix: use io.Copy instead of ReadLine for http streaming

Current behavior may cause unintended response stream since it reads line and concatenate them with newline. This commit fixes the issue by using `io.Copy` to stream response body instead of using `ReadLine`. Signed-off-by: Sunghoon Kang <[email protected]>
akuity · Sep 12, 2024 · 0044f11 · 0044f11
1 parent e933104
commit 0044f11
Show file tree

Hide file tree

Showing 10 changed files with 532 additions and 96 deletions.
diff --git a/api/proto/testv1/test.proto b/api/proto/testv1/test.proto
@@ -27,6 +27,11 @@ service TestService {
       get: "/download-invitations"
     };
   }
+  rpc DownloadLargeFile(DownloadLargeFileRequest) returns (stream google.api.HttpBody) {
+    option (google.api.http) = {
+      get: "/download-large-file"
+    };
+  }
 }
 
 message InvitationMetadata {
@@ -78,3 +83,7 @@ message TrackInvitationResponse {
 message DownloadInvitationsRequest {
   optional EventType type = 1;
 }
+
+message DownloadLargeFileRequest {
+  // Explicitly empty
+}
diff --git a/internal/assets/assets.go b/internal/assets/assets.go
@@ -0,0 +1,8 @@
+package assets
+
+import _ "embed"
+
+var (
+	//go:embed large_file.txt
+	LargeFile string
+)
diff --git a/internal/assets/large_file.txt b/internal/assets/large_file.txt
diff --git a/internal/test/gen/testv1/test.gw.client.go b/internal/test/gen/testv1/test.gw.client.go
diff --git a/internal/test/gen/testv1/test.pb.go b/internal/test/gen/testv1/test.pb.go
diff --git a/internal/test/gen/testv1/test.pb.gw.go b/internal/test/gen/testv1/test.pb.gw.go
diff --git a/internal/test/gen/testv1/test_grpc.pb.go b/internal/test/gen/testv1/test_grpc.pb.go
diff --git a/internal/test/server/server.go b/internal/test/server/server.go
@@ -1,11 +1,14 @@
 package server
 
 import (
+	_ "embed"
+
 	"bytes"
 	"context"
 	"encoding/base64"
 	"time"
 
+	"github.com/akuity/grpc-gateway-client/internal/assets"
 	"github.com/bufbuild/protoyaml-go"
 	"google.golang.org/genproto/googleapis/api/httpbody"
 
@@ -76,3 +79,13 @@ func (s *testServiceServer) DownloadInvitations(req *testv1.DownloadInvitationsR
 	}
 	return nil
 }
+
+func (s *testServiceServer) DownloadLargeFile(
+	req *testv1.DownloadLargeFileRequest,
+	srv testv1.TestService_DownloadLargeFileServer,
+) error {
+	return srv.Send(&httpbody.HttpBody{
+		ContentType: "text/plain",
+		Data:        []byte(assets.LargeFile),
+	})
+}
diff --git a/pkg/grpc/gateway/request.go b/pkg/grpc/gateway/request.go
@@ -1,7 +1,6 @@
 package gateway
 
 import (
-	"bufio"
 	"bytes"
 	"context"
 	"encoding/json"
@@ -164,31 +163,15 @@ func doHTTPStreamingRequest(ctx context.Context, c Client, req *resty.Request) (
 		contentType := res.Header().Get("Content-Type")
 		body := res.RawBody()
 		defer func() { _ = body.Close() }()
-		r := bufio.NewReader(body)
-		for {
-			var data bytes.Buffer
-			for {
-				line, isPrefix, err := r.ReadLine()
-				if err != nil {
-					if errors.Is(err, io.EOF) {
-						close(resCh)
-						return
-					}
-					errCh <- fmt.Errorf("read line: %w", err)
-					return
-				}
-				if _, err := data.Write(append(line, '\n')); err != nil {
-					errCh <- fmt.Errorf("write line: %w", err)
-					return
-				}
-				if !isPrefix {
-					break
-				}
-			}
-			resCh <- &httpbody.HttpBody{
-				ContentType: contentType,
-				Data:        data.Bytes(),
-			}
+
+		var data bytes.Buffer
+		if _, err := io.Copy(&data, body); err != nil {
+			errCh <- fmt.Errorf("copy body: %w", err)
+			return
+		}
+		resCh <- &httpbody.HttpBody{
+			ContentType: contentType,
+			Data:        data.Bytes(),
 		}
 	}()
 	return resCh, errCh, nil

diff --git a/pkg/grpc/gateway/request_test.go b/pkg/grpc/gateway/request_test.go
@@ -1,6 +1,8 @@
 package gateway_test
 
 import (
+	_ "embed"
+
 	"bytes"
 	"context"
 	"net"
@@ -10,8 +12,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/akuity/grpc-gateway-client/internal/assets"
 	"github.com/bufbuild/protoyaml-go"
 	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
+	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 	"google.golang.org/genproto/googleapis/api/httpbody"
 	"google.golang.org/grpc"
@@ -141,6 +145,32 @@ read:
 	}
 }
 
+func (s *RequestTestSuite) TestDownloadLargeFileRequest() {
+	ctx, cancel := context.WithTimeout(context.TODO(), 100*time.Millisecond)
+	defer cancel()
+	req := s.client.NewRequest(http.MethodGet, "/download-large-file")
+	resCh, errCh, err := gateway.DoStreamingRequest[httpbody.HttpBody](ctx, s.client, req)
+	s.Require().NoError(err)
+	var buf bytes.Buffer
+
+read:
+	for {
+		select {
+		case <-ctx.Done():
+			break read
+		case err := <-errCh:
+			s.Require().NoError(err)
+		case data, ok := <-resCh:
+			if !ok {
+				break read
+			}
+			buf.Write(data.GetData())
+		}
+	}
+
+	require.Equal(s.T(), strings.TrimSpace(assets.LargeFile), strings.TrimSpace(buf.String()))
+}
+
 func (s *RequestTestSuite) TestDownloadRequest_Error() {
 	ctx, cancel := context.WithTimeout(context.TODO(), 500*time.Millisecond)
 	defer cancel()