TUN-5300: Define RPC to register UDP sessions

connection/quic.go

@@ -10,6 +10,7 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/google/uuid"
 	"github.com/lucas-clemente/quic-go"
 	"github.com/pkg/errors"
 	"github.com/rs/zerolog"
@@ -34,6 +35,7 @@ type QUICConnection struct {
 	httpProxy         OriginProxy
 	gracefulShutdownC <-chan struct{}
 	stoppedGracefully bool
+	udpSessions       *udpSessions
 }
 
 // NewQUICConnection returns a new instance of QUICConnection.
@@ -64,9 +66,10 @@ func NewQUICConnection(
 	}
 
 	return &QUICConnection{
-		session:   session,
-		httpProxy: httpProxy,
-		logger:    observer.log,
+		session:     session,
+		httpProxy:   httpProxy,
+		logger:      observer.log,
+		udpSessions: newUDPSessions(),
 	}, nil
 }
 
@@ -99,7 +102,30 @@ func (q *QUICConnection) Close() {
 }
 
 func (q *QUICConnection) handleStream(stream quic.Stream) error {
-	connectRequest, err := quicpogs.ReadConnectRequestData(stream)
+	signature, err := quicpogs.DetermineProtocol(stream)
+	if err != nil {
+		return err
+	}
+	switch signature {
+	case quicpogs.DataStreamProtocolSignature:
+		reqServerStream, err := quicpogs.NewRequestServerStream(stream, signature)
+		if err != nil {
+			return nil
+		}
+		return q.handleDataStream(reqServerStream)
+	case quicpogs.RPCStreamProtocolSignature:
+		rpcStream, err := quicpogs.NewRPCServerStream(stream, signature)
+		if err != nil {
+			return err
+		}
+		return q.handleRPCStream(rpcStream)
+	default:
+		return fmt.Errorf("Unknown protocol %v", signature)
+	}
+}
+
+func (q *QUICConnection) handleDataStream(stream *quicpogs.RequestServerStream) error {
+	connectRequest, err := stream.ReadConnectRequestData()
 	if err != nil {
 		return err
 	}
@@ -114,32 +140,38 @@ func (q *QUICConnection) handleStream(stream quic.Stream) error {
 		w := newHTTPResponseAdapter(stream)
 		return q.httpProxy.ProxyHTTP(w, req, connectRequest.Type == quicpogs.ConnectionTypeWebsocket)
 	case quicpogs.ConnectionTypeTCP:
-		rwa := &streamReadWriteAcker{
-			ReadWriter: stream,
-		}
+		rwa := &streamReadWriteAcker{stream}
 		return q.httpProxy.ProxyTCP(context.Background(), rwa, &TCPRequest{Dest: connectRequest.Dest})
 	}
 	return nil
 }
 
+func (q *QUICConnection) handleRPCStream(rpcStream *quicpogs.RPCServerStream) error {
+	return rpcStream.Serve(q, q.logger)
+}
+
+func (q *QUICConnection) RegisterUdpSession(ctx context.Context, sessionID uuid.UUID, dstIP net.IP, dstPort uint16) error {
+	return q.udpSessions.register(sessionID, dstIP, dstPort)
+}
+
 // streamReadWriteAcker is a light wrapper over QUIC streams with a callback to send response back to
 // the client.
 type streamReadWriteAcker struct {
-	io.ReadWriter
+	*quicpogs.RequestServerStream
 }
 
 // AckConnection acks response back to the proxy.
 func (s *streamReadWriteAcker) AckConnection() error {
-	return quicpogs.WriteConnectResponseData(s, nil)
+	return s.WriteConnectResponseData(nil)
 }
 
 // httpResponseAdapter translates responses written by the HTTP Proxy into ones that can be used in QUIC.
 type httpResponseAdapter struct {
-	io.Writer
+	*quicpogs.RequestServerStream
 }
 
-func newHTTPResponseAdapter(w io.Writer) httpResponseAdapter {
-	return httpResponseAdapter{w}
+func newHTTPResponseAdapter(s *quicpogs.RequestServerStream) httpResponseAdapter {
+	return httpResponseAdapter{s}
 }
 
 func (hrw httpResponseAdapter) WriteRespHeaders(status int, header http.Header) error {
@@ -151,11 +183,11 @@ func (hrw httpResponseAdapter) WriteRespHeaders(status int, header http.Header)
 			metadata = append(metadata, quicpogs.Metadata{Key: httpHeaderKey, Val: v})
 		}
 	}
-	return quicpogs.WriteConnectResponseData(hrw, nil, metadata...)
+	return hrw.WriteConnectResponseData(nil, metadata...)
 }
 
 func (hrw httpResponseAdapter) WriteErrorResponse(err error) {
-	quicpogs.WriteConnectResponseData(hrw, err, quicpogs.Metadata{Key: "HttpStatus", Val: strconv.Itoa(http.StatusBadGateway)})
+	hrw.WriteConnectResponseData(err, quicpogs.Metadata{Key: "HttpStatus", Val: strconv.Itoa(http.StatusBadGateway)})
 }
 
 func buildHTTPRequest(connectRequest *quicpogs.ConnectRequest, body io.ReadCloser) (*http.Request, error) {

connection/quic_test.go

@@ -26,7 +26,6 @@ import (
 	"github.com/stretchr/testify/require"
 
 	quicpogs "github.com/cloudflare/cloudflared/quic"
-	"github.com/cloudflare/cloudflared/tunnelrpc/pogs"
 	tunnelpogs "github.com/cloudflare/cloudflared/tunnelrpc/pogs"
 )
 
@@ -76,15 +75,15 @@ func TestQUICServer(t *testing.T) {
 			dest:           "/ok",
 			connectionType: quicpogs.ConnectionTypeHTTP,
 			metadata: []quicpogs.Metadata{
-				quicpogs.Metadata{
+				{
 					Key: "HttpHeader:Cf-Ray",
 					Val: "123123123",
 				},
-				quicpogs.Metadata{
+				{
 					Key: "HttpHost",
 					Val: "cf.host",
 				},
-				quicpogs.Metadata{
+				{
 					Key: "HttpMethod",
 					Val: "GET",
 				},
@@ -96,19 +95,19 @@ func TestQUICServer(t *testing.T) {
 			dest:           "/echo_body",
 			connectionType: quicpogs.ConnectionTypeHTTP,
 			metadata: []quicpogs.Metadata{
-				quicpogs.Metadata{
+				{
 					Key: "HttpHeader:Cf-Ray",
 					Val: "123123123",
 				},
-				quicpogs.Metadata{
+				{
 					Key: "HttpHost",
 					Val: "cf.host",
 				},
-				quicpogs.Metadata{
+				{
 					Key: "HttpMethod",
 					Val: "POST",
 				},
-				quicpogs.Metadata{
+				{
 					Key: "HttpHeader:Content-Length",
 					Val: "24",
 				},
@@ -121,19 +120,19 @@ func TestQUICServer(t *testing.T) {
 			dest:           "/ok",
 			connectionType: quicpogs.ConnectionTypeWebsocket,
 			metadata: []quicpogs.Metadata{
-				quicpogs.Metadata{
+				{
 					Key: "HttpHeader:Cf-Cloudflared-Proxy-Connection-Upgrade",
 					Val: "Websocket",
 				},
-				quicpogs.Metadata{
+				{
 					Key: "HttpHeader:Another-Header",
 					Val: "Misc",
 				},
-				quicpogs.Metadata{
+				{
 					Key: "HttpHost",
 					Val: "cf.host",
 				},
-				quicpogs.Metadata{
+				{
 					Key: "HttpMethod",
 					Val: "get",
 				},
@@ -171,7 +170,7 @@ func TestQUICServer(t *testing.T) {
 				udpListener.LocalAddr(),
 				tlsClientConfig,
 				originProxy,
-				&pogs.ConnectionOptions{},
+				&tunnelpogs.ConnectionOptions{},
 				controlStream,
 				NewObserver(&log, &log, false),
 			)
@@ -218,10 +217,11 @@ func quicServer(
 	stream, err := session.OpenStreamSync(context.Background())
 	require.NoError(t, err)
 
-	err = quicpogs.WriteConnectRequestData(stream, dest, connectionType, metadata...)
+	reqClientStream := quicpogs.RequestClientStream{ReadWriteCloser: stream}
+	err = reqClientStream.WriteConnectRequestData(dest, connectionType, metadata...)
 	require.NoError(t, err)
 
-	_, err = quicpogs.ReadConnectResponseData(stream)
+	_, err = reqClientStream.ReadConnectResponseData()
 	require.NoError(t, err)
 
 	if message != nil {
@@ -309,23 +309,23 @@ func TestBuildHTTPRequest(t *testing.T) {
 			connectRequest: &quicpogs.ConnectRequest{
 				Dest: "http://test.com",
 				Metadata: []quicpogs.Metadata{
-					quicpogs.Metadata{
+					{
 						Key: "HttpHeader:Cf-Cloudflared-Proxy-Connection-Upgrade",
 						Val: "Websocket",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpHeader:Content-Length",
 						Val: "514",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpHeader:Another-Header",
 						Val: "Misc",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpHost",
 						Val: "cf.host",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpMethod",
 						Val: "get",
 					},
@@ -355,19 +355,19 @@ func TestBuildHTTPRequest(t *testing.T) {
 			connectRequest: &quicpogs.ConnectRequest{
 				Dest: "http://test.com",
 				Metadata: []quicpogs.Metadata{
-					quicpogs.Metadata{
+					{
 						Key: "HttpHeader:Cf-Cloudflared-Proxy-Connection-Upgrade",
 						Val: "Websocket",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpHeader:Another-Header",
 						Val: "Misc",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpHost",
 						Val: "cf.host",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpMethod",
 						Val: "get",
 					},
@@ -396,19 +396,19 @@ func TestBuildHTTPRequest(t *testing.T) {
 			connectRequest: &quicpogs.ConnectRequest{
 				Dest: "http://test.com",
 				Metadata: []quicpogs.Metadata{
-					quicpogs.Metadata{
+					{
 						Key: "HttpHeader:Another-Header",
 						Val: "Misc",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpHeader:Transfer-Encoding",
 						Val: "chunked",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpHost",
 						Val: "cf.host",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpMethod",
 						Val: "get",
 					},
@@ -438,19 +438,19 @@ func TestBuildHTTPRequest(t *testing.T) {
 			connectRequest: &quicpogs.ConnectRequest{
 				Dest: "http://test.com",
 				Metadata: []quicpogs.Metadata{
-					quicpogs.Metadata{
+					{
 						Key: "HttpHeader:Another-Header",
 						Val: "Misc",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpHeader:Transfer-Encoding",
 						Val: "gzip,chunked",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpHost",
 						Val: "cf.host",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpMethod",
 						Val: "get",
 					},
@@ -481,15 +481,15 @@ func TestBuildHTTPRequest(t *testing.T) {
 				Type: quicpogs.ConnectionTypeWebsocket,
 				Dest: "http://test.com",
 				Metadata: []quicpogs.Metadata{
-					quicpogs.Metadata{
+					{
 						Key: "HttpHeader:Another-Header",
 						Val: "Misc",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpHost",
 						Val: "cf.host",
 					},
-					quicpogs.Metadata{
+					{
 						Key: "HttpMethod",
 						Val: "get",
 					},

connection/udp_session.go

@@ -0,0 +1,43 @@
+package connection
+
+import (
+	"net"
+	"sync"
+
+	"github.com/google/uuid"
+)
+
+// TODO: TUN-5422 Unregister session
+type udpSessions struct {
+	lock     sync.Mutex
+	sessions map[uuid.UUID]*net.UDPConn
+}
+
+func newUDPSessions() *udpSessions {
+	return &udpSessions{
+		sessions: make(map[uuid.UUID]*net.UDPConn),
+	}
+}
+
+func (us *udpSessions) register(id uuid.UUID, dstIP net.IP, dstPort uint16) error {
+	us.lock.Lock()
+	defer us.lock.Unlock()
+	dstAddr := &net.UDPAddr{
+		IP:   dstIP,
+		Port: int(dstPort),
+	}
+	conn, err := net.DialUDP("udp", us.localAddr(), dstAddr)
+	if err != nil {
+		return err
+	}
+	us.sessions[id] = conn
+	return nil
+}
+
+func (ud *udpSessions) localAddr() *net.UDPAddr {
+	// TODO: Determine the IP to bind to
+	return &net.UDPAddr{
+		IP:   net.IPv4zero,
+		Port: 0,
+	}
+}