Update order of operation for request interceptor and serialization in gremlin-go (#3358)

* add proper error handling for non-graphbinary cases
* Update order of interceptor and serialization of request
* export request and rename to RequestMessage for clarity and public access like other GLVs
* added wg to ensure in-flight goroutines complete & added response body drain before close to prevent TCP RST errors

Author: xiazcy

CHANGELOG.asciidoc

@@ -41,8 +41,10 @@ image::https://raw.githubusercontent.com/apache/tinkerpop/master/docs/static/ima
 * Removed deprecated `Graph.traversal()` method in JS in favor of the anonymous `traversal()` function.
 * Replace `Bytecode` with `GremlinLang` & update serialization to GraphBinary 4 for `gremlin-go`.
 * Added `RequestInterceptor` to `gremlin-go` with `auth` reference implementations to replace `authInfo`.
-* Refactored GraphBinary serializers to use `io.Writer` and `io.Reader` instead of `*bytes.Buffer` for streaming capacities.
+* Refactored GraphBinary serializers in `gremlin-go` to use `io.Writer` and `io.Reader` instead of `*bytes.Buffer` for streaming capacities.
 * Refactored `httpProtocol` and `httpTransport` in `gremlin-go` into single `connection.go` that handles HTTP request and response.
+* Reordered interceptor chain in `gremlin-go` so interceptors access raw request before serialization.
+* Exported `request` in `gremlin-go` as `RequestMessage` with public `Gremlin`/`Fields` for clarity, access and consistency.
 * Refactored result handling in `gremlin-driver` by merging `ResultQueue` into `ResultSet`.
 * Replace `Bytecode` with `GremlinLang` in `gremlin-dotnet`.
 * Replace `WebSocket` with `HTTP` (non-streaming) in `gremlin-dotnet`.

gremlin-go/driver/auth.go

@@ -22,6 +22,7 @@ package gremlingo
 import (
 	"context"
 	"encoding/base64"
+	"fmt"
 	"sync"
 	"time"
 
@@ -39,26 +40,40 @@ func BasicAuth(username, password string) RequestInterceptor {
 	}
 }
 
-// Sigv4Auth returns a RequestInterceptor that signs requests using AWS SigV4.
+// SigV4Auth returns a RequestInterceptor that signs requests using AWS SigV4.
 // It uses the default AWS credential chain (env vars, shared config, IAM role, etc.)
-func Sigv4Auth(region, service string) RequestInterceptor {
-	return Sigv4AuthWithCredentials(region, service, nil)
+func SigV4Auth(region, service string) RequestInterceptor {
+	return SigV4AuthWithCredentials(region, service, nil)
 }
 
-// Sigv4AuthWithCredentials returns a RequestInterceptor that signs requests using AWS SigV4
+// SigV4AuthWithCredentials returns a RequestInterceptor that signs requests using AWS SigV4
 // with the provided credentials provider. If provider is nil, uses default credential chain.
+// If the request body has not been serialized yet (*RequestMessage), it is automatically
+// serialized to GraphBinary before signing.
 //
 // Caches the signer and credentials provider for efficiency.
-func Sigv4AuthWithCredentials(region, service string, credentialsProvider aws.CredentialsProvider) RequestInterceptor {
+func SigV4AuthWithCredentials(region, service string, credentialsProvider aws.CredentialsProvider) RequestInterceptor {
 	// Create signer once - it's stateless and safe to reuse
 	signer := v4.NewSigner()
+	serialize := SerializeRequest()
 
 	// Cache for resolved credentials provider (lazy initialization)
 	var cachedProvider aws.CredentialsProvider
 	var providerOnce sync.Once
 	var providerErr error
 
 	return func(req *HttpRequest) error {
+		// If Body is still *RequestMessage, serialize it to GraphBinary before signing.
+		if _, ok := req.Body.(*RequestMessage); ok {
+			if err := serialize(req); err != nil {
+				return fmt.Errorf("SigV4 auto-serialization failed: %w", err)
+			}
+		}
+
+		if _, ok := req.Body.([]byte); !ok {
+			return fmt.Errorf("SigV4 signing requires body to be []byte; got %T", req.Body)
+		}
+
 		ctx := context.Background()
 
 		// Resolve credentials provider once if not provided

gremlin-go/driver/auth_test.go

@@ -30,7 +30,7 @@ import (
 )
 
 func createMockRequest() *HttpRequest {
-	req, _ := NewHttpRequest("POST", "https://localhost:8182/gremlin")
+	req, _ := NewHttpRequest("POST", "https://test_url:8182/gremlin")
 	req.Headers.Set("Content-Type", graphBinaryMimeType)
 	req.Headers.Set("Accept", graphBinaryMimeType)
 	req.Body = []byte(`{"gremlin":"g.V()"}`)
@@ -72,24 +72,24 @@ func (m *mockCredentialsProvider) Retrieve(ctx context.Context) (aws.Credentials
 	}, nil
 }
 
-func TestSigv4Auth(t *testing.T) {
+func TestSigV4Auth(t *testing.T) {
 	t.Run("adds signed headers", func(t *testing.T) {
 		req := createMockRequest()
 		assert.Empty(t, req.Headers.Get("Authorization"))
 		assert.Empty(t, req.Headers.Get("X-Amz-Date"))
 
 		provider := &mockCredentialsProvider{
-			accessKey: "MOCK_ACCESS_KEY",
-			secretKey: "MOCK_SECRET_KEY",
+			accessKey: "MOCK_ID",
+			secretKey: "MOCK_KEY",
 		}
-		interceptor := Sigv4AuthWithCredentials("us-west-2", "neptune-db", provider)
+		interceptor := SigV4AuthWithCredentials("gremlin-east-1", "tinkerpop-sigv4", provider)
 		err := interceptor(req)
 
 		assert.NoError(t, err)
 		assert.NotEmpty(t, req.Headers.Get("X-Amz-Date"))
 		authHeader := req.Headers.Get("Authorization")
-		assert.True(t, strings.HasPrefix(authHeader, "AWS4-HMAC-SHA256 Credential=MOCK_ACCESS_KEY"))
-		assert.Contains(t, authHeader, "us-west-2/neptune-db/aws4_request")
+		assert.True(t, strings.HasPrefix(authHeader, "AWS4-HMAC-SHA256 Credential=MOCK_ID"))
+		assert.Contains(t, authHeader, "gremlin-east-1/tinkerpop-sigv4/aws4_request")
 		assert.Contains(t, authHeader, "Signature=")
 	})
 
@@ -98,17 +98,17 @@ func TestSigv4Auth(t *testing.T) {
 		assert.Empty(t, req.Headers.Get("X-Amz-Security-Token"))
 
 		provider := &mockCredentialsProvider{
-			accessKey:    "MOCK_ACCESS_KEY",
-			secretKey:    "MOCK_SECRET_KEY",
-			sessionToken: "MOCK_SESSION_TOKEN",
+			accessKey:    "MOCK_ID",
+			secretKey:    "MOCK_KEY",
+			sessionToken: "MOCK_TOKEN",
 		}
-		interceptor := Sigv4AuthWithCredentials("us-west-2", "neptune-db", provider)
+		interceptor := SigV4AuthWithCredentials("gremlin-east-1", "tinkerpop-sigv4", provider)
 		err := interceptor(req)
 
 		assert.NoError(t, err)
-		assert.Equal(t, "MOCK_SESSION_TOKEN", req.Headers.Get("X-Amz-Security-Token"))
+		assert.Equal(t, "MOCK_TOKEN", req.Headers.Get("X-Amz-Security-Token"))
 		authHeader := req.Headers.Get("Authorization")
 		assert.True(t, strings.HasPrefix(authHeader, "AWS4-HMAC-SHA256 Credential="))
-		assert.Contains(t, authHeader, "Signature=")
+		assert.Contains(t, authHeader, "gremlin-east-1/tinkerpop-sigv4/aws4_request")
 	})
 }

gremlin-go/driver/connection.go

@@ -22,69 +22,17 @@ package gremlingo
 import (
 	"bytes"
 	"compress/zlib"
-	"crypto/sha256"
 	"crypto/tls"
-	"encoding/hex"
+	"encoding/json"
+	"fmt"
 	"io"
 	"net"
 	"net/http"
-	"net/url"
+	"strings"
+	"sync"
 	"time"
 )
 
-// Common HTTP header keys
-const (
-	HeaderContentType    = "Content-Type"
-	HeaderAccept         = "Accept"
-	HeaderUserAgent      = "User-Agent"
-	HeaderAcceptEncoding = "Accept-Encoding"
-	HeaderAuthorization  = "Authorization"
-)
-
-// HttpRequest represents an HTTP request that can be modified by interceptors.
-type HttpRequest struct {
-	Method  string
-	URL     *url.URL
-	Headers http.Header
-	Body    []byte
-}
-
-// NewHttpRequest creates a new HttpRequest with the given method and URL.
-func NewHttpRequest(method, rawURL string) (*HttpRequest, error) {
-	u, err := url.Parse(rawURL)
-	if err != nil {
-		return nil, err
-	}
-	return &HttpRequest{
-		Method:  method,
-		URL:     u,
-		Headers: make(http.Header),
-	}, nil
-}
-
-// ToStdRequest converts HttpRequest to a standard http.Request for signing.
-// Returns nil if the request cannot be created (invalid method or URL).
-func (r *HttpRequest) ToStdRequest() (*http.Request, error) {
-	req, err := http.NewRequest(r.Method, r.URL.String(), bytes.NewReader(r.Body))
-	if err != nil {
-		return nil, err
-	}
-	req.Header = r.Headers
-	return req, nil
-}
-
-// PayloadHash returns the SHA256 hash of the request body for SigV4 signing.
-func (r *HttpRequest) PayloadHash() string {
-	if len(r.Body) == 0 {
-		return "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" // SHA256 of empty string
-	}
-	h := sha256.Sum256(r.Body)
-	return hex.EncodeToString(h[:])
-}
-
-// RequestInterceptor is a function that modifies an HTTP request before it is sent.
-type RequestInterceptor func(*HttpRequest) error
-
 // connectionSettings holds configuration for the connection.
 type connectionSettings struct {
 	tlsConfig                *tls.Config
@@ -106,6 +54,7 @@ type connection struct {
 	logHandler   *logHandler
 	serializer   *GraphBinarySerializer
 	interceptors []RequestInterceptor
+	wg           sync.WaitGroup
 }
 
 // Connection pool defaults aligned with Java driver
@@ -171,21 +120,19 @@ func (c *connection) AddInterceptor(interceptor RequestInterceptor) {
 }
 
 // submit sends request and streams results directly to ResultSet
-func (c *connection) submit(req *request) (ResultSet, error) {
+func (c *connection) submit(req *RequestMessage) (ResultSet, error) {
 	rs := newChannelResultSet()
 
-	data, err := c.serializer.SerializeMessage(req)
-	if err != nil {
-		rs.Close()
-		return rs, err
-	}
-
-	go c.executeAndStream(data, rs)
+	c.wg.Add(1)
+	go func() {
+		defer c.wg.Done()
+		c.executeAndStream(req, rs)
+	}()
 
 	return rs, nil
 }
 
-func (c *connection) executeAndStream(data []byte, rs ResultSet) {
+func (c *connection) executeAndStream(req *RequestMessage, rs ResultSet) {
 	defer rs.Close()
 
 	// Create HttpRequest for interceptors
@@ -195,12 +142,15 @@ func (c *connection) executeAndStream(data []byte, rs ResultSet) {
 		rs.setError(err)
 		return
 	}
-	httpReq.Body = data
 
 	// Set default headers before interceptors
 	c.setHttpRequestHeaders(httpReq)
 
-	// Apply interceptors
+	// Set Body to the raw *RequestMessage so interceptors can inspect/modify it
+	httpReq.Body = req
+
+	// Apply interceptors — they see *RequestMessage in Body (pre-serialization).
+	// Interceptors may replace Body with []byte, io.Reader, or *http.Request.
 	for _, interceptor := range c.interceptors {
 		if err := interceptor(httpReq); err != nil {
 			c.logHandler.logf(Error, failedToSendRequest, err.Error())
@@ -209,27 +159,90 @@ func (c *connection) executeAndStream(data []byte, rs ResultSet) {
 		}
 	}
 
-	// Create actual http.Request from HttpRequest
-	req, err := http.NewRequest(httpReq.Method, httpReq.URL.String(), bytes.NewReader(httpReq.Body))
-	if err != nil {
-		c.logHandler.logf(Error, failedToSendRequest, err.Error())
-		rs.setError(err)
+	// After interceptors, serialize if Body is still *RequestMessage
+	if r, ok := httpReq.Body.(*RequestMessage); ok {
+		if c.serializer != nil {
+			data, err := c.serializer.SerializeMessage(r)
+			if err != nil {
+				c.logHandler.logf(Error, failedToSendRequest, err.Error())
+				rs.setError(err)
+				return
+			}
+			httpReq.Body = data
+		} else {
+			errMsg := "request body was not serialized; either provide a serializer or add an interceptor that serializes the request"
+			c.logHandler.logf(Error, failedToSendRequest, errMsg)
+			rs.setError(fmt.Errorf("%s", errMsg))
+			return
+		}
+	}
+
+	// Create actual http.Request from HttpRequest based on Body type
+	var httpGoReq *http.Request
+	switch body := httpReq.Body.(type) {
+	case []byte:
+		httpGoReq, err = http.NewRequest(httpReq.Method, httpReq.URL.String(), bytes.NewReader(body))
+		if err != nil {
+			c.logHandler.logf(Error, failedToSendRequest, err.Error())
+			rs.setError(err)
+			return
+		}
+		httpGoReq.Header = httpReq.Headers
+	case io.Reader:
+		httpGoReq, err = http.NewRequest(httpReq.Method, httpReq.URL.String(), body)
+		if err != nil {
+			c.logHandler.logf(Error, failedToSendRequest, err.Error())
+			rs.setError(err)
+			return
+		}
+		httpGoReq.Header = httpReq.Headers
+	case *http.Request:
+		httpGoReq = body
+	default:
+		errMsg := fmt.Sprintf("unsupported body type after interceptors: %T", body)
+		c.logHandler.logf(Error, failedToSendRequest, errMsg)
+		rs.setError(fmt.Errorf("%s", errMsg))
 		return
 	}
-	req.Header = httpReq.Headers
 
-	resp, err := c.httpClient.Do(req)
+	resp, err := c.httpClient.Do(httpGoReq)
 	if err != nil {
 		c.logHandler.logf(Error, failedToSendRequest, err.Error())
 		rs.setError(err)
 		return
 	}
 	defer func() {
+		// Drain any unread bytes so the connection can be reused gracefully.
+		// Without this, Go's HTTP client sends a TCP RST instead of FIN,
+		// causing "Connection reset by peer" errors on the server.
+		io.Copy(io.Discard, resp.Body)
 		if err := resp.Body.Close(); err != nil {
 			c.logHandler.logf(Debug, failedToCloseResponseBody, err.Error())
 		}
 	}()
 
+	// If the HTTP status indicates an error and the response is not GraphBinary,
+	// read the body as a text/JSON error message instead of attempting binary
+	// deserialization which would produce cryptic errors.
+	contentType := resp.Header.Get(HeaderContentType)
+	if resp.StatusCode >= 400 && !strings.Contains(contentType, graphBinaryMimeType) {
+		bodyBytes, readErr := io.ReadAll(resp.Body)
+		if readErr != nil {
+			c.logHandler.logf(Error, failedToReceiveResponse, readErr.Error())
+			rs.setError(fmt.Errorf("Gremlin Server returned HTTP %d and failed to read body: %w",
+				resp.StatusCode, readErr))
+			return
+		}
+		errorBody := string(bodyBytes)
+		errorMsg := tryExtractJSONError(errorBody)
+		if errorMsg == "" {
+			errorMsg = fmt.Sprintf("Gremlin Server returned HTTP %d: %s", resp.StatusCode, errorBody)
+		}
+		c.logHandler.logf(Error, failedToReceiveResponse, errorMsg)
+		rs.setError(fmt.Errorf("%s", errorMsg))
+		return
+	}
+
 	reader, zlibReader, err := c.getReader(resp)
 	if err != nil {
 		c.logHandler.logf(Error, failedToReceiveResponse, err.Error())
@@ -308,6 +321,23 @@ func (c *connection) streamToResultSet(reader io.Reader, rs ResultSet) {
 	}
 }
 
+// tryExtractJSONError attempts to extract an error message from a JSON response body.
+// The server sometimes responds with a JSON object containing a "message" field
+// even when it cannot produce a GraphBinary response.
+func tryExtractJSONError(body string) string {
+	var obj map[string]interface{}
+	if err := json.Unmarshal([]byte(body), &obj); err != nil {
+		return ""
+	}
+	if msg, ok := obj["message"]; ok {
+		if s, ok := msg.(string); ok {
+			return s
+		}
+	}
+	return ""
+}
+
 func (c *connection) close() {
+	c.wg.Wait()
 	c.httpClient.CloseIdleConnections()
 }