refs #95: implemented websocket support;

client.go

@@ -37,19 +37,19 @@ type HttpClient interface {
 	Post(url, contentType string, body io.Reader) (resp *http.Response, err error)
 
 	GetBandwidthTracker() bandwidth.BandwidthTracker
+	GetDialer() proxy.ContextDialer
 }
 
 // Interface guards are a cheap way to make sure all methods are implemented, this is a static check and does not affect runtime performance.
 var _ HttpClient = (*httpClient)(nil)
 
 type httpClient struct {
-	logger Logger
-
+	http.Client
+	logger           Logger
 	bandwidthTracker bandwidth.BandwidthTracker
 	config           *httpClientConfig
-
-	http.Client
-	headerLck sync.Mutex
+	headerLck        sync.Mutex
+	dialer           proxy.ContextDialer
 }
 
 var DefaultTimeoutSeconds = 30
@@ -87,7 +87,7 @@ func NewHttpClient(logger Logger, options ...HttpClientOption) (HttpClient, erro
 		return nil, err
 	}
 
-	client, bandwidthTracker, clientProfile, err := buildFromConfig(logger, config)
+	client, dialer, bandwidthTracker, clientProfile, err := buildFromConfig(logger, config)
 	if err != nil {
 		return nil, err
 	}
@@ -112,6 +112,7 @@ func NewHttpClient(logger Logger, options ...HttpClientOption) (HttpClient, erro
 		config:           config,
 		headerLck:        sync.Mutex{},
 		bandwidthTracker: bandwidthTracker,
+		dialer:           dialer,
 	}, nil
 }
 
@@ -143,14 +144,14 @@ func validateConfig(config *httpClientConfig) error {
 	return nil
 }
 
-func buildFromConfig(logger Logger, config *httpClientConfig) (*http.Client, bandwidth.BandwidthTracker, profiles.ClientProfile, error) {
+func buildFromConfig(logger Logger, config *httpClientConfig) (*http.Client, proxy.ContextDialer, bandwidth.BandwidthTracker, profiles.ClientProfile, error) {
 	var dialer proxy.ContextDialer
 	dialer = newDirectDialer(config.timeout, config.localAddr, config.dialer)
 
 	if config.proxyUrl != "" && config.proxyDialerFactory == nil {
 		proxyDialer, err := newConnectDialer(config.proxyUrl, config.timeout, config.localAddr, config.dialer, config.connectHeaders, logger)
 		if err != nil {
-			return nil, nil, profiles.ClientProfile{}, err
+			return nil, nil, nil, profiles.ClientProfile{}, err
 		}
 
 		dialer = proxyDialer
@@ -159,7 +160,7 @@ func buildFromConfig(logger Logger, config *httpClientConfig) (*http.Client, ban
 	if config.proxyDialerFactory != nil {
 		proxyDialer, err := config.proxyDialerFactory(config.proxyUrl, config.timeout, config.localAddr, config.connectHeaders, logger)
 		if err != nil {
-			return nil, nil, profiles.ClientProfile{}, err
+			return nil, nil, nil, profiles.ClientProfile{}, err
 		}
 
 		dialer = proxyDialer
@@ -187,7 +188,7 @@ func buildFromConfig(logger Logger, config *httpClientConfig) (*http.Client, ban
 
 	transport, err := newRoundTripper(clientProfile, config.transportOptions, config.serverNameOverwrite, config.insecureSkipVerify, config.withRandomTlsExtensionOrder, config.forceHttp1, config.disableHttp3, config.enableProtocolRacing, config.certificatePins, config.badPinHandler, config.disableIPV6, config.disableIPV4, bandwidthTracker, dialer)
 	if err != nil {
-		return nil, nil, clientProfile, err
+		return nil, nil, nil, clientProfile, err
 	}
 
 	client := &http.Client{
@@ -200,14 +201,19 @@ func buildFromConfig(logger Logger, config *httpClientConfig) (*http.Client, ban
 		client.Jar = config.cookieJar
 	}
 
-	return client, bandwidthTracker, clientProfile, nil
+	return client, dialer, bandwidthTracker, clientProfile, nil
 }
 
 // CloseIdleConnections closes all idle connections of the underlying http client.
 func (c *httpClient) CloseIdleConnections() {
 	c.Client.CloseIdleConnections()
 }
 
+// GetDialer() returns the underlying Dialer
+func (c *httpClient) GetDialer() proxy.ContextDialer {
+	return c.dialer
+}
+
 // SetFollowRedirect configures the client's HTTP redirect following policy.
 func (c *httpClient) SetFollowRedirect(followRedirect bool) {
 	c.logger.Debug("set follow redirect from %v to %v", c.config.followRedirects, followRedirect)

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/bogdanfinn/fhttp v0.6.6
 	github.com/bogdanfinn/quic-go-utls v1.0.7-utls
 	github.com/bogdanfinn/utls v1.7.7-barnius
+	github.com/bogdanfinn/websocket v1.5.4-barnius
 	github.com/google/uuid v1.6.0
 	github.com/stretchr/testify v1.11.1
 	github.com/tam7t/hpkp v0.0.0-20160821193359-2b70b4024ed5
@@ -31,4 +32,6 @@ require (
 
 // replace github.com/bogdanfinn/quic-go-utls => ../quic-go-utls
 
+// replace github.com/bogdanfinn/websocket => ../websocket
+
 // replace github.com/bogdanfinn/fhttp => ../fhttp

go.sum

@@ -10,6 +10,8 @@ github.com/bogdanfinn/quic-go-utls v1.0.7-utls h1:opxU/wt2C6FcD3rkGSOwfpQgfGSFx9
 github.com/bogdanfinn/quic-go-utls v1.0.7-utls/go.mod h1:bk8QMY2KypO8A6LzHJ7C4+bdB0ksLOd6NZt600wXYe8=
 github.com/bogdanfinn/utls v1.7.7-barnius h1:OuJ497cc7F3yKNVHRsYPQdGggmk5x6+V5ZlrCR7fOLU=
 github.com/bogdanfinn/utls v1.7.7-barnius/go.mod h1:aAK1VZQlpKZClF1WEQeq6kyclbkPq4hz6xTbB5xSlmg=
+github.com/bogdanfinn/websocket v1.5.4-barnius h1:qn3DU/KMHMNPNnwDCtA/IN3QqmzV98DVsvQkQkxruKw=
+github.com/bogdanfinn/websocket v1.5.4-barnius/go.mod h1:IWHoWLZd+5/o9340A/m/ApAOBmWDSTUXuL7zSVmRHQk=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=

tests/websocket_test.go

@@ -0,0 +1,172 @@
+package tests
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/bogdanfinn/fhttp"
+	"github.com/bogdanfinn/fhttp/httptest"
+	tls_client "github.com/bogdanfinn/tls-client"
+	"github.com/bogdanfinn/tls-client/profiles"
+	gorillaWebsocket "github.com/bogdanfinn/websocket"
+	"github.com/stretchr/testify/require"
+)
+
+var upgrader = gorillaWebsocket.Upgrader{}
+
+func echoHandler(w http.ResponseWriter, r *http.Request) {
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		return
+	}
+	defer conn.Close()
+
+	for {
+		mt, msg, err := conn.ReadMessage()
+		if err != nil {
+			return
+		}
+
+		if err := conn.WriteMessage(mt, msg); err != nil {
+			return
+		}
+	}
+}
+
+func TestWebSocketEcho(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(echoHandler))
+	defer server.Close()
+
+	url := "ws" + server.URL[len("http"):]
+
+	options := []tls_client.HttpClientOption{
+		tls_client.WithClientProfile(profiles.Chrome_133),
+		tls_client.WithRandomTLSExtensionOrder(),
+	}
+
+	client, err := tls_client.NewHttpClient(nil, options...)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	websocketOptions := []tls_client.WebsocketOption{
+		tls_client.WithTlsClient(client),
+		tls_client.WithUrl(url),
+		tls_client.WithHeaders(http.Header{}),
+		tls_client.WithHandshakeTimeoutMilliseconds(1000),
+	}
+
+	ws, err := tls_client.New(nil, websocketOptions...)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	wsConnection, err := ws.Connect(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer wsConnection.Close()
+
+	expected := "hello world"
+	err = wsConnection.WriteMessage(gorillaWebsocket.TextMessage, []byte(expected))
+	require.NoError(t, err)
+
+	_, msg, err := wsConnection.ReadMessage()
+	require.NoError(t, err)
+	require.Equal(t, expected, string(msg))
+
+	wsConnection.SetReadDeadline(time.Now().Add(2 * time.Second))
+}
+
+func TestWebSocketWithHeaderOrder(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(echoHandler))
+	defer server.Close()
+
+	url := "ws" + server.URL[len("http"):]
+
+	options := []tls_client.HttpClientOption{
+		tls_client.WithClientProfile(profiles.Chrome_133),
+		tls_client.WithRandomTLSExtensionOrder(),
+		tls_client.WithForceHttp1(),
+	}
+
+	client, err := tls_client.NewHttpClient(nil, options...)
+	require.NoError(t, err)
+
+	customHeaders := http.Header{
+		"User-Agent":    {"CustomBot/1.0"},
+		"Custom-Header": {"CustomValue"},
+		http.HeaderOrderKey: {
+			"host",
+			"upgrade",
+			"connection",
+			"sec-websocket-key",
+			"sec-websocket-version",
+			"user-agent",
+			"custom-header",
+		},
+	}
+
+	websocketOptions := []tls_client.WebsocketOption{
+		tls_client.WithTlsClient(client),
+		tls_client.WithUrl(url),
+		tls_client.WithHeaders(customHeaders),
+		tls_client.WithHandshakeTimeoutMilliseconds(1000),
+	}
+
+	ws, err := tls_client.New(nil, websocketOptions...)
+	require.NoError(t, err)
+
+	wsConnection, err := ws.Connect(context.Background())
+	require.NoError(t, err)
+	defer wsConnection.Close()
+
+	expected := "header order test"
+	err = wsConnection.WriteMessage(gorillaWebsocket.TextMessage, []byte(expected))
+	require.NoError(t, err)
+
+	_, msg, err := wsConnection.ReadMessage()
+	require.NoError(t, err)
+	require.Equal(t, expected, string(msg))
+}
+
+func TestWebSocketWithoutHeaderOrder(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(echoHandler))
+	defer server.Close()
+
+	url := "ws" + server.URL[len("http"):]
+
+	options := []tls_client.HttpClientOption{
+		tls_client.WithClientProfile(profiles.Chrome_133),
+		tls_client.WithRandomTLSExtensionOrder(),
+	}
+
+	client, err := tls_client.NewHttpClient(nil, options...)
+	require.NoError(t, err)
+
+	websocketOptions := []tls_client.WebsocketOption{
+		tls_client.WithTlsClient(client),
+		tls_client.WithUrl(url),
+		tls_client.WithHeaders(http.Header{
+			"User-Agent": {"TestBot/1.0"},
+		}),
+		tls_client.WithHandshakeTimeoutMilliseconds(1000),
+	}
+
+	ws, err := tls_client.New(nil, websocketOptions...)
+	require.NoError(t, err)
+
+	wsConnection, err := ws.Connect(context.Background())
+	require.NoError(t, err)
+	defer wsConnection.Close()
+
+	expected := "no header order test"
+	err = wsConnection.WriteMessage(gorillaWebsocket.TextMessage, []byte(expected))
+	require.NoError(t, err)
+
+	_, msg, err := wsConnection.ReadMessage()
+	require.NoError(t, err)
+	require.Equal(t, expected, string(msg))
+}

websocket.go

@@ -0,0 +1,86 @@
+package tls_client
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/bogdanfinn/websocket"
+)
+
+type Websocket struct {
+	config websocketConfig
+	dialer *websocket.Dialer
+}
+
+// New creates a new WebSocket wrapper that uses tls-client for connections.
+// This allows WebSocket connections to use the same TLS fingerprinting and
+// configuration as regular HTTP requests.
+//
+// Example usage:
+//
+//	// Create HTTP client with ForceHttp1 (required for WebSocket!)
+//	client, _ := NewHttpClient(nil,
+//	    WithClientProfile(profiles.Chrome_133),
+//	    WithForceHttp1(),
+//	)
+//
+//	// Create WebSocket with optional header ordering
+//	headers := http.Header{
+//	    "User-Agent": {"MyBot/1.0"},
+//	    http.HeaderOrderKey: {"host", "upgrade", "connection", "user-agent"},
+//	}
+//
+//	ws, _ := New(nil,
+//	    WithTlsClient(client),
+//	    WithUrl("wss://example.com/ws"),
+//	    WithHeaders(headers),
+//	)
+//
+//	conn, _ := ws.Connect(context.Background())
+//	defer conn.Close()
+func New(logger Logger, options ...WebsocketOption) (*Websocket, error) {
+	config := &websocketConfig{}
+
+	for _, opt := range options {
+		opt(config)
+	}
+
+	if err := validateWebsocketConfig(config); err != nil {
+		return nil, err
+	}
+
+	dialer := &websocket.Dialer{
+		HandshakeTimeout:  config.handshakeTimeout,
+		Jar:               config.cookieJar,
+		ReadBufferSize:    config.readBufferSize,
+		WriteBufferSize:   config.writeBufferSize,
+		NetDialTLSContext: config.tlsClient.GetDialer().DialContext,
+		NetDialContext:    config.tlsClient.GetDialer().DialContext,
+	}
+
+	return &Websocket{
+		config: *config,
+		dialer: dialer,
+	}, nil
+}
+
+func (w *Websocket) Connect(ctx context.Context) (*websocket.Conn, error) {
+	c, _, err := w.dialer.DialContext(ctx, w.config.url, w.config.headers)
+	if err != nil {
+		return nil, err
+	}
+
+	return c, nil
+}
+
+func validateWebsocketConfig(config *websocketConfig) error {
+	if config.tlsClient == nil {
+		return fmt.Errorf("tlsClient cannot be nil for websocket connection")
+	}
+
+	if config.url == "" {
+		return fmt.Errorf("url cannot be empty for websocket connection")
+	}
+
+	return nil
+}