From 7e47667b0806d308c8fea021d5cdc49def13584c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Oliveirinha?= <joliveirinha@cloudflare.com>
Date: Tue, 30 Nov 2021 10:27:33 +0000
Subject: [PATCH] TUN-5481: Create abstraction for Origin UDP Connection

Creates an abstraction over UDP Conn for origin "connection" which can
be useful for future support of complex protocols that may require
changing ports during protocol negotiation (eg. SIP, TFTP)
In addition, it removes a dependency from ingress on connection package.
---
 connection/quic.go                | 12 ++----------
 ingress/origin_connection_test.go | 21 ++++++++++++++++++---
 ingress/origin_udp_proxy.go       | 27 +++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 13 deletions(-)
 create mode 100644 ingress/origin_udp_proxy.go

diff --git a/connection/quic.go b/connection/quic.go
index 933cb197..190173f4 100644
--- a/connection/quic.go
+++ b/connection/quic.go
@@ -17,6 +17,7 @@ import (
 	"golang.org/x/sync/errgroup"
 
 	"github.com/cloudflare/cloudflared/datagramsession"
+	"github.com/cloudflare/cloudflared/ingress"
 	quicpogs "github.com/cloudflare/cloudflared/quic"
 	tunnelpogs "github.com/cloudflare/cloudflared/tunnelrpc/pogs"
 )
@@ -176,7 +177,7 @@ func (q *QUICConnection) handleRPCStream(rpcStream *quicpogs.RPCServerStream) er
 func (q *QUICConnection) RegisterUdpSession(ctx context.Context, sessionID uuid.UUID, dstIP net.IP, dstPort uint16) error {
 	// Each session is a series of datagram from an eyeball to a dstIP:dstPort.
 	// (src port, dst IP, dst port) uniquely identifies a session, so it needs a dedicated connected socket.
-	originProxy, err := q.newUDPProxy(dstIP, dstPort)
+	originProxy, err := ingress.DialUDP(dstIP, dstPort)
 	if err != nil {
 		q.logger.Err(err).Msgf("Failed to create udp proxy to %s:%d", dstIP, dstPort)
 		return err
@@ -292,15 +293,6 @@ func isTransferEncodingChunked(req *http.Request) bool {
 	return strings.Contains(strings.ToLower(transferEncodingVal), "chunked")
 }
 
-// TODO: TUN-5303: Define an UDPProxy in ingress package
-func (q *QUICConnection) newUDPProxy(dstIP net.IP, dstPort uint16) (*net.UDPConn, error) {
-	dstAddr := &net.UDPAddr{
-		IP:   dstIP,
-		Port: int(dstPort),
-	}
-	return net.DialUDP("udp", nil, dstAddr)
-}
-
 // TODO: TUN-5303: Find the local IP once in ingress package
 // TODO: TUN-5421 allow user to specify which IP to bind to
 func getLocalIP() (net.IP, error) {
diff --git a/ingress/origin_connection_test.go b/ingress/origin_connection_test.go
index e89201d7..07bcb500 100644
--- a/ingress/origin_connection_test.go
+++ b/ingress/origin_connection_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net"
 	"net/http"
@@ -19,7 +20,6 @@ import (
 	"golang.org/x/net/proxy"
 	"golang.org/x/sync/errgroup"
 
-	"github.com/cloudflare/cloudflared/connection"
 	"github.com/cloudflare/cloudflared/logger"
 	"github.com/cloudflare/cloudflared/socks"
 	"github.com/cloudflare/cloudflared/websocket"
@@ -192,8 +192,10 @@ func TestSocksStreamWSOverTCPConnection(t *testing.T) {
 
 func TestWsConnReturnsBeforeStreamReturns(t *testing.T) {
 	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		eyeballConn, err := connection.NewHTTP2RespWriter(r, w, connection.TypeWebsocket)
-		assert.NoError(t, err)
+		eyeballConn := &readWriter{
+			w: w,
+			r: r.Body,
+		}
 
 		cfdConn, originConn := net.Pipe()
 		tcpOverWSConn := tcpOverWSConnection{
@@ -319,3 +321,16 @@ func echoTCPOrigin(t *testing.T, conn net.Conn) {
 	_, err = conn.Write(testResponse)
 	assert.NoError(t, err)
 }
+
+type readWriter struct {
+	w io.Writer
+	r io.Reader
+}
+
+func (r *readWriter) Read(p []byte) (n int, err error) {
+	return r.r.Read(p)
+}
+
+func (r *readWriter) Write(p []byte) (n int, err error) {
+	return r.w.Write(p)
+}
diff --git a/ingress/origin_udp_proxy.go b/ingress/origin_udp_proxy.go
new file mode 100644
index 00000000..3f2b86c7
--- /dev/null
+++ b/ingress/origin_udp_proxy.go
@@ -0,0 +1,27 @@
+package ingress
+
+import (
+	"fmt"
+	"io"
+	"net"
+)
+
+type UDPProxy struct {
+	io.ReadWriteCloser
+}
+
+func DialUDP(dstIP net.IP, dstPort uint16) (*UDPProxy, error) {
+	dstAddr := &net.UDPAddr{
+		IP:   dstIP,
+		Port: int(dstPort),
+	}
+
+	// We use nil as local addr to force runtime to find the best suitable local address IP given the destination
+	// address as context.
+	udpConn, err := net.DialUDP("udp", nil, dstAddr)
+	if err != nil {
+		return nil, fmt.Errorf("unable to create UDP proxy to origin (%v:%v): %w", dstIP, dstPort, err)
+	}
+
+	return &UDPProxy{udpConn}, nil
+}