From 7e3a8d3a04fe94de544184353a2545b2f269c34e Mon Sep 17 00:00:00 2001 From: yuhan6665 <1588741+yuhan6665@users.noreply.github.com> Date: Fri, 29 Mar 2024 02:27:02 -0400 Subject: [PATCH] Add separate host config for websocket --- infra/conf/transport_internet.go | 18 ++++++++++++------ transport/internet/websocket/config.go | 1 + transport/internet/websocket/hub.go | 6 ++++++ 3 files changed, 19 insertions(+), 6 deletions(-) diff --git a/infra/conf/transport_internet.go b/infra/conf/transport_internet.go index bcb33937..805a674e 100644 --- a/infra/conf/transport_internet.go +++ b/infra/conf/transport_internet.go @@ -146,6 +146,7 @@ func (c *TCPConfig) Build() (proto.Message, error) { } type WebSocketConfig struct { + Host string `json:"host"` Path string `json:"path"` Headers map[string]string `json:"headers"` AcceptProxyProtocol bool `json:"acceptProxyProtocol"` @@ -154,10 +155,6 @@ type WebSocketConfig struct { // Build implements Buildable. func (c *WebSocketConfig) Build() (proto.Message, error) { path := c.Path - header := make(map[string]string); - for key, value := range c.Headers { - header[key] = value; - } var ed uint32 if u, err := url.Parse(path); err == nil { if q := u.Query(); q.Get("ed") != "" { @@ -168,9 +165,18 @@ func (c *WebSocketConfig) Build() (proto.Message, error) { path = u.String() } } + // If http host is not set in the Host field, but in headers field, we add it to Host Field here. + // If we don't do that, http host will be overwritten as address. + // Host priority: Host field > headers field > address. + if c.Host == "" && c.Headers["host"] != "" { + c.Host = c.Headers["host"] + } else if c.Host == "" && c.Headers["Host"] != "" { + c.Host = c.Headers["Host"] + } config := &websocket.Config{ Path: path, - Header: header, + Host: c.Host, + Header: c.Headers, AcceptProxyProtocol: c.AcceptProxyProtocol, Ed: ed, } @@ -178,8 +184,8 @@ func (c *WebSocketConfig) Build() (proto.Message, error) { } type HttpUpgradeConfig struct { - Path string `json:"path"` Host string `json:"host"` + Path string `json:"path"` Headers map[string]string `json:"headers"` AcceptProxyProtocol bool `json:"acceptProxyProtocol"` } diff --git a/transport/internet/websocket/config.go b/transport/internet/websocket/config.go index 2e3cfea5..8cb4a855 100644 --- a/transport/internet/websocket/config.go +++ b/transport/internet/websocket/config.go @@ -25,6 +25,7 @@ func (c *Config) GetRequestHeader() http.Header { for k, v := range c.Header { header.Add(k, v) } + header.Set("Host", c.Host) return header } diff --git a/transport/internet/websocket/hub.go b/transport/internet/websocket/hub.go index b3fe9ff4..8a860c30 100644 --- a/transport/internet/websocket/hub.go +++ b/transport/internet/websocket/hub.go @@ -21,6 +21,7 @@ import ( ) type requestHandler struct { + host string path string ln *Listener } @@ -37,6 +38,10 @@ var upgrader = &websocket.Upgrader{ } func (h *requestHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) { + if len(h.host) > 0 && request.Host != h.host { + writer.WriteHeader(http.StatusNotFound) + return + } if request.URL.Path != h.path { writer.WriteHeader(http.StatusNotFound) return @@ -125,6 +130,7 @@ func ListenWS(ctx context.Context, address net.Address, port net.Port, streamSet l.server = http.Server{ Handler: &requestHandler{ + host: wsSettings.Host, path: wsSettings.GetNormalizedPath(), ln: l, },