refactor(model/netx.go): TLSHandhaker now returns a TLSConn (#1281)

I am making progress with ooni/probe#2531 and
I want to reactor model/netx.go such that the TLSHandshaker returns a
model.TLSConn rather than a net.Conn.

Returning a net.Conn and documenting it is a model.TLSConn is bad
compared to returning a model.TLSConn directly.

Note that we cannot apply the same transformation to netxlite's
TLSDialer.DialTLSContext because such a method must be assignable to
net/http and github.com/ooni/oohttp's Transport function also called
DialTLSContext.

The fact that we need code to be assignable to the Transport function is
what historically led the TLSHandshaker to return a net.Conn as well.
But it was quite clear from the get go that this choice led to some
quirks (and, in fact, this behavior was explicitly documented as such).

While there, slightly refactor `internal/experiment/echcheck/utls.go` to
avoid storing the conn inside the handshaker and make sure the test
coverage does not drop for this experiment.

While there, note that ooni/probe#2538 exists
and commit a mitigation.

Author: bassosimone

internal/dslx/tls.go

@@ -124,27 +124,22 @@ func (f *tlsHandshakeFunc) Apply(
 	defer cancel()
 
 	// handshake
-	conn, tlsState, err := handshaker.Handshake(ctx, input.Conn, config)
+	conn, err := handshaker.Handshake(ctx, input.Conn, config)
 
 	// possibly register established conn for late close
 	f.Pool.MaybeTrack(conn)
 
 	// stop the operation logger
 	ol.Stop(err)
 
-	var tlsConn netxlite.TLSConn
-	if conn != nil {
-		tlsConn = conn.(netxlite.TLSConn) // guaranteed to work
-	}
-
 	state := &TLSConnection{
 		Address:     input.Address,
-		Conn:        tlsConn, // possibly nil
+		Conn:        conn, // possibly nil
 		Domain:      input.Domain,
 		IDGenerator: input.IDGenerator,
 		Logger:      input.Logger,
 		Network:     input.Network,
-		TLSState:    tlsState,
+		TLSState:    netxlite.MaybeTLSConnectionState(conn),
 		Trace:       trace,
 		ZeroTime:    input.ZeroTime,
 	}

internal/dslx/tls_test.go

@@ -70,17 +70,22 @@ func TestTLSHandshake(t *testing.T) {
 				return nil
 			},
 		}
-		tlsConn := &mocks.TLSConn{Conn: tcpConn}
+		tlsConn := &mocks.TLSConn{
+			Conn: tcpConn,
+			MockConnectionState: func() tls.ConnectionState {
+				return tls.ConnectionState{}
+			},
+		}
 
 		eofHandshaker := &mocks.TLSHandshaker{
-			MockHandshake: func(ctx context.Context, conn net.Conn, config *tls.Config) (net.Conn, tls.ConnectionState, error) {
-				return nil, tls.ConnectionState{}, io.EOF
+			MockHandshake: func(ctx context.Context, conn net.Conn, config *tls.Config) (model.TLSConn, error) {
+				return nil, io.EOF
 			},
 		}
 
 		goodHandshaker := &mocks.TLSHandshaker{
-			MockHandshake: func(ctx context.Context, conn net.Conn, config *tls.Config) (net.Conn, tls.ConnectionState, error) {
-				return tlsConn, tls.ConnectionState{}, nil
+			MockHandshake: func(ctx context.Context, conn net.Conn, config *tls.Config) (model.TLSConn, error) {
+				return tlsConn, nil
 			},
 		}
 

internal/experiment/echcheck/handshake.go

@@ -41,9 +41,10 @@ func handshakeWithExtension(ctx context.Context, conn net.Conn, zeroTime time.Ti
 	tracedHandshaker := handshakerConstructor(log.Log, &utls.HelloFirefox_Auto)
 
 	start := time.Now()
-	_, connState, err := tracedHandshaker.Handshake(ctx, conn, tlsConfig)
+	maybeTLSConn, err := tracedHandshaker.Handshake(ctx, conn, tlsConfig)
 	finish := time.Now()
 
+	connState := netxlite.MaybeTLSConnectionState(maybeTLSConn)
 	return measurexlite.NewArchivalTLSOrQUICHandshakeResult(0, start.Sub(zeroTime), "tcp", address, tlsConfig,
 		connState, err, finish.Sub(zeroTime))
 }

internal/experiment/echcheck/measure_test.go

@@ -84,12 +84,10 @@ func TestMeasurementSuccess(t *testing.T) {
 	}
 
 	summary, err := measurer.GetSummaryKeys(&model.Measurement{})
-
+	if err != nil {
+		t.Fatal(err)
+	}
 	if summary.(SummaryKeys).IsAnomaly != false {
 		t.Fatal("expected false")
 	}
 }
-
-func newsession() model.ExperimentSession {
-	return &mockable.Session{MockableLogger: log.Log}
-}

internal/experiment/echcheck/utls.go

@@ -12,7 +12,6 @@ import (
 )
 
 type tlsHandshakerWithExtensions struct {
-	conn       *netxlite.UTLSConn
 	extensions []utls.TLSExtension
 	dl         model.DebugLogger
 	id         *utls.ClientHelloID
@@ -32,18 +31,19 @@ func newHandshakerWithExtensions(extensions []utls.TLSExtension) func(dl model.D
 	}
 }
 
-func (t *tlsHandshakerWithExtensions) Handshake(ctx context.Context, conn net.Conn, tlsConfig *tls.Config) (
-	net.Conn, tls.ConnectionState, error) {
-	var err error
-	t.conn, err = netxlite.NewUTLSConn(conn, tlsConfig, t.id)
+func (t *tlsHandshakerWithExtensions) Handshake(
+	ctx context.Context, tcpConn net.Conn, tlsConfig *tls.Config) (model.TLSConn, error) {
+	tlsConn, err := netxlite.NewUTLSConn(tcpConn, tlsConfig, t.id)
 	runtimex.Assert(err == nil, "unexpected error when creating UTLSConn")
 
 	if t.extensions != nil && len(t.extensions) != 0 {
-		t.conn.BuildHandshakeState()
-		t.conn.Extensions = append(t.conn.Extensions, t.extensions...)
+		tlsConn.BuildHandshakeState()
+		tlsConn.Extensions = append(tlsConn.Extensions, t.extensions...)
 	}
 
-	err = t.conn.Handshake()
+	if err := tlsConn.Handshake(); err != nil {
+		return nil, err
+	}
 
-	return t.conn.NetConn(), t.conn.ConnectionState(), err
+	return tlsConn, nil
 }

internal/experiment/echcheck/utls_test.go

@@ -0,0 +1,41 @@
+package echcheck
+
+import (
+	"context"
+	"crypto/tls"
+	"errors"
+	"testing"
+
+	"github.com/ooni/probe-cli/v3/internal/mocks"
+	"github.com/ooni/probe-cli/v3/internal/model"
+	utls "gitlab.com/yawning/utls.git"
+)
+
+func TestTLSHandshakerWithExtension(t *testing.T) {
+	t.Run("when the TLS handshake fails", func(t *testing.T) {
+		thx := &tlsHandshakerWithExtensions{
+			extensions: []utls.TLSExtension{},
+			dl:         model.DiscardLogger,
+			id:         &utls.HelloChrome_70,
+		}
+
+		expected := errors.New("mocked error")
+		tcpConn := &mocks.Conn{
+			MockWrite: func(b []byte) (int, error) {
+				return 0, expected
+			},
+		}
+
+		tlsConfig := &tls.Config{
+			InsecureSkipVerify: true,
+		}
+
+		tlsConn, err := thx.Handshake(context.Background(), tcpConn, tlsConfig)
+		if !errors.Is(err, expected) {
+			t.Fatal(err)
+		}
+		if tlsConn != nil {
+			t.Fatal("expected nil tls conn")
+		}
+	})
+}

internal/experiment/tlsmiddlebox/tracing.go

@@ -97,7 +97,7 @@ func (m *Measurer) handshakeWithTTL(ctx context.Context, index int64, zeroTime t
 	if clientId > 0 {
 		thx = trace.NewTLSHandshakerUTLS(logger, ClientIDs[clientId])
 	}
-	_, _, err = thx.Handshake(ctx, conn, genTLSConfig(sni))
+	_, err = thx.Handshake(ctx, conn, genTLSConfig(sni))
 	ol.Stop(err)
 	soErr := extractSoError(conn)
 	// 4. reset the TTL value to ensure that conn closes successfully

internal/experiment/tlsping/tlsping.go

@@ -189,7 +189,7 @@ func (m *Measurer) tlsConnectAndHandshake(ctx context.Context, index int64,
 		RootCAs:    nil,
 		ServerName: sni,
 	}
-	_, _, err = thx.Handshake(ctx, conn, config)
+	_, err = thx.Handshake(ctx, conn, config)
 	ol.Stop(err)
 	sp.TLSHandshake = trace.FirstTLSHandshakeOrNil() // record the first handshake from the buffer
 	sp.NetworkEvents = trace.NetworkEvents()

internal/experiment/webconnectivitylte/secureflow.go

@@ -154,14 +154,15 @@ func (t *SecureFlow) Run(parentCtx context.Context, index int64) error {
 	const tlsTimeout = 10 * time.Second
 	tlsCtx, tlsCancel := context.WithTimeout(parentCtx, tlsTimeout)
 	defer tlsCancel()
-	tlsConn, tlsConnState, err := tlsHandshaker.Handshake(tlsCtx, tcpConn, tlsConfig)
+	tlsConn, err := tlsHandshaker.Handshake(tlsCtx, tcpConn, tlsConfig)
 	t.TestKeys.AppendTLSHandshakes(trace.TLSHandshakes()...)
 	if err != nil {
 		ol.Stop(err)
 		return err
 	}
 	defer tlsConn.Close()
 
+	tlsConnState := netxlite.MaybeTLSConnectionState(tlsConn)
 	alpn := tlsConnState.NegotiatedProtocol
 
 	// Determine whether we're allowed to fetch the webpage
@@ -177,8 +178,7 @@ func (t *SecureFlow) Run(parentCtx context.Context, index int64) error {
 	httpTransport := netxlite.NewHTTPTransport(
 		t.Logger,
 		netxlite.NewNullDialer(),
-		// note: netxlite guarantees that here tlsConn is a netxlite.TLSConn
-		netxlite.NewSingleUseTLSDialer(tlsConn.(netxlite.TLSConn)),
+		netxlite.NewSingleUseTLSDialer(tlsConn),
 	)
 
 	// create HTTP request

internal/legacy/measurex/measurer.go

@@ -355,13 +355,12 @@ func (mx *Measurer) TLSConnectAndHandshakeWithDB(ctx context.Context,
 	ctx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 	th := mx.WrapTLSHandshaker(db, mx.TLSHandshaker)
-	tlsConn, _, err := th.Handshake(ctx, conn, config)
+	tlsConn, err := th.Handshake(ctx, conn, config)
 	ol.Stop(err)
 	if err != nil {
 		return nil, err
 	}
-	// cast safe according to the docs of netxlite's handshaker
-	return tlsConn.(netxlite.TLSConn), nil
+	return tlsConn, nil
 }
 
 // QUICHandshake connects and TLS handshakes with a QUIC endpoint.

internal/legacy/measurex/tls.go

@@ -11,7 +11,6 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"errors"
-	"net"
 	"time"
 
 	"github.com/ooni/probe-cli/v3/internal/model"
@@ -53,13 +52,13 @@ type QUICTLSHandshakeEvent struct {
 	Started         float64
 }
 
-func (thx *tlsHandshakerDB) Handshake(ctx context.Context,
-	conn Conn, config *tls.Config) (net.Conn, tls.ConnectionState, error) {
+func (thx *tlsHandshakerDB) Handshake(ctx context.Context, conn Conn, config *tls.Config) (model.TLSConn, error) {
 	network := conn.RemoteAddr().Network()
 	remoteAddr := conn.RemoteAddr().String()
 	started := time.Since(thx.begin).Seconds()
-	tconn, state, err := thx.TLSHandshaker.Handshake(ctx, conn, config)
+	tconn, err := thx.TLSHandshaker.Handshake(ctx, conn, config)
 	finished := time.Since(thx.begin).Seconds()
+	tstate := netxlite.MaybeTLSConnectionState(tconn)
 	thx.db.InsertIntoTLSHandshake(&QUICTLSHandshakeEvent{
 		Network:         network,
 		RemoteAddr:      remoteAddr,
@@ -70,12 +69,12 @@ func (thx *tlsHandshakerDB) Handshake(ctx context.Context,
 		Finished:        finished,
 		Failure:         NewFailure(err),
 		Oddity:          thx.computeOddity(err),
-		TLSVersion:      netxlite.TLSVersionString(state.Version),
-		CipherSuite:     netxlite.TLSCipherSuiteString(state.CipherSuite),
-		NegotiatedProto: state.NegotiatedProtocol,
-		PeerCerts:       peerCerts(err, &state),
+		TLSVersion:      netxlite.TLSVersionString(tstate.Version),
+		CipherSuite:     netxlite.TLSCipherSuiteString(tstate.CipherSuite),
+		NegotiatedProto: tstate.NegotiatedProtocol,
+		PeerCerts:       peerCerts(err, &tstate),
 	})
-	return tconn, state, err
+	return tconn, err
 }
 
 func (thx *tlsHandshakerDB) computeOddity(err error) Oddity {

internal/legacy/tracex/tls.go

@@ -42,7 +42,7 @@ func (s *Saver) WrapTLSHandshaker(thx model.TLSHandshaker) model.TLSHandshaker {
 
 // Handshake implements model.TLSHandshaker.Handshake
 func (h *TLSHandshakerSaver) Handshake(
-	ctx context.Context, conn net.Conn, config *tls.Config) (net.Conn, tls.ConnectionState, error) {
+	ctx context.Context, conn net.Conn, config *tls.Config) (model.TLSConn, error) {
 	proto := conn.RemoteAddr().Network()
 	remoteAddr := conn.RemoteAddr().String()
 	start := time.Now()
@@ -54,23 +54,24 @@ func (h *TLSHandshakerSaver) Handshake(
 		TLSServerName: config.ServerName,
 		Time:          start,
 	}})
-	tlsconn, state, err := h.TLSHandshaker.Handshake(ctx, conn, config)
+	tlsconn, err := h.TLSHandshaker.Handshake(ctx, conn, config)
 	stop := time.Now()
+	tstate := netxlite.MaybeTLSConnectionState(tlsconn)
 	h.Saver.Write(&EventTLSHandshakeDone{&EventValue{
 		Address:            remoteAddr,
 		Duration:           stop.Sub(start),
 		Err:                NewFailureStr(err),
 		NoTLSVerify:        config.InsecureSkipVerify,
 		Proto:              proto,
-		TLSCipherSuite:     netxlite.TLSCipherSuiteString(state.CipherSuite),
-		TLSNegotiatedProto: state.NegotiatedProtocol,
+		TLSCipherSuite:     netxlite.TLSCipherSuiteString(tstate.CipherSuite),
+		TLSNegotiatedProto: tstate.NegotiatedProtocol,
 		TLSNextProtos:      config.NextProtos,
-		TLSPeerCerts:       tlsPeerCerts(state, err),
+		TLSPeerCerts:       tlsPeerCerts(tstate, err),
 		TLSServerName:      config.ServerName,
-		TLSVersion:         netxlite.TLSVersionString(state.Version),
+		TLSVersion:         netxlite.TLSVersionString(tstate.Version),
 		Time:               stop,
 	}})
-	return tlsconn, state, err
+	return tlsconn, err
 }
 
 var _ model.TLSHandshaker = &TLSHandshakerSaver{}

internal/legacy/tracex/tls_test.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/ooni/probe-cli/v3/internal/mocks"
+	"github.com/ooni/probe-cli/v3/internal/model"
 )
 
 func TestWrapTLSHandshaker(t *testing.T) {
@@ -98,9 +99,8 @@ func TestTLSHandshakerSaver(t *testing.T) {
 				},
 			}
 			thx := saver.WrapTLSHandshaker(&mocks.TLSHandshaker{
-				MockHandshake: func(ctx context.Context, conn net.Conn,
-					config *tls.Config) (net.Conn, tls.ConnectionState, error) {
-					return returnedConn, returnedConnState, nil
+				MockHandshake: func(ctx context.Context, conn net.Conn, config *tls.Config) (model.TLSConn, error) {
+					return returnedConn, nil
 				},
 			})
 			ctx := context.Background()
@@ -121,7 +121,7 @@ func TestTLSHandshakerSaver(t *testing.T) {
 					}
 				},
 			}
-			conn, _, err := thx.Handshake(ctx, tcpConn, tlsConfig)
+			conn, err := thx.Handshake(ctx, tcpConn, tlsConfig)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -161,9 +161,8 @@ func TestTLSHandshakerSaver(t *testing.T) {
 			expected := errors.New("mocked error")
 			saver := &Saver{}
 			thx := saver.WrapTLSHandshaker(&mocks.TLSHandshaker{
-				MockHandshake: func(ctx context.Context, conn net.Conn,
-					config *tls.Config) (net.Conn, tls.ConnectionState, error) {
-					return nil, tls.ConnectionState{}, expected
+				MockHandshake: func(ctx context.Context, conn net.Conn, config *tls.Config) (model.TLSConn, error) {
+					return nil, expected
 				},
 			})
 			ctx := context.Background()
@@ -184,7 +183,7 @@ func TestTLSHandshakerSaver(t *testing.T) {
 					}
 				},
 			}
-			conn, _, err := thx.Handshake(ctx, tcpConn, tlsConfig)
+			conn, err := thx.Handshake(ctx, tcpConn, tlsConfig)
 			if !errors.Is(err, expected) {
 				t.Fatal("unexpected err", err)
 			}

internal/measurexlite/tls.go

@@ -35,7 +35,7 @@ var _ model.TLSHandshaker = &tlsHandshakerTrace{}
 
 // Handshake implements model.TLSHandshaker.Handshake.
 func (thx *tlsHandshakerTrace) Handshake(
-	ctx context.Context, conn net.Conn, tlsConfig *tls.Config) (net.Conn, tls.ConnectionState, error) {
+	ctx context.Context, conn net.Conn, tlsConfig *tls.Config) (model.TLSConn, error) {
 	return thx.thx.Handshake(netxlite.ContextWithTrace(ctx, thx.tx), conn, tlsConfig)
 }