<commit_before>package rpc

import (
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"github.com/keybase/backoff"
	"golang.org/x/net/context"
)

// DisconnectStatus is the connection information passed to
// ConnectionHandler.OnDisconnected().
type DisconnectStatus int

const (
	// skip 0
	_ = iota
	// UsingExistingConnection means that an existing
	// connection will be used.
	UsingExistingConnection DisconnectStatus = iota
	// StartingFirstConnection means that a connection will be
	// started, and this is the first one.
	StartingFirstConnection
	// StartingNonFirstConnection means that a connection will be
	// started, and this is not the first one.
	StartingNonFirstConnection
)

// ConnectionTransport is a container for an underlying transport to be
// used by a Connection instance.
type ConnectionTransport interface {
	// Dial is called to connect to the server.
	Dial(ctx context.Context) (Transporter, error)

	// IsConnected is called to check for connection status.
	IsConnected() bool

	// Finalize is used to indicate the result of Dial is complete.
	Finalize()

	// Close is used to close any open connection.
	Close()
}

type connTransport struct {
	uri             *FMPURI
	l               LogFactory
	wef             WrapErrorFunc
	conn            net.Conn
	transport       Transporter
	stagedTransport Transporter
}

var _ ConnectionTransport = (*connTransport)(nil)

// NewConnectionTransport creates a ConnectionTransport for a given FMPURI.
func NewConnectionTransport(uri *FMPURI, l LogFactory, wef WrapErrorFunc) ConnectionTransport {
	return &connTransport{
		uri: uri,
		l:   l,
		wef: wef,
	}
}

func (t *connTransport) Dial(context.Context) (Transporter, error) {
	var err error
	t.conn, err = t.uri.Dial()
	if err != nil {
		return nil, err
	}

	// Disable SIGPIPE on platforms that require it (Darwin). See sigpipe_bsd.go.
	if err = DisableSigPipe(t.conn); err != nil {
		return nil, err
	}

	t.stagedTransport = NewTransport(t.conn, t.l, t.wef)
	return t.stagedTransport, nil
}

func (t *connTransport) IsConnected() bool {
	return t.transport != nil && t.transport.IsConnected()
}

func (t *connTransport) Finalize() {
	t.transport = t.stagedTransport
	t.stagedTransport = nil
}

func (t *connTransport) Close() {
	t.conn.Close()
	t.transport = nil
	t.stagedTransport = nil
}

// ConnectionHandler is the callback interface for interacting with the connection.
type ConnectionHandler interface {
	// OnConnect is called immediately after a connection has been
	// established.  An implementation would likely log something,
	// register served protocols, and/or perform authentication.
	OnConnect(context.Context, *Connection, GenericClient, *Server) error

	// OnConnectError is called whenever there is an error during connection.
	OnConnectError(err error, reconnectThrottleDuration time.Duration)

	// OnDoCommandError is called whenever there is an error during DoCommand
	OnDoCommandError(err error, nextTime time.Duration)

	// OnDisconnected is called whenever the connection notices it
	// is disconnected.
	OnDisconnected(ctx context.Context, status DisconnectStatus)

	// ShouldRetry is called whenever an error is returned by
	// an RPC function passed to Connection.DoCommand(), and
	// should return whether or not that error signifies that that
	// RPC should retried (with backoff)
	ShouldRetry(name string, err error) bool

	// ShouldRetryOnConnect is called whenever an error is returned
	// during connection establishment, and should return whether or
	// not the connection should be established again.
	ShouldRetryOnConnect(err error) bool

	// HandlerName returns a string representing the type of the connection
	// handler.
	HandlerName() string
}

// ConnectionTransportTLS is a ConnectionTransport implementation that
// uses TLS+rpc.
type ConnectionTransportTLS struct {
	rootCerts []byte
	srvAddr   string
	tlsConfig *tls.Config

	// Protects everything below.
	mutex           sync.Mutex
	transport       Transporter
	stagedTransport Transporter
	conn            net.Conn
	logFactory      LogFactory
	wef             WrapErrorFunc
}

// Test that ConnectionTransportTLS fully implements the ConnectionTransport interface.
var _ ConnectionTransport = (*ConnectionTransportTLS)(nil)

// Dial is an implementation of the ConnectionTransport interface.
func (ct *ConnectionTransportTLS) Dial(ctx context.Context) (
	Transporter, error) {
	var conn net.Conn
	err := runUnlessCanceled(ctx, func() error {
		config := ct.tlsConfig
		host, _, err := net.SplitHostPort(ct.srvAddr)
		if err != nil {
			return err
		}

		// If we didn't specify a tls.Config, but we did specify
		// explicit rootCerts, then populate a new tls.Config here.
		// Otherwise, we're using the defaults via `nil` tls.Config.
		if config == nil && ct.rootCerts != nil {
			// load CA certificate
			certs := x509.NewCertPool()
			if !certs.AppendCertsFromPEM(ct.rootCerts) {
				return errors.New("Unable to load root certificates")
			}
			config = &tls.Config{
				RootCAs:    certs,
				ServerName: host,
			}
		}
		// Final check to make sure we have a TLS config since tls.Client requires
		// either ServerName or InsecureSkipVerify to be set.
		if config == nil {
			config = &tls.Config{ServerName: host}
		}

		// connect
		dialer := net.Dialer{
			KeepAlive: 10 * time.Second,
		}
		baseConn, err := dialer.Dial("tcp", ct.srvAddr)
		if err != nil {
			return err
		}
		conn = tls.Client(baseConn, config)

		// Disable SIGPIPE on platforms that require it (Darwin). See sigpipe_bsd.go.
		return DisableSigPipe(baseConn)
	})
	if err != nil {
		return nil, err
	}

	ct.mutex.Lock()
	defer ct.mutex.Unlock()
	transport := NewTransport(conn, ct.logFactory, ct.wef)
	ct.conn = conn
	ct.stagedTransport = transport<commit_msg>Fix sporadic connection hangs<commit_after>package rpc

import (
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"github.com/keybase/backoff"
	"golang.org/x/net/context"
)

// DisconnectStatus is the connection information passed to
// ConnectionHandler.OnDisconnected().
type DisconnectStatus int

const (
	// skip 0
	_ = iota
	// UsingExistingConnection means that an existing
	// connection will be used.
	UsingExistingConnection DisconnectStatus = iota
	// StartingFirstConnection means that a connection will be
	// started, and this is the first one.
	StartingFirstConnection
	// StartingNonFirstConnection means that a connection will be
	// started, and this is not the first one.
	StartingNonFirstConnection
)

// ConnectionTransport is a container for an underlying transport to be
// used by a Connection instance.
type ConnectionTransport interface {
	// Dial is called to connect to the server.
	Dial(ctx context.Context) (Transporter, error)

	// IsConnected is called to check for connection status.
	IsConnected() bool

	// Finalize is used to indicate the result of Dial is complete.
	Finalize()

	// Close is used to close any open connection.
	Close()
}

type connTransport struct {
	uri             *FMPURI
	l               LogFactory
	wef             WrapErrorFunc
	conn            net.Conn
	transport       Transporter
	stagedTransport Transporter
}

var _ ConnectionTransport = (*connTransport)(nil)

// NewConnectionTransport creates a ConnectionTransport for a given FMPURI.
func NewConnectionTransport(uri *FMPURI, l LogFactory, wef WrapErrorFunc) ConnectionTransport {
	return &connTransport{
		uri: uri,
		l:   l,
		wef: wef,
	}
}

func (t *connTransport) Dial(context.Context) (Transporter, error) {
	var err error
	t.conn, err = t.uri.Dial()
	if err != nil {
		return nil, err
	}

	// Disable SIGPIPE on platforms that require it (Darwin). See sigpipe_bsd.go.
	if err = DisableSigPipe(t.conn); err != nil {
		return nil, err
	}

	t.stagedTransport = NewTransport(t.conn, t.l, t.wef)
	return t.stagedTransport, nil
}

func (t *connTransport) IsConnected() bool {
	return t.transport != nil && t.transport.IsConnected()
}

func (t *connTransport) Finalize() {
	t.transport = t.stagedTransport
	t.stagedTransport = nil
}

func (t *connTransport) Close() {
	t.conn.Close()
	t.transport = nil
	t.stagedTransport = nil
}

// ConnectionHandler is the callback interface for interacting with the connection.
type ConnectionHandler interface {
	// OnConnect is called immediately after a connection has been
	// established.  An implementation would likely log something,
	// register served protocols, and/or perform authentication.
	OnConnect(context.Context, *Connection, GenericClient, *Server) error

	// OnConnectError is called whenever there is an error during connection.
	OnConnectError(err error, reconnectThrottleDuration time.Duration)

	// OnDoCommandError is called whenever there is an error during DoCommand
	OnDoCommandError(err error, nextTime time.Duration)

	// OnDisconnected is called whenever the connection notices it
	// is disconnected.
	OnDisconnected(ctx context.Context, status DisconnectStatus)

	// ShouldRetry is called whenever an error is returned by
	// an RPC function passed to Connection.DoCommand(), and
	// should return whether or not that error signifies that that
	// RPC should retried (with backoff)
	ShouldRetry(name string, err error) bool

	// ShouldRetryOnConnect is called whenever an error is returned
	// during connection establishment, and should return whether or
	// not the connection should be established again.
	ShouldRetryOnConnect(err error) bool

	// HandlerName returns a string representing the type of the connection
	// handler.
	HandlerName() string
}

// ConnectionTransportTLS is a ConnectionTransport implementation that
// uses TLS+rpc.
type ConnectionTransportTLS struct {
	rootCerts []byte
	srvAddr   string
	tlsConfig *tls.Config

	// Protects everything below.
	mutex           sync.Mutex
	transport       Transporter
	stagedTransport Transporter
	conn            net.Conn
	logFactory      LogFactory
	wef             WrapErrorFunc
}

// Test that ConnectionTransportTLS fully implements the ConnectionTransport interface.
var _ ConnectionTransport = (*ConnectionTransportTLS)(nil)

// Dial is an implementation of the ConnectionTransport interface.
func (ct *ConnectionTransportTLS) Dial(ctx context.Context) (
	Transporter, error) {
	var conn net.Conn
	err := runUnlessCanceled(ctx, func() error {
		config := ct.tlsConfig
		host, _, err := net.SplitHostPort(ct.srvAddr)
		if err != nil {
			return err
		}

		// If we didn't specify a tls.Config, but we did specify
		// explicit rootCerts, then populate a new tls.Config here.
		// Otherwise, we're using the defaults via `nil` tls.Config.
		if config == nil && ct.rootCerts != nil {
			// load CA certificate
			certs := x509.NewCertPool()
			if !certs.AppendCertsFromPEM(ct.rootCerts) {
				return errors.New("Unable to load root certificates")
			}
			config = &tls.Config{
				RootCAs:    certs,
				ServerName: host,
			}
		}
		// Final check to make sure we have a TLS config since tls.Client requires
		// either ServerName or InsecureSkipVerify to be set.
		if config == nil {
			config = &tls.Config{ServerName: host}
		}

		// connect
		dialer := net.Dialer{
			KeepAlive: 10 * time.Second,
		}
		baseConn, err := dialer.Dial("tcp", ct.srvAddr)
		if err != nil {
			return err
		}
		conn = tls.Client(baseConn, config)
		if err := conn.(*tls.Conn).Handshake(); err != nil {
			return err
		}

		// Disable SIGPIPE on platforms that require it (Darwin). See sigpipe_bsd.go.
		return DisableSigPipe(baseConn)
	})
	if err != nil {
		return nil, err
	}

	ct.mutex.Lock()
	defer ct.mutex.Unlock()
	transport := NewTransport(conn, ct.logFactory, ct.wef)
	ct.conn = conn
	ct.stagedTransport = transport
{'fraction_non_alphanumeric': 0.08682660291983267, 'fraction_numerical': 0.0015367540339793393, 'mean_word_length': 3.609996064541519, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, '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': '5585718', 'n_tokens_mistral': 3662, 'n_tokens_neox': 3447, 'n_words': 1207}