Add separate host config for websocket

This commit is contained in:
yuhan6665 2024-03-29 02:27:02 -04:00 committed by RPRX
parent e2302b421c
commit 7e3a8d3a04
3 changed files with 19 additions and 6 deletions

View File

@ -146,6 +146,7 @@ func (c *TCPConfig) Build() (proto.Message, error) {
} }
type WebSocketConfig struct { type WebSocketConfig struct {
Host string `json:"host"`
Path string `json:"path"` Path string `json:"path"`
Headers map[string]string `json:"headers"` Headers map[string]string `json:"headers"`
AcceptProxyProtocol bool `json:"acceptProxyProtocol"` AcceptProxyProtocol bool `json:"acceptProxyProtocol"`
@ -154,10 +155,6 @@ type WebSocketConfig struct {
// Build implements Buildable. // Build implements Buildable.
func (c *WebSocketConfig) Build() (proto.Message, error) { func (c *WebSocketConfig) Build() (proto.Message, error) {
path := c.Path path := c.Path
header := make(map[string]string);
for key, value := range c.Headers {
header[key] = value;
}
var ed uint32 var ed uint32
if u, err := url.Parse(path); err == nil { if u, err := url.Parse(path); err == nil {
if q := u.Query(); q.Get("ed") != "" { if q := u.Query(); q.Get("ed") != "" {
@ -168,9 +165,18 @@ func (c *WebSocketConfig) Build() (proto.Message, error) {
path = u.String() 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{ config := &websocket.Config{
Path: path, Path: path,
Header: header, Host: c.Host,
Header: c.Headers,
AcceptProxyProtocol: c.AcceptProxyProtocol, AcceptProxyProtocol: c.AcceptProxyProtocol,
Ed: ed, Ed: ed,
} }
@ -178,8 +184,8 @@ func (c *WebSocketConfig) Build() (proto.Message, error) {
} }
type HttpUpgradeConfig struct { type HttpUpgradeConfig struct {
Path string `json:"path"`
Host string `json:"host"` Host string `json:"host"`
Path string `json:"path"`
Headers map[string]string `json:"headers"` Headers map[string]string `json:"headers"`
AcceptProxyProtocol bool `json:"acceptProxyProtocol"` AcceptProxyProtocol bool `json:"acceptProxyProtocol"`
} }

View File

@ -25,6 +25,7 @@ func (c *Config) GetRequestHeader() http.Header {
for k, v := range c.Header { for k, v := range c.Header {
header.Add(k, v) header.Add(k, v)
} }
header.Set("Host", c.Host)
return header return header
} }

View File

@ -21,6 +21,7 @@ import (
) )
type requestHandler struct { type requestHandler struct {
host string
path string path string
ln *Listener ln *Listener
} }
@ -37,6 +38,10 @@ var upgrader = &websocket.Upgrader{
} }
func (h *requestHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) { 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 { if request.URL.Path != h.path {
writer.WriteHeader(http.StatusNotFound) writer.WriteHeader(http.StatusNotFound)
return return
@ -125,6 +130,7 @@ func ListenWS(ctx context.Context, address net.Address, port net.Port, streamSet
l.server = http.Server{ l.server = http.Server{
Handler: &requestHandler{ Handler: &requestHandler{
host: wsSettings.Host,
path: wsSettings.GetNormalizedPath(), path: wsSettings.GetNormalizedPath(),
ln: l, ln: l,
}, },