<commit_before>package rpcc

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"sync"

	"github.com/gorilla/websocket"
)

var (
	// ErrConnClosing indicates that the operation is illegal because
	// the connection is closing.
	ErrConnClosing = errors.New("rpcc: the connection is closing")
)

const (
	defaultWriteBufferSize = 4096

	// Fixed header size used by gorilla/websocket.
	websocketMaxHeaderSize = 2 + 8 + 4

	// Maximum message size accepted by Chrome is 1MB. It silently
	// drops the websocket connection for larger messages.
	chromeMessageSizeLimit = 1024 * 1024
)

// DialOption represents a dial option passed to Dial.
type DialOption func(*dialOptions)

// WithCodec returns a DialOption that sets the codec responsible for
// encoding and decoding requests and responses onto the connection.
// This option overrides the default json codec.
func WithCodec(f func(conn io.ReadWriter) Codec) DialOption {
	return func(o *dialOptions) {
		o.codec = f
	}
}

// WithDialer returns a DialOption that sets the dialer for the underlying
// net.Conn. This option overrides the default WebSocket dialer.
func WithDialer(f func(ctx context.Context, addr string) (net.Conn, error)) DialOption {
	return func(o *dialOptions) {
		o.dialer = f
	}
}

// WithWriteBufferSize returns a DialOption that sets the size of the write
// buffer for the underlying websocket connection. Messages larger than this
// size are fragmented according to the websocket specification.
func WithWriteBufferSize(n int) DialOption {
	return func(o *dialOptions) {
		o.wsDialer.WriteBufferSize = n
	}
}

// WithCompression returns a DialOption that enables compression for the
// underlying websocket connection. Use SetCompressionLevel on Conn to
// change the default compression level for subsequent writes.
func WithCompression() DialOption {
	return func(o *dialOptions) {
		o.wsDialer.EnableCompression = true
	}
}

type dialOptions struct {
	codec    func(io.ReadWriter) Codec
	dialer   func(context.Context, string) (net.Conn, error)
	wsDialer websocket.Dialer
}

// Dial connects to target and returns an active connection.
// The target should be a WebSocket URL, format:
// "ws://localhost:9222/target".
func Dial(target string, opts ...DialOption) (*Conn, error) {
	return DialContext(context.Background(), target, opts...)
}

// DialContext is like Dial, with a caller provided Context.
// A nil Context will panic.
func DialContext(ctx context.Context, target string, opts ...DialOption) (conn *Conn, err error) {
	if ctx == nil {
		panic("nil Context")
	}

	c := &Conn{
		pending: make(map[uint64]*rpcCall),
		streams: make(map[string]*streamClients),
	}
	c.ctx, c.cancel = context.WithCancel(context.Background())

	defer func() {
		select {
		case <-ctx.Done():
			conn, err = nil, ctx.Err()
		default:
		}
		if err != nil {
			c.Close()
		}
	}()

	for _, o := range opts {
		o(&c.dialOpts)
	}

	netDial := c.dialOpts.dialer
	if netDial == nil {
		netDial = func(ctx context.Context, addr string) (net.Conn, error) {
			ws := &c.dialOpts.wsDialer

			if ws.WriteBufferSize == 0 {
				// Set the default size for use in writeLimiter.
				ws.WriteBufferSize = defaultWriteBufferSize
			}

			// Base writeLimit on the buffer size allocated by the
			// websocket package, we assume a message of this size
			// is fragmented and subtract one.
			writeLimit := ws.WriteBufferSize + websocketMaxHeaderSize - 1
			if writeLimit > chromeMessageSizeLimit {
				writeLimit = chromeMessageSizeLimit
			}

			// Set NetDial to dial with context, this action will
			// override the HandshakeTimeout setting.
			ws.NetDial = func(network, addr string) (net.Conn, error) {
				conn, err := (&net.Dialer{}).DialContext(ctx, network, addr)
				return &writeLimiter{
					limit: writeLimit,
					Conn:  conn,
				}, err
			}

			wsConn, _, err := ws.Dial(addr, nil)
			if err != nil {
				return nil, err
			}

			if ws.EnableCompression {
				c.compressionLevel = wsConn.SetCompressionLevel
			}

			return &wsNetConn{conn: wsConn}, nil
		}
	}

	c.conn, err = netDial(ctx, target)
	if err != nil {
		return nil, err
	}
	newCodec := c.dialOpts.codec
	if newCodec == nil {
		newCodec = func(conn io.ReadWriter) Codec {
			return &jsonCodec{
				enc: json.NewEncoder(conn),
				dec: json.NewDecoder(conn),
			}
		}
	}
	c.codec = newCodec(c.conn)

	recvDone := func(err error) {
		// When we receive Inspector.detached the remote will close
		// the connection afterwards and recvDone will return. Maybe
		// we could give the user time to react to the event before
		// closing?
		// TODO(mafredri): Do we want to close here, like this?
		c.close(err)
	}
	go c.recv(c.notify, recvDone)

	return c, nil
}

// writeLimiter wraps a net.Conn and prevents writes of greater length
// than limit. Works around Chrome's lack of support for large or
// fragmented websocket messages and prevents sudden termination of the
// websocket connection. Gives the user an actionable error message when
// writes exceed limit.
type writeLimiter struct {
	limit int
	net.Conn
}

// BUG(mafredri): Chrome does not support websocket fragmentation
// (continuation messages) or messages that exceed 1 MB in size.
// See https://github.com/mafredri/cdp/issues/4.
func (c *writeLimiter) Write(b []byte) (n int, err error) {
	if len(b) > c.limit {
		return 0, errors.New("rpcc: message too large (increase write buffer size or enable compression)")
	}
	return c.Conn.Write(b)
}
func (c *writeLimiter) Read(b []byte) (n int, err error) {
	return c.Conn.Read(b)
}

// Codec is used by recv and dispatcher to
// send and receive RPC communication.
type Codec interface {
	// WriteRequest encodes and writes the request onto the
	// underlying connection. Request is re-used between writes and
	// references to it should not be kept.
	WriteRequest(*Request) error
	// ReadResponse decodes a response from the underlying
	// connection. Response is re-used between reads and references
	// to it should not be kept.<commit_msg>rpcc: Remove hard-coded 1MB write buffer size limit<commit_after>package rpcc

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"sync"

	"github.com/gorilla/websocket"
)

var (
	// ErrConnClosing indicates that the operation is illegal because
	// the connection is closing.
	ErrConnClosing = errors.New("rpcc: the connection is closing")
)

const (
	defaultWriteBufferSize = 4096

	// Fixed header size used by gorilla/websocket.
	websocketMaxHeaderSize = 2 + 8 + 4
)

// DialOption represents a dial option passed to Dial.
type DialOption func(*dialOptions)

// WithCodec returns a DialOption that sets the codec responsible for
// encoding and decoding requests and responses onto the connection.
// This option overrides the default json codec.
func WithCodec(f func(conn io.ReadWriter) Codec) DialOption {
	return func(o *dialOptions) {
		o.codec = f
	}
}

// WithDialer returns a DialOption that sets the dialer for the underlying
// net.Conn. This option overrides the default WebSocket dialer.
func WithDialer(f func(ctx context.Context, addr string) (net.Conn, error)) DialOption {
	return func(o *dialOptions) {
		o.dialer = f
	}
}

// WithWriteBufferSize returns a DialOption that sets the size of the write
// buffer for the underlying websocket connection. Messages larger than this
// size are fragmented according to the websocket specification.
//
// The maximum buffer size for recent versions of Chrome is 104857586 (~100MB),
// for older versions a maximum of 1048562 (~1MB) can be used. This is because
// Chrome does not support websocket fragmentation.
func WithWriteBufferSize(n int) DialOption {
	return func(o *dialOptions) {
		o.wsDialer.WriteBufferSize = n
	}
}

// WithCompression returns a DialOption that enables compression for the
// underlying websocket connection. Use SetCompressionLevel on Conn to
// change the default compression level for subsequent writes.
func WithCompression() DialOption {
	return func(o *dialOptions) {
		o.wsDialer.EnableCompression = true
	}
}

type dialOptions struct {
	codec    func(io.ReadWriter) Codec
	dialer   func(context.Context, string) (net.Conn, error)
	wsDialer websocket.Dialer
}

// Dial connects to target and returns an active connection.
// The target should be a WebSocket URL, format:
// "ws://localhost:9222/target".
func Dial(target string, opts ...DialOption) (*Conn, error) {
	return DialContext(context.Background(), target, opts...)
}

// DialContext is like Dial, with a caller provided Context.
// A nil Context will panic.
func DialContext(ctx context.Context, target string, opts ...DialOption) (conn *Conn, err error) {
	if ctx == nil {
		panic("nil Context")
	}

	c := &Conn{
		pending: make(map[uint64]*rpcCall),
		streams: make(map[string]*streamClients),
	}
	c.ctx, c.cancel = context.WithCancel(context.Background())

	defer func() {
		select {
		case <-ctx.Done():
			conn, err = nil, ctx.Err()
		default:
		}
		if err != nil {
			c.Close()
		}
	}()

	for _, o := range opts {
		o(&c.dialOpts)
	}

	netDial := c.dialOpts.dialer
	if netDial == nil {
		netDial = func(ctx context.Context, addr string) (net.Conn, error) {
			ws := &c.dialOpts.wsDialer

			if ws.WriteBufferSize == 0 {
				// Set the default size for use in writeLimiter.
				ws.WriteBufferSize = defaultWriteBufferSize
			}

			// Base writeLimit on the buffer size allocated by the
			// websocket package, we assume a message of this size
			// is fragmented and subtract one.
			writeLimit := ws.WriteBufferSize + websocketMaxHeaderSize - 1

			// Set NetDial to dial with context, this action will
			// override the HandshakeTimeout setting.
			ws.NetDial = func(network, addr string) (net.Conn, error) {
				conn, err := (&net.Dialer{}).DialContext(ctx, network, addr)
				return &writeLimiter{
					limit: writeLimit,
					Conn:  conn,
				}, err
			}

			wsConn, _, err := ws.Dial(addr, nil)
			if err != nil {
				return nil, err
			}

			if ws.EnableCompression {
				c.compressionLevel = wsConn.SetCompressionLevel
			}

			return &wsNetConn{conn: wsConn}, nil
		}
	}

	c.conn, err = netDial(ctx, target)
	if err != nil {
		return nil, err
	}
	newCodec := c.dialOpts.codec
	if newCodec == nil {
		newCodec = func(conn io.ReadWriter) Codec {
			return &jsonCodec{
				enc: json.NewEncoder(conn),
				dec: json.NewDecoder(conn),
			}
		}
	}
	c.codec = newCodec(c.conn)

	recvDone := func(err error) {
		// When we receive Inspector.detached the remote will close
		// the connection afterwards and recvDone will return. Maybe
		// we could give the user time to react to the event before
		// closing?
		// TODO(mafredri): Do we want to close here, like this?
		c.close(err)
	}
	go c.recv(c.notify, recvDone)

	return c, nil
}

// writeLimiter wraps a net.Conn and prevents writes of greater length
// than limit. Works around Chrome's lack of support for large or
// fragmented websocket messages and prevents sudden termination of the
// websocket connection. Gives the user an actionable error message when
// writes exceed limit.
type writeLimiter struct {
	limit int
	net.Conn
}

// BUG(mafredri): Chrome does not support websocket fragmentation
// (continuation messages) or messages that exceed 1MB in size.
// This limit was bumped in more recent versions of Chrome which can
// receive messages up to 100MB in size.
// See https://github.com/mafredri/cdp/issues/4 and
// https://github.com/ChromeDevTools/devtools-protocol/issues/24.
func (c *writeLimiter) Write(b []byte) (n int, err error) {
	if len(b) > c.limit {
		return 0, errors.New("rpcc: message too large (increase write buffer size or enable compression)")
	}
	return c.Conn.Write(b)
}
func (c *writeLimiter) Read(b []byte) (n int, err error) {
	return c.Conn.Read(b)
}

// Codec is used by recv and dispatcher to
// send and receive RPC communication.
type Codec interface {
	// WriteRequest encodes and writes the request onto the
	// underlying connection. Request is re-used between writes and
	// references to it should not be kept.
	WriteRequest(*Request) error
	// ReadResponse decodes a response from the underlying
	// connection. Response is re-used between reads and references
	// to it should not be kept.
{'fraction_non_alphanumeric': 0.10093473269924566, 'fraction_numerical': 0.005821580846179075, 'mean_word_length': 3.5443368107302535, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1619974', 'n_tokens_mistral': 4148, 'n_tokens_neox': 3783, 'n_words': 1367}