Modularize connection handling

service/metrics/metrics.go

@@ -53,8 +53,13 @@ func (c *measuredConn) Write(b []byte) (int, error) {
 	return n, err
 }
 
-func (c *measuredConn) ReadFrom(r io.Reader) (int64, error) {
-	n, err := io.Copy(c.StreamConn, r)
+func (c *measuredConn) ReadFrom(r io.Reader) (n int64, err error) {
+	if rf, ok := c.StreamConn.(io.ReaderFrom); ok {
+		// Prefer ReadFrom if we are calling ReadFrom. Otherwise io.Copy will try WriteTo first.
+		n, err = rf.ReadFrom(r)
+	} else {
+		n, err = io.Copy(c.StreamConn, r)
+	}
 	*c.writeCount += n
 	return n, err
 }

service/tcp.go

@@ -239,25 +239,22 @@ func (h *tcpHandler) Handle(ctx context.Context, clientConn transport.StreamConn
 	logger.Debugf("Done with status %v, duration %v", status, connDuration)
 }
 
-func (h *tcpHandler) handleConnection(ctx context.Context, listenerPort int, clientConn transport.StreamConn, proxyMetrics *metrics.ProxyMetrics) (string, *onet.ConnectionError) {
-	// Set a deadline to receive the address to the target.
-	clientConn.SetReadDeadline(time.Now().Add(h.readTimeout))
-
-	// 1. Find the cipher and acess key id.
+func (h *tcpHandler) authenticate(clientConn transport.StreamConn, proxyMetrics *metrics.ProxyMetrics) (string, transport.StreamConn, *onet.ConnectionError) {
+	// TODO(fortuna): Offer alternative transports.
+	// Find the cipher and acess key id.
 	cipherEntry, clientReader, clientSalt, timeToCipher, keyErr := findAccessKey(clientConn, remoteIP(clientConn), h.ciphers)
 	h.m.AddTCPCipherSearch(keyErr == nil, timeToCipher)
 	if keyErr != nil {
 		logger.Debugf("Failed to find a valid cipher after reading %v bytes: %v", proxyMetrics.ClientProxy, keyErr)
 		const status = "ERR_CIPHER"
-		h.absorbProbe(listenerPort, clientConn, status, proxyMetrics)
-		return "", onet.NewConnectionError(status, "Failed to find a valid cipher", keyErr)
+		return "", nil, onet.NewConnectionError(status, "Failed to find a valid cipher", keyErr)
 	}
 	var id string
 	if cipherEntry != nil {
 		id = cipherEntry.ID
 	}
 
-	// 2. Check if the connection is a replay.
+	// Check if the connection is a replay.
 	isServerSalt := cipherEntry.SaltGenerator.IsServerSalt(clientSalt)
 	// Only check the cache if findAccessKey succeeded and the salt is unrecognized.
 	if isServerSalt || !h.replayCache.Add(cipherEntry.ID, clientSalt) {
@@ -267,38 +264,39 @@ func (h *tcpHandler) handleConnection(ctx context.Context, listenerPort int, cli
 		} else {
 			status = "ERR_REPLAY_CLIENT"
 		}
-		h.absorbProbe(listenerPort, clientConn, status, proxyMetrics)
 		logger.Debugf(status+": %v sent %d bytes", clientConn.RemoteAddr(), proxyMetrics.ClientProxy)
-		return id, onet.NewConnectionError(status, "Replay detected", nil)
+		return id, nil, onet.NewConnectionError(status, "Replay detected", nil)
 	}
-
-	// 3. Read target address and dial it.
 	ssr := shadowsocks.NewReader(clientReader, cipherEntry.CryptoKey)
-	tgtAddr, err := socks.ReadAddr(ssr)
+	ssw := shadowsocks.NewWriter(clientConn, cipherEntry.CryptoKey)
+	ssw.SetSaltGenerator(cipherEntry.SaltGenerator)
+	return id, transport.WrapConn(clientConn, ssr, ssw), nil
+}
 
-	// Clear the deadline for the target address
-	clientConn.SetReadDeadline(time.Time{})
+func getProxyRequest(clientConn transport.StreamConn) (string, error) {
+	// TODO(fortuna): Use Shadowsocks proxy, HTTP CONNECT or SOCKS5 based on first byte:
+	// case 1, 3 or 4: Shadowsocks (address type)
+	// case 5: SOCKS5 (protocol version)
+	// case "C": HTTP CONNECT (first char of method)
+	tgtAddr, err := socks.ReadAddr(clientConn)
 	if err != nil {
-		// Drain to prevent a close on cipher error.
-		io.Copy(io.Discard, clientConn)
-		return id, onet.NewConnectionError("ERR_READ_ADDRESS", "Failed to get target address", err)
+		return "", err
 	}
-	tgtConn, dialErr := h.dialer.DialStream(ctx, tgtAddr.String())
+	return tgtAddr.String(), nil
+}
+
+func proxyConnection(ctx context.Context, dialer transport.StreamDialer, tgtAddr string, clientConn transport.StreamConn) *onet.ConnectionError {
+	tgtConn, dialErr := dialer.DialStream(ctx, tgtAddr)
 	if dialErr != nil {
 		// We don't drain so dial errors and invalid addresses are communicated quickly.
-		return id, ensureConnectionError(dialErr, "ERR_CONNECT", "Failed to connect to target")
+		return ensureConnectionError(dialErr, "ERR_CONNECT", "Failed to connect to target")
 	}
-	tgtConn = metrics.MeasureConn(tgtConn, &proxyMetrics.ProxyTarget, &proxyMetrics.TargetProxy)
 	defer tgtConn.Close()
-
-	// 4. Bridge the client and target connections
 	logger.Debugf("proxy %s <-> %s", clientConn.RemoteAddr().String(), tgtConn.RemoteAddr().String())
-	ssw := shadowsocks.NewWriter(clientConn, cipherEntry.CryptoKey)
-	ssw.SetSaltGenerator(cipherEntry.SaltGenerator)
 
 	fromClientErrCh := make(chan error)
 	go func() {
-		_, fromClientErr := ssr.WriteTo(tgtConn)
+		_, fromClientErr := io.Copy(tgtConn, clientConn)
 		if fromClientErr != nil {
 			// Drain to prevent a close in the case of a cipher error.
 			io.Copy(io.Discard, clientConn)
@@ -310,19 +308,58 @@ func (h *tcpHandler) handleConnection(ctx context.Context, listenerPort int, cli
 		tgtConn.CloseWrite()
 		fromClientErrCh <- fromClientErr
 	}()
-	_, fromTargetErr := ssw.ReadFrom(tgtConn)
+	_, fromTargetErr := io.Copy(clientConn, tgtConn)
 	// Send FIN to client.
 	clientConn.CloseWrite()
 	tgtConn.CloseRead()
 
 	fromClientErr := <-fromClientErrCh
 	if fromClientErr != nil {
-		return id, onet.NewConnectionError("ERR_RELAY_CLIENT", "Failed to relay traffic from client", fromClientErr)
+		return onet.NewConnectionError("ERR_RELAY_CLIENT", "Failed to relay traffic from client", fromClientErr)
 	}
 	if fromTargetErr != nil {
-		return id, onet.NewConnectionError("ERR_RELAY_TARGET", "Failed to relay traffic from target", fromTargetErr)
+		return onet.NewConnectionError("ERR_RELAY_TARGET", "Failed to relay traffic from target", fromTargetErr)
 	}
-	return id, nil
+	return nil
+}
+
+func (h *tcpHandler) handleConnection(ctx context.Context, listenerPort int, outerConn transport.StreamConn, proxyMetrics *metrics.ProxyMetrics) (string, *onet.ConnectionError) {
+	// Set a deadline to receive the address to the target.
+	readDeadline := time.Now().Add(h.readTimeout)
+	if deadline, ok := ctx.Deadline(); ok {
+		outerConn.SetDeadline(deadline)
+		if deadline.Before(readDeadline) {
+			readDeadline = deadline
+		}
+	}
+	outerConn.SetReadDeadline(readDeadline)
+
+	id, innerConn, authErr := h.authenticate(outerConn, proxyMetrics)
+	if authErr != nil {
+		// Drain to protect against probing attacks.
+		h.absorbProbe(listenerPort, outerConn, authErr.Status, proxyMetrics)
+		return id, authErr
+	}
+
+	// Read target address and dial it.
+	tgtAddr, err := getProxyRequest(innerConn)
+	// Clear the deadline for the target address
+	outerConn.SetReadDeadline(time.Time{})
+	if err != nil {
+		// Drain to prevent a close on cipher error.
+		io.Copy(io.Discard, outerConn)
+		return id, onet.NewConnectionError("ERR_READ_ADDRESS", "Failed to get target address", err)
+	}
+
+	dialer := transport.FuncStreamDialer(func(ctx context.Context, addr string) (transport.StreamConn, error) {
+		tgtConn, err := h.dialer.DialStream(ctx, tgtAddr)
+		if err != nil {
+			return nil, err
+		}
+		tgtConn = metrics.MeasureConn(tgtConn, &proxyMetrics.ProxyTarget, &proxyMetrics.TargetProxy)
+		return tgtConn, nil
+	})
+	return id, proxyConnection(ctx, dialer, tgtAddr, innerConn)
 }
 
 // Keep the connection open until we hit the authentication deadline to protect against probing attacks