From 75231604c7ff24a3cc475ccde740bba4dfc79535 Mon Sep 17 00:00:00 2001
From: Shelikhoo <xiaokangwang@outlook.com>
Date: Sat, 1 May 2021 01:31:41 +0100
Subject: [PATCH] add support for header based early data on server side

---
 transport/internet/websocket/hub.go | 25 ++++++++++++++++++-------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/transport/internet/websocket/hub.go b/transport/internet/websocket/hub.go
index 167938a73..69eb29ae8 100644
--- a/transport/internet/websocket/hub.go
+++ b/transport/internet/websocket/hub.go
@@ -24,9 +24,10 @@ import (
 )
 
 type requestHandler struct {
-	path             string
-	ln               *Listener
-	earlyDataEnabled bool
+	path                string
+	ln                  *Listener
+	earlyDataEnabled    bool
+	earlyDataHeaderName string
 }
 
 var upgrader = &websocket.Upgrader{
@@ -45,6 +46,13 @@ func (h *requestHandler) ServeHTTP(writer http.ResponseWriter, request *http.Req
 			writer.WriteHeader(http.StatusNotFound)
 			return
 		}
+	} else if h.earlyDataHeaderName != "" {
+		if request.URL.Path != h.path {
+			writer.WriteHeader(http.StatusNotFound)
+			return
+		}
+		earlyDataStr := request.Header.Get(h.earlyDataHeaderName)
+		earlyData = base64.NewDecoder(base64.RawURLEncoding, bytes.NewReader([]byte(earlyDataStr)))
 	} else {
 		if strings.HasPrefix(request.URL.RequestURI(), h.path) {
 			earlyDataStr := request.URL.RequestURI()[len(h.path):]
@@ -135,18 +143,21 @@ func ListenWS(ctx context.Context, address net.Address, port net.Port, streamSet
 
 	l.listener = listener
 	var useEarlyData = false
+	var earlyDataHeaderName = ""
 	if wsSettings.MaxEarlyData != 0 {
 		useEarlyData = true
+		earlyDataHeaderName = wsSettings.EarlyDataHeaderName
 	}
 
 	l.server = http.Server{
 		Handler: &requestHandler{
-			path:             wsSettings.GetNormalizedPath(),
-			ln:               l,
-			earlyDataEnabled: useEarlyData,
+			path:                wsSettings.GetNormalizedPath(),
+			ln:                  l,
+			earlyDataEnabled:    useEarlyData,
+			earlyDataHeaderName: earlyDataHeaderName,
 		},
 		ReadHeaderTimeout: time.Second * 4,
-		MaxHeaderBytes:    2048,
+		MaxHeaderBytes:    http.DefaultMaxHeaderBytes,
 	}
 
 	go func() {