feat(x): use endpoint instea of dialer for websocket

Jigsaw-Code · Jan 30, 2025 · c597671 · c597671
1 parent 80b6430
commit c597671
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/x/websocket/endpoint.go b/x/websocket/endpoint.go
@@ -36,13 +36,13 @@ import (
 // Websockets, with each Write becoming a separate message. Half-close is supported:
 // CloseRead will not close the Websocket connection, while CloseWrite sends a Websocket
 // close but continues reading until a close is received from the server.
-func NewStreamEndpoint(urlStr string, sd transport.StreamDialer, opts ...Option) (func(context.Context) (transport.StreamConn, error), error) {
-	return newEndpoint(urlStr, sd, func(gc *gorillaConn) transport.StreamConn { return gc }, opts...)
+func NewStreamEndpoint(urlStr string, se transport.StreamEndpoint, opts ...Option) (func(context.Context) (transport.StreamConn, error), error) {
+	return newEndpoint(urlStr, se, func(gc *gorillaConn) transport.StreamConn { return gc }, opts...)
 }
 
 // NewPacketEndpoint creates a new Websocket Packet Endpoint. Each packet is exchanged as a Websocket message.
-func NewPacketEndpoint(urlStr string, sd transport.StreamDialer, opts ...Option) (func(context.Context) (net.Conn, error), error) {
-	return newEndpoint(urlStr, sd, func(gc *gorillaConn) net.Conn { return gc }, opts...)
+func NewPacketEndpoint(urlStr string, se transport.StreamEndpoint, opts ...Option) (func(context.Context) (net.Conn, error), error) {
+	return newEndpoint(urlStr, se, func(gc *gorillaConn) net.Conn { return gc }, opts...)
 }
 
 type options struct {
@@ -68,7 +68,7 @@ func WithHTTPHeaders(headers http.Header) Option {
 	}
 }
 
-func newEndpoint[ConnType net.Conn](urlStr string, sd transport.StreamDialer, wsToConn func(*gorillaConn) ConnType, opts ...Option) (func(context.Context) (ConnType, error), error) {
+func newEndpoint[ConnType net.Conn](urlStr string, se transport.StreamEndpoint, wsToConn func(*gorillaConn) ConnType, opts ...Option) (func(context.Context) (ConnType, error), error) {
 	_, err := url.Parse(urlStr)
 	if err != nil {
 		return nil, fmt.Errorf("url is invalid: %w", err)
@@ -84,11 +84,11 @@ func newEndpoint[ConnType net.Conn](urlStr string, sd transport.StreamDialer, ws
 
 	wsDialer := &websocket.Dialer{
 		TLSClientConfig: resolvedOpts.tlsConfig,
-		NetDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+		NetDialContext: func(ctx context.Context, network, _ string) (net.Conn, error) {
 			if !strings.HasPrefix(network, "tcp") {
 				return nil, fmt.Errorf("websocket dialer does not support network type %v", network)
 			}
-			return sd.DialStream(ctx, addr)
+			return se.ConnectStream(ctx)
 		},
 	}
 	return func(ctx context.Context) (ConnType, error) {

diff --git a/x/websocket/endpoint_test.go b/x/websocket/endpoint_test.go
@@ -91,7 +91,8 @@ func Test_NewStreamEndpoint(t *testing.T) {
 	// 	},
 	// }
 	client := ts.Client()
-	connect, err := NewStreamEndpoint("wss"+ts.URL[5:]+"/tcp", &transport.TCPDialer{}, WithTLSConfig(client.Transport.(*http.Transport).TLSClientConfig))
+	endpoint := &transport.TCPEndpoint{Address: ts.Listener.Addr().String()}
+	connect, err := NewStreamEndpoint("wss"+ts.URL[5:]+"/tcp", endpoint, WithTLSConfig(client.Transport.(*http.Transport).TLSClientConfig))
 	require.NoError(t, err)
 	require.NotNil(t, connect)
 
@@ -142,7 +143,8 @@ func Test_NewPacketEndpoint(t *testing.T) {
 	defer ts.Close()
 
 	client := ts.Client()
-	connect, err := NewPacketEndpoint("wss"+ts.URL[5:]+"/udp", &transport.TCPDialer{}, WithTLSConfig(client.Transport.(*http.Transport).TLSClientConfig))
+	endpoint := &transport.TCPEndpoint{Address: ts.Listener.Addr().String()}
+	connect, err := NewPacketEndpoint("wss"+ts.URL[5:]+"/udp", endpoint, WithTLSConfig(client.Transport.(*http.Transport).TLSClientConfig))
 	require.NoError(t, err)
 	require.NotNil(t, connect)