snapshot of https://github.com/ent/ent.git (master, 2026-09-15; blobs > 131072B omitted)

dialect/gremlin/internal/ws/conn.go

@@ -0,0 +1,315 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ws
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"sync"
+	"time"
+
+	"entgo.io/ent/dialect/gremlin"
+	"entgo.io/ent/dialect/gremlin/encoding"
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+
+	"github.com/gorilla/websocket"
+	"golang.org/x/sync/errgroup"
+)
+
+const (
+	// Time allowed to write a message to the peer.
+	writeWait = 5 * time.Second
+
+	// Time allowed to read the next pong message from the peer.
+	pongWait = 10 * time.Second
+
+	// Send pings to peer with this period. Must be less than pongWait.
+	pingPeriod = (pongWait * 9) / 10
+)
+
+type (
+	// A Dialer contains options for connecting to Gremlin server.
+	Dialer struct {
+		// Underlying websocket dialer.
+		websocket.Dialer
+
+		// Gremlin server basic auth credentials.
+		user, pass string
+	}
+
+	// Conn performs operations on a gremlin server.
+	Conn struct {
+		// Underlying websocket connection.
+		conn *websocket.Conn
+
+		// Credentials for basic authentication.
+		user, pass string
+
+		// Goroutine tracking.
+		ctx context.Context
+		grp *errgroup.Group
+
+		// Channel of outbound requests.
+		send chan io.Reader
+
+		// Map of in flight requests.
+		inflight sync.Map
+	}
+
+	// inflight tracks request state.
+	inflight struct {
+		// partially received data
+		frags []graphson.RawMessage
+
+		// response channel
+		result chan<- result
+	}
+
+	// represents an execution result.
+	result struct {
+		rsp *gremlin.Response
+		err error
+	}
+)
+
+var (
+	// DefaultDialer is a dialer with all fields set to the default values.
+	DefaultDialer = &Dialer{
+		Dialer: websocket.Dialer{
+			Proxy:            http.ProxyFromEnvironment,
+			HandshakeTimeout: 5 * time.Second,
+			WriteBufferSize:  8192,
+			ReadBufferSize:   8192,
+		},
+	}
+
+	// ErrConnClosed is returned by the Conn's Execute method when
+	// the underlying gremlin server connection is closed.
+	ErrConnClosed = errors.New("gremlin: server connection closed")
+
+	// ErrDuplicateRequest is returned by the Conns Execute method on
+	// request identifier key collision.
+	ErrDuplicateRequest = errors.New("gremlin: duplicate request")
+)
+
+// Dial creates a new connection by calling DialContext with a background context.
+func (d *Dialer) Dial(uri string) (*Conn, error) {
+	return d.DialContext(context.Background(), uri)
+}
+
+// DialContext creates a new Gremlin connection.
+func (d *Dialer) DialContext(ctx context.Context, uri string) (*Conn, error) {
+	c, rsp, err := d.Dialer.DialContext(ctx, uri, nil)
+	if err != nil {
+		return nil, fmt.Errorf("gremlin: dialing uri %s: %w", uri, err)
+	}
+	defer rsp.Body.Close()
+
+	conn := &Conn{
+		conn: c,
+		user: d.user,
+		pass: d.pass,
+		send: make(chan io.Reader),
+	}
+	conn.grp, conn.ctx = errgroup.WithContext(context.Background())
+
+	conn.grp.Go(conn.sender)
+	conn.grp.Go(conn.receiver)
+
+	return conn, nil
+}
+
+// Execute executes a request against a Gremlin server.
+func (c *Conn) Execute(ctx context.Context, req *gremlin.Request) (*gremlin.Response, error) {
+	// buffered result channel prevents receiver block on context cancellation
+	result := make(chan result, 1)
+
+	// request id must be unique across inflight request
+	if _, loaded := c.inflight.LoadOrStore(req.RequestID, &inflight{result: result}); loaded {
+		return nil, ErrDuplicateRequest
+	}
+
+	pr, pw := io.Pipe()
+	defer pr.Close()
+
+	// stream graphson encoding into request
+	c.grp.Go(func() error {
+		err := graphson.NewEncoder(pw).Encode(req)
+		if err != nil {
+			err = fmt.Errorf("encoding request: %w", err)
+		}
+		pw.CloseWithError(err)
+		return err
+	})
+
+	// local copy for single write
+	send := c.send
+
+	for {
+		select {
+		case <-c.ctx.Done():
+			c.inflight.Delete(req.RequestID)
+			return nil, ErrConnClosed
+		case <-ctx.Done():
+			c.inflight.Delete(req.RequestID)
+			return nil, ctx.Err()
+		case send <- pr:
+			send = nil
+		case result := <-result:
+			return result.rsp, result.err
+		}
+	}
+}
+
+// Close connection with a Gremlin server.
+func (c *Conn) Close() error {
+	c.grp.Go(func() error { return ErrConnClosed })
+	_ = c.grp.Wait()
+	return nil
+}
+
+func (c *Conn) sender() error {
+	pinger := time.NewTicker(pingPeriod)
+	defer pinger.Stop()
+
+	// closing connection terminates receiver
+	defer c.conn.Close()
+
+	for {
+		select {
+		case r := <-c.send:
+			// ensure write completes within a window
+			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
+
+			// fetch next message writer
+			w, err := c.conn.NextWriter(websocket.BinaryMessage)
+			if err != nil {
+				return fmt.Errorf("getting message writer: %w", err)
+			}
+
+			// write mime header
+			if _, err := w.Write(encoding.GraphSON3Mime); err != nil {
+				return fmt.Errorf("writing mime header: %w", err)
+			}
+
+			// write request body
+			if _, err := io.Copy(w, r); err != nil {
+				return fmt.Errorf("writing request: %w", err)
+			}
+
+			// finish message write
+			if err := w.Close(); err != nil {
+				return fmt.Errorf("closing message writer: %w", err)
+			}
+		case <-c.ctx.Done():
+			// connection closing
+			return c.conn.WriteControl(
+				websocket.CloseMessage,
+				websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
+				time.Time{},
+			)
+		case <-pinger.C:
+			// periodic connection keepalive
+			if err := c.conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(writeWait)); err != nil {
+				return fmt.Errorf("writing ping message: %w", err)
+			}
+		}
+	}
+}
+
+func (c *Conn) receiver() error {
+	// handle keepalive responses
+	c.conn.SetReadDeadline(time.Now().Add(pongWait))
+	c.conn.SetPongHandler(func(string) error {
+		return c.conn.SetReadDeadline(time.Now().Add(pongWait))
+	})
+
+	// complete all in flight requests on termination
+	defer c.inflight.Range(func(id, ifr any) bool {
+		ifr.(*inflight).result <- result{err: ErrConnClosed}
+		c.inflight.Delete(id)
+		return true
+	})
+
+	for {
+		// rely on sender connection close during termination
+		_, r, err := c.conn.NextReader()
+		if err != nil {
+			return fmt.Errorf("writing ping message: %w", err)
+		}
+
+		// decode received response
+		var rsp gremlin.Response
+		if err := graphson.NewDecoder(r).Decode(&rsp); err != nil {
+			return fmt.Errorf("reading response: %w", err)
+		}
+
+		ifr, ok := c.inflight.Load(rsp.RequestID)
+		if !ok {
+			// context cancellation aborts inflight requests
+			continue
+		}
+
+		// handle incoming response
+		if done := c.receive(ifr.(*inflight), &rsp); done {
+			// stop tracking finished requests
+			c.inflight.Delete(rsp.RequestID)
+		}
+	}
+}
+
+func (c *Conn) receive(ifr *inflight, rsp *gremlin.Response) bool {
+	result := result{rsp: rsp}
+	switch rsp.Status.Code {
+	case gremlin.StatusSuccess:
+		// quickly handle non fragmented responses
+		if ifr.frags == nil {
+			break
+		}
+		// handle fragment
+		fallthrough
+	case gremlin.StatusPartialContent:
+		// append received fragment
+		var frag []graphson.RawMessage
+		if err := graphson.Unmarshal(rsp.Result.Data, &frag); err != nil {
+			result.err = fmt.Errorf("decoding response fragment: %w", err)
+			break
+		}
+		ifr.frags = append(ifr.frags, frag...)
+
+		// partial response requires additional fragments
+		if rsp.Status.Code == gremlin.StatusPartialContent {
+			return false
+		}
+
+		// reassemble fragmented response
+		if rsp.Result.Data, result.err = graphson.Marshal(ifr.frags); result.err != nil {
+			result.err = fmt.Errorf("assembling fragmented response: %w", result.err)
+		}
+	case gremlin.StatusAuthenticate:
+		// receiver should never block
+		c.grp.Go(func() error {
+			var buf bytes.Buffer
+			if err := graphson.NewEncoder(&buf).Encode(
+				gremlin.NewAuthRequest(rsp.RequestID, c.user, c.pass),
+			); err != nil {
+				return fmt.Errorf("encoding auth request: %w", err)
+			}
+			select {
+			case c.send <- &buf:
+			case <-c.ctx.Done():
+			}
+			return c.ctx.Err()
+		})
+		return false
+	}
+
+	ifr.result <- result
+	return true
+}

dialect/gremlin/internal/ws/conn_test.go

@@ -0,0 +1,390 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ws
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"sync"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin"
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type conn struct{ *websocket.Conn }
+
+func (c conn) ReadRequest() (*gremlin.Request, error) {
+	_, data, err := c.ReadMessage()
+	if err != nil {
+		return nil, err
+	}
+	var req gremlin.Request
+	if err := graphson.Unmarshal(data[data[0]+1:], &req); err != nil {
+		return nil, err
+	}
+	return &req, nil
+}
+
+func (c conn) WriteResponse(rsp *gremlin.Response) error {
+	data, err := graphson.Marshal(rsp)
+	if err != nil {
+		return err
+	}
+	return c.WriteMessage(websocket.BinaryMessage, data)
+}
+
+func serve(handler func(conn)) *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		upgrader := websocket.Upgrader{ReadBufferSize: 1024, WriteBufferSize: 1024}
+		c, _ := upgrader.Upgrade(w, r, nil)
+		defer c.Close()
+		handler(conn{c})
+		for {
+			_, _, err := c.ReadMessage()
+			if err != nil {
+				break
+			}
+		}
+	}))
+}
+
+func TestConnectClosure(t *testing.T) {
+	var wg sync.WaitGroup
+	wg.Add(1)
+	defer wg.Wait()
+
+	srv := serve(func(conn conn) {
+		defer wg.Done()
+		_, _, err := conn.ReadMessage()
+		assert.True(t, websocket.IsCloseError(err, websocket.CloseNormalClosure))
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+
+	err = conn.Close()
+	assert.NoError(t, err)
+
+	_, err = conn.Execute(context.Background(), gremlin.NewEvalRequest("g.V()"))
+	assert.EqualError(t, err, ErrConnClosed.Error())
+}
+
+func TestSimpleQuery(t *testing.T) {
+	srv := serve(func(conn conn) {
+		typ, data, err := conn.ReadMessage()
+		require.NoError(t, err)
+		assert.Equal(t, websocket.BinaryMessage, typ)
+
+		var req gremlin.Request
+		err = graphson.Unmarshal(data[data[0]+1:], &req)
+		require.NoError(t, err)
+		assert.Equal(t, "g.V()", req.Arguments["gremlin"])
+
+		rsp := gremlin.Response{RequestID: req.RequestID}
+		rsp.Status.Code = gremlin.StatusNoContent
+		err = conn.WriteResponse(&rsp)
+		require.NoError(t, err)
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer assert.Condition(t, func() bool { return assert.NoError(t, conn.Close()) })
+
+	rsp, err := conn.Execute(context.Background(), gremlin.NewEvalRequest("g.V()"))
+	assert.NoError(t, err)
+	require.NotNil(t, rsp)
+	assert.Equal(t, gremlin.StatusNoContent, rsp.Status.Code)
+}
+
+func TestDuplicateRequest(t *testing.T) {
+	// skip until flakiness will be fixed.
+	t.SkipNow()
+	srv := serve(func(conn conn) {
+		req, err := conn.ReadRequest()
+		require.NoError(t, err)
+
+		rsp := gremlin.Response{RequestID: req.RequestID}
+		rsp.Status.Code = gremlin.StatusNoContent
+		err = conn.WriteResponse(&rsp)
+		require.NoError(t, err)
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	var errors [2]error
+	req := gremlin.NewEvalRequest("g.V()")
+
+	var wg sync.WaitGroup
+	wg.Add(len(errors))
+
+	for i := range errors {
+		go func(i int) {
+			_, errors[i] = conn.Execute(context.Background(), req)
+			wg.Done()
+		}(i)
+	}
+	wg.Wait()
+
+	err = errors[0]
+	if err == nil {
+		err = errors[1]
+	}
+	assert.EqualError(t, err, ErrDuplicateRequest.Error())
+}
+
+func TestConnectCancellation(t *testing.T) {
+	srv := serve(func(conn) {})
+	defer srv.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	conn, err := DefaultDialer.DialContext(ctx, "ws://"+srv.Listener.Addr().String())
+	assert.Error(t, err)
+	assert.Nil(t, conn)
+}
+
+func TestQueryCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	srv := serve(func(conn conn) {
+		if _, _, err := conn.ReadMessage(); err == nil {
+			cancel()
+		}
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	_, err = conn.Execute(ctx, gremlin.NewEvalRequest("g.E()"))
+	assert.EqualError(t, err, context.Canceled.Error())
+}
+
+func TestBadResponse(t *testing.T) {
+	tests := []struct {
+		name   string
+		mangle func(*gremlin.Response) *gremlin.Response
+	}{
+		{
+			name: "NoStatus",
+			mangle: func(rsp *gremlin.Response) *gremlin.Response {
+				return rsp
+			},
+		},
+		{
+			name: "Malformed",
+			mangle: func(rsp *gremlin.Response) *gremlin.Response {
+				rsp.Status.Code = gremlin.StatusMalformedRequest
+				rsp.Status.Message = "bad request"
+				return rsp
+			},
+		},
+		{
+			name: "Unknown",
+			mangle: func(rsp *gremlin.Response) *gremlin.Response {
+				rsp.Status.Code = 424242
+				return rsp
+			},
+		},
+	}
+
+	srv := serve(func(conn conn) {
+		for {
+			req, err := conn.ReadRequest()
+			if err != nil {
+				break
+			}
+
+			idx, err := strconv.ParseInt(req.Arguments["gremlin"].(string), 10, 0)
+			require.NoError(t, err)
+
+			err = conn.WriteResponse(tests[idx].mangle(&gremlin.Response{RequestID: req.RequestID}))
+			require.NoError(t, err)
+		}
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	var wg sync.WaitGroup
+	wg.Add(len(tests))
+
+	ctx := context.Background()
+	for i, tc := range tests {
+		i, tc := i, tc
+		t.Run(tc.name, func(t *testing.T) {
+			defer wg.Done()
+			rsp, err := conn.Execute(ctx, gremlin.NewEvalRequest(strconv.FormatInt(int64(i), 10)))
+			assert.NoError(t, err)
+			assert.True(t, rsp.IsErr())
+		})
+	}
+	wg.Wait()
+}
+
+func TestServerHangup(t *testing.T) {
+	// skip until flakiness will be fixed.
+	t.SkipNow()
+	srv := serve(func(conn conn) { _ = conn.Close() })
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	_, err = conn.Execute(context.Background(), gremlin.NewEvalRequest("g.V()"))
+	assert.EqualError(t, err, ErrConnClosed.Error())
+	assert.Error(t, conn.ctx.Err())
+}
+
+func TestCanceledLongRequest(t *testing.T) {
+	// skip until flakiness will be fixed.
+	t.SkipNow()
+	ctx, cancel := context.WithCancel(context.Background())
+	srv := serve(func(conn conn) {
+		var responses [3]*gremlin.Response
+		for i := 0; i < len(responses); i++ {
+			req, err := conn.ReadRequest()
+			require.NoError(t, err)
+
+			rsp := gremlin.Response{RequestID: req.RequestID}
+			rsp.Status.Code = gremlin.StatusSuccess
+			rsp.Result.Data = graphson.RawMessage(`"ok"`)
+			responses[i] = &rsp
+		}
+
+		cancel()
+
+		responses[0], responses[2] = responses[2], responses[0]
+		for i := 0; i < len(responses); i++ {
+			err := conn.WriteResponse(responses[i])
+			require.NoError(t, err)
+		}
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	var wg sync.WaitGroup
+	wg.Add(3)
+	defer wg.Wait()
+
+	for i := 0; i < 3; i++ {
+		go func(ctx context.Context, idx int) {
+			defer wg.Done()
+			rsp, err := conn.Execute(ctx, gremlin.NewEvalRequest("g.V()"))
+			if idx > 0 {
+				assert.NoError(t, err)
+				assert.EqualValues(t, []byte(`"ok"`), rsp.Result.Data)
+			} else {
+				assert.EqualError(t, err, context.Canceled.Error())
+			}
+		}(ctx, i)
+		ctx = context.Background()
+	}
+}
+
+func TestPartialResponse(t *testing.T) {
+	type kv struct {
+		Key   string
+		Value int
+	}
+	kvs := []kv{
+		{"one", 1},
+		{"two", 2},
+		{"three", 3},
+	}
+	srv := serve(func(conn conn) {
+		req, err := conn.ReadRequest()
+		require.NoError(t, err)
+
+		for i := range kvs {
+			data, err := graphson.Marshal([]kv{kvs[i]})
+			require.NoError(t, err)
+
+			rsp := gremlin.Response{RequestID: req.RequestID}
+			rsp.Result.Data = graphson.RawMessage(data)
+
+			if i != len(kvs)-1 {
+				rsp.Status.Code = gremlin.StatusPartialContent
+			} else {
+				rsp.Status.Code = gremlin.StatusSuccess
+			}
+
+			err = conn.WriteResponse(&rsp)
+			require.NoError(t, err)
+		}
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	rsp, err := conn.Execute(context.Background(), gremlin.NewEvalRequest("g.E()"))
+	assert.NoError(t, err)
+
+	var result []kv
+	err = graphson.Unmarshal(rsp.Result.Data, &result)
+	require.NoError(t, err)
+	assert.Equal(t, kvs, result)
+}
+
+func TestAuthentication(t *testing.T) {
+	user, pass := "username", "password"
+	srv := serve(func(conn conn) {
+		req, err := conn.ReadRequest()
+		require.NoError(t, err)
+
+		rsp := gremlin.Response{RequestID: req.RequestID}
+		rsp.Status.Code = gremlin.StatusAuthenticate
+		err = conn.WriteResponse(&rsp)
+		require.NoError(t, err)
+
+		areq, err := conn.ReadRequest()
+		require.NoError(t, err)
+
+		var acreds gremlin.Credentials
+		err = acreds.UnmarshalText([]byte(areq.Arguments["sasl"].(string)))
+		assert.NoError(t, err)
+		areq.Arguments["sasl"] = acreds
+		assert.Equal(t, gremlin.NewAuthRequest(req.RequestID, user, pass), areq)
+
+		rsp = gremlin.Response{RequestID: req.RequestID}
+		rsp.Status.Code = gremlin.StatusNoContent
+		err = conn.WriteResponse(&rsp)
+		require.NoError(t, err)
+	})
+	defer srv.Close()
+
+	dialer := *DefaultDialer
+	dialer.user = user
+	dialer.pass = pass
+
+	client, err := dialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer client.Close()
+
+	_, err = client.Execute(context.Background(), gremlin.NewEvalRequest("g.E().drop()"))
+	assert.NoError(t, err)
+}