From 3da1c25471df3148b12dc766b3ee9c2440b9b0c5 Mon Sep 17 00:00:00 2001
From: Igor Postelnik <igor@cloudflare.com>
Date: Mon, 13 Jun 2022 12:46:52 -0500
Subject: [PATCH] TUN-6381: Write error data on QUIC stream when we fail to
 talk to the origin; separate logging for protocol errors vs. origin errors.

---
 connection/quic.go | 29 ++++++++++++++++++++---------
 quic/pogs.go       | 14 ++++++++++++++
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/connection/quic.go b/connection/quic.go
index 00765651..fc138e90 100644
--- a/connection/quic.go
+++ b/connection/quic.go
@@ -160,7 +160,7 @@ func (q *QUICConnection) handleStream(stream io.ReadWriteCloser) error {
 	case quicpogs.DataStreamProtocolSignature:
 		reqServerStream, err := quicpogs.NewRequestServerStream(stream, signature)
 		if err != nil {
-			return nil
+			return err
 		}
 		return q.handleDataStream(reqServerStream)
 	case quicpogs.RPCStreamProtocolSignature:
@@ -175,30 +175,41 @@ func (q *QUICConnection) handleStream(stream io.ReadWriteCloser) error {
 }
 
 func (q *QUICConnection) handleDataStream(stream *quicpogs.RequestServerStream) error {
-	connectRequest, err := stream.ReadConnectRequestData()
+	request, err := stream.ReadConnectRequestData()
 	if err != nil {
 		return err
 	}
 
+	if err := q.dispatchRequest(stream, err, request); err != nil {
+		_ = stream.WriteConnectResponseData(err)
+		q.logger.Err(err).Str("type", request.Type.String()).Str("dest", request.Dest).Msg("Request failed")
+	}
+
+	return nil
+}
+
+func (q *QUICConnection) dispatchRequest(stream *quicpogs.RequestServerStream, err error, request *quicpogs.ConnectRequest) error {
 	originProxy, err := q.orchestrator.GetOriginProxy()
 	if err != nil {
 		return err
 	}
 
-	switch connectRequest.Type {
+	switch request.Type {
 	case quicpogs.ConnectionTypeHTTP, quicpogs.ConnectionTypeWebsocket:
-		tracedReq, err := buildHTTPRequest(connectRequest, stream)
+		tracedReq, err := buildHTTPRequest(request, stream)
 		if err != nil {
 			return err
 		}
-
 		w := newHTTPResponseAdapter(stream)
-		return originProxy.ProxyHTTP(w, tracedReq, connectRequest.Type == quicpogs.ConnectionTypeWebsocket)
+		return originProxy.ProxyHTTP(w, tracedReq, request.Type == quicpogs.ConnectionTypeWebsocket)
+
 	case quicpogs.ConnectionTypeTCP:
 		rwa := &streamReadWriteAcker{stream}
-		metadata := connectRequest.MetadataMap()
-		return originProxy.ProxyTCP(context.Background(), rwa, &TCPRequest{Dest: connectRequest.Dest,
-			FlowID: metadata[QUICMetadataFlowID]})
+		metadata := request.MetadataMap()
+		return originProxy.ProxyTCP(context.Background(), rwa, &TCPRequest{
+			Dest:   request.Dest,
+			FlowID: metadata[QUICMetadataFlowID],
+		})
 	}
 	return nil
 }
diff --git a/quic/pogs.go b/quic/pogs.go
index e82fbe47..262e6561 100644
--- a/quic/pogs.go
+++ b/quic/pogs.go
@@ -1,6 +1,8 @@
 package quic
 
 import (
+	"fmt"
+
 	capnp "zombiezen.com/go/capnproto2"
 	"zombiezen.com/go/capnproto2/pogs"
 
@@ -16,6 +18,18 @@ const (
 	ConnectionTypeTCP
 )
 
+func (c ConnectionType) String() string {
+	switch c {
+	case ConnectionTypeHTTP:
+		return "http"
+	case ConnectionTypeWebsocket:
+		return "ws"
+	case ConnectionTypeTCP:
+		return "tcp"
+	}
+	panic(fmt.Sprintf("invalid ConnectionType: %d", c))
+}
+
 // ConnectRequest is the representation of metadata sent at the start of a QUIC application handshake.
 type ConnectRequest struct {
 	Dest     string         `capnp:"dest"`