From 599b9bf91cee00e675640e2f0737baf812007c94 Mon Sep 17 00:00:00 2001 From: durck Date: Thu, 9 Jul 2026 13:10:51 +0300 Subject: [PATCH] feat: add configurable web transport paths --- README.md | 19 +++++++ cmd/client/main.go | 16 ++++++ cmd/server/main.go | 8 ++- docker-entrypoint.sh | 20 ++++++- internal/client/client.go | 10 +++- internal/client/http.go | 29 ++++++++-- internal/client/http_test.go | 41 ++++++++++++++ internal/server/commands/link.go | 12 +++++ internal/server/server.go | 4 +- internal/server/webserver/buildmanager.go | 4 +- pkg/mux/multiplexer.go | 65 +++++++++++++++++++++-- pkg/mux/multiplexer_test.go | 59 ++++++++++++++++++++ pkg/transport/paths.go | 31 +++++++++++ pkg/transport/paths_test.go | 47 ++++++++++++++++ 14 files changed, 349 insertions(+), 16 deletions(-) create mode 100644 internal/client/http_test.go create mode 100644 pkg/mux/multiplexer_test.go create mode 100644 pkg/transport/paths.go create mode 100644 pkg/transport/paths_test.go diff --git a/README.md b/README.md index fc262661..d50b2160 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,8 @@ services: - EXTERNAL_ADDRESS=:3232 - RSSH_CONSOLE_LABEL=c2.label - RSSH_LOG_LEVEL=INFO # DISABLED, INFO, WARNING, ERROR, FATAL + - RSSH_WS_PATH=/ws + - RSSH_PUSH_PATH=/push - SEED_AUTHORIZED_KEYS=${SSH_PUBLIC_KEY} volumes: - ./data:/data @@ -221,6 +223,7 @@ This requires the web server component has been enabled. --ntlm-proxy-creds Set NTLM proxy credentials in format DOMAIN\\USER:PASS --owners Set owners of client, if unset client is public all users. E.g --owners jsmith,ldavidson --proxy Set connect proxy address to bake it + --push-path Set HTTP(S) polling base path to bake into the client --raw-download Download over raw TCP, outputs bash downloader rather than http --shared-object Generate shared object file --sni When TLS is in use, set a custom SNI for the client to connect with @@ -230,6 +233,7 @@ This requires the web server component has been enabled. --use-kerberos Instruct client to try and use kerberos ticket when using a proxy --working-directory Set download/working directory for automatic script (i.e doing curl https://.sh) --ws Use plain http websockets as the underlying transport + --ws-path Set WebSocket transport path to bake into the client --wss Use TLS websockets as the underlying transport -C Comment to add as the public key (acts as the name) -l List currently active download links @@ -275,6 +279,21 @@ Or by baking it in with the `link` command. ssh your.rssh.server -p 3232 link --ws --name test ``` +WebSocket and HTTP(S) polling transports default to `/ws` and `/push`. If a +reverse proxy needs different paths, configure the same paths on the server and +client: + +```sh +reverse_ssh --ws-path /socket --push-path /api/push :3232 +./client -d ws://your.rssh.server:3232 --ws-path /socket +./client -d http://your.rssh.server:3232 --push-path /api/push +catcher$ link --ws --ws-path /socket --name test +catcher$ link --http --push-path /api/push --name test-http +``` + +The Docker entrypoint accepts these server options through `RSSH_WS_PATH` and +`RSSH_PUSH_PATH`. + ### Bash autocomplete The RSSH server has the `autocomplete` command which integrates nicely with bash so that you can have autocompletions when not using the server console. diff --git a/cmd/client/main.go b/cmd/client/main.go index 4f5fbd42..644ce07b 100644 --- a/cmd/client/main.go +++ b/cmd/client/main.go @@ -44,6 +44,8 @@ var ( proxy string ignoreInput string customSNI string + wsPath string + pushPath string // golang can only embed strings using the compile time linker useHostKerberos string logLevel string @@ -64,6 +66,8 @@ func printHelp() { fmt.Println("\t\t--ntlm-proxy-creds\tNTLM proxy credentials in format DOMAIN\\USER:PASS") fmt.Println("\t\t--process_name\tProcess name shown in tasklist/process list") fmt.Println("\t\t--sni\tWhen using TLS set the clients requested SNI to this value") + fmt.Println("\t\t--ws-path\tWebSocket transport path, defaults to /ws") + fmt.Println("\t\t--push-path\tHTTP(S) polling transport base path, defaults to /push") fmt.Println("\t\t--log-level\tChange logging output levels, [INFO,WARNING,ERROR,FATAL,DISABLED]") fmt.Println("\t\t--version-string\tSSH version string to use, i.e SSH-VERSION, defaults to internal.Version-runtime.GOOS_runtime.GOARCH") fmt.Println("\t\t--private-key-path\tOptional path to unencrypted SSH key to use for connecting") @@ -82,6 +86,8 @@ func makeInitialSettings() (*client.Settings, error) { Addr: destination, ProxyUseHostKerberos: useHostKerberos == "true", SNI: customSNI, + WSPath: wsPath, + PushPath: pushPath, VersionString: versionString, } @@ -175,6 +181,16 @@ func main() { settings.SNI = userSpecifiedSNI } + userSpecifiedWSPath, err := line.GetArgString("ws-path") + if err == nil { + settings.WSPath = userSpecifiedWSPath + } + + userSpecifiedPushPath, err := line.GetArgString("push-path") + if err == nil { + settings.PushPath = userSpecifiedPushPath + } + timeoutInt := 180 timeout, err := line.GetArgString("connect-timeout") if err == nil { diff --git a/cmd/server/main.go b/cmd/server/main.go index 014278f4..928322fd 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -31,6 +31,8 @@ func printHelp() { fmt.Println("\t--webserver\t\t(Depreciated) Enable webserver on the listen_address port") fmt.Println("\t--enable-client-downloads\t\tEnable webserver and raw TCP to download clients") fmt.Println("\t--external_address\tIf the external IP and port of the RSSH server is different from the listening address, set that here") + fmt.Println("\t--ws-path\t\tWebSocket transport path, defaults to /ws") + fmt.Println("\t--push-path\t\tHTTP(S) polling transport base path, defaults to /push") fmt.Println("\t--timeout\t\tSet rssh client timeout (when a client is considered disconnected) defaults, in seconds, defaults to 5, if set to 0 timeout is disabled") fmt.Println(" Utility") fmt.Println("\t--fingerprint\t\tPrint fingerprint and exit. (Will generate server key if none exists)") @@ -50,6 +52,8 @@ func main() { "fingerprint": true, "webserver": true, // deprecated "enable-client-downloads": true, + "ws-path": true, + "push-path": true, "datadir": true, "h": true, "help": true, @@ -166,6 +170,8 @@ func main() { tlskey, _ := options.GetArgString("tlskey") enabledDownloads := options.IsSet("webserver") || options.IsSet("enable-client-downloads") + wsPath, _ := options.GetArgString("ws-path") + pushPath, _ := options.GetArgString("push-path") if options.IsSet("webserver") { log.Println("[WARNING] --webserver is deprecated, use --enable-client-downloads") @@ -211,5 +217,5 @@ func main() { log.Println("connect back: ", connectBackAddress) - server.Run(listenAddress, dataDir, connectBackAddress, autogeneratedConnectBack, tlscert, tlskey, insecure, enabledDownloads, tls, openproxy, timeout) + server.Run(listenAddress, dataDir, connectBackAddress, autogeneratedConnectBack, tlscert, tlskey, insecure, enabledDownloads, tls, openproxy, timeout, wsPath, pushPath) } diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh index e13189cc..7c8c6a99 100644 --- a/docker-entrypoint.sh +++ b/docker-entrypoint.sh @@ -24,4 +24,22 @@ if [ ! -z "$SEED_AUTHORIZED_KEYS" ]; then fi cd /app/bin -exec ./server --datadir /data --enable-client-downloads --tls --external_address $EXTERNAL_ADDRESS :2222 \ No newline at end of file + +server_args=( + --datadir /data + --enable-client-downloads + --tls + --external_address "$EXTERNAL_ADDRESS" +) + +if [ -n "$RSSH_WS_PATH" ]; then + server_args+=(--ws-path "$RSSH_WS_PATH") +fi + +if [ -n "$RSSH_PUSH_PATH" ]; then + server_args+=(--push-path "$RSSH_PUSH_PATH") +fi + +server_args+=(:2222) + +exec ./server "${server_args[@]}" diff --git a/internal/client/client.go b/internal/client/client.go index 641404bd..8bfc7fb9 100644 --- a/internal/client/client.go +++ b/internal/client/client.go @@ -23,6 +23,7 @@ import ( "github.com/NHAS/reverse_ssh/internal/client/handlers" "github.com/NHAS/reverse_ssh/internal/client/keys" "github.com/NHAS/reverse_ssh/pkg/logger" + "github.com/NHAS/reverse_ssh/pkg/transport" "github.com/bodgit/ntlmssp" "golang.org/x/crypto/ssh" socks "golang.org/x/net/proxy" @@ -305,6 +306,8 @@ type Settings struct { Fingerprint string ProxyAddr string SNI string + WSPath string + PushPath string ProxyUseHostKerberos bool @@ -344,6 +347,9 @@ func Run(settings *Settings) { l := logger.NewLog("client") var err error + settings.WSPath = transport.NormalizePath(settings.WSPath, transport.DefaultWSPath) + settings.PushPath = transport.NormalizePath(settings.PushPath, transport.DefaultPushPath) + settings.ProxyAddr, err = GetProxyDetails(settings.ProxyAddr) if err != nil { log.Fatal("Invalid proxy details", settings.ProxyAddr, ":", err) @@ -462,7 +468,7 @@ func Run(settings *Settings) { switch scheme { case "wss", "ws": - c, err := websocket.NewConfig("ws://"+realAddr+"/ws", "ws://"+realAddr) + c, err := websocket.NewConfig("ws://"+realAddr+settings.WSPath, "ws://"+realAddr) if err != nil { log.Println("Could not create websockets configuration: ", err) <-time.After(10 * time.Second) @@ -483,7 +489,7 @@ func Run(settings *Settings) { conn = wsConn case "http", "https": - conn, err = NewHTTPConn(scheme+"://"+realAddr, func() (net.Conn, error) { + conn, err = NewHTTPConnWithPushPath(scheme+"://"+realAddr, settings.PushPath, func() (net.Conn, error) { return Connect(realAddr, settings.ProxyAddr, settings.ConnectTimeout, settings.ProxyUseHostKerberos, settings.ntlm) }) diff --git a/internal/client/http.go b/internal/client/http.go index 5a4a823b..628465ca 100644 --- a/internal/client/http.go +++ b/internal/client/http.go @@ -16,6 +16,7 @@ import ( "github.com/NHAS/reverse_ssh/internal/client/keys" "github.com/NHAS/reverse_ssh/pkg/mux" + "github.com/NHAS/reverse_ssh/pkg/transport" ) type HTTPConn struct { @@ -29,16 +30,22 @@ type HTTPConn struct { // Cache buster for middleware proxies start int - client *http.Client + pushPath string + client *http.Client } func NewHTTPConn(address string, connector func() (net.Conn, error)) (*HTTPConn, error) { + return NewHTTPConnWithPushPath(address, transport.DefaultPushPath, connector) +} + +func NewHTTPConnWithPushPath(address, pushPath string, connector func() (net.Conn, error)) (*HTTPConn, error) { result := &HTTPConn{ done: make(chan interface{}), readBuffer: mux.NewSyncBuffer(8096), address: address, start: mathrand.Int(), + pushPath: transport.NormalizePath(pushPath, transport.DefaultPushPath), } result.client = &http.Client{ @@ -62,9 +69,9 @@ func NewHTTPConn(address string, connector func() (net.Conn, error)) (*HTTPConn, publicKeyBytes := s.PublicKey().Marshal() - resp, err := result.client.Head(address + "/push?key=" + hex.EncodeToString(publicKeyBytes)) + resp, err := result.client.Head(result.initURL(hex.EncodeToString(publicKeyBytes))) if err != nil { - return nil, fmt.Errorf("failed to connect to %s/push?key=%s, err: %s", address, hex.EncodeToString(publicKeyBytes), err) + return nil, fmt.Errorf("failed to connect to %s, err: %s", result.initURL(hex.EncodeToString(publicKeyBytes)), err) } resp.Body.Close() @@ -90,6 +97,18 @@ func NewHTTPConn(address string, connector func() (net.Conn, error)) (*HTTPConn, return result, nil } +func (c *HTTPConn) initURL(publicKeyHex string) string { + return c.address + c.pushPath + "?key=" + publicKeyHex +} + +func (c *HTTPConn) readURL() string { + return c.address + transport.JoinPushPath(c.pushPath, strconv.Itoa(c.start)) + "?id=" + c.ID +} + +func (c *HTTPConn) writeURL() string { + return c.address + c.pushPath + "?id=" + c.ID +} + func (c *HTTPConn) startReadLoop() { for { select { @@ -99,7 +118,7 @@ func (c *HTTPConn) startReadLoop() { default: } - resp, err := c.client.Get(c.address + "/push/" + strconv.Itoa(c.start) + "?id=" + c.ID) + resp, err := c.client.Get(c.readURL()) if err != nil { log.Println("error getting data: ", err) c.Close() @@ -142,7 +161,7 @@ func (c *HTTPConn) Write(b []byte) (n int, err error) { default: } - resp, err := c.client.Post(c.address+"/push?id="+c.ID, "application/octet-stream", bytes.NewBuffer(b)) + resp, err := c.client.Post(c.writeURL(), "application/octet-stream", bytes.NewBuffer(b)) if err != nil { c.Close() return 0, err diff --git a/internal/client/http_test.go b/internal/client/http_test.go new file mode 100644 index 00000000..2a23e113 --- /dev/null +++ b/internal/client/http_test.go @@ -0,0 +1,41 @@ +package client + +import "testing" + +func TestHTTPConnURLsUseDefaultPushPath(t *testing.T) { + c := &HTTPConn{ + ID: "session", + address: "https://example.internal", + pushPath: "/push", + start: 7, + } + + if got := c.initURL("abc123"); got != "https://example.internal/push?key=abc123" { + t.Fatalf("initURL = %q", got) + } + if got := c.readURL(); got != "https://example.internal/push/7?id=session" { + t.Fatalf("readURL = %q", got) + } + if got := c.writeURL(); got != "https://example.internal/push?id=session" { + t.Fatalf("writeURL = %q", got) + } +} + +func TestHTTPConnURLsUseCustomPushPath(t *testing.T) { + c := &HTTPConn{ + ID: "session", + address: "https://example.internal", + pushPath: "/api/poll", + start: 9, + } + + if got := c.initURL("abc123"); got != "https://example.internal/api/poll?key=abc123" { + t.Fatalf("initURL = %q", got) + } + if got := c.readURL(); got != "https://example.internal/api/poll/9?id=session" { + t.Fatalf("readURL = %q", got) + } + if got := c.writeURL(); got != "https://example.internal/api/poll?id=session" { + t.Fatalf("writeURL = %q", got) + } +} diff --git a/internal/server/commands/link.go b/internal/server/commands/link.go index 85c36a5f..ae1bc73e 100644 --- a/internal/server/commands/link.go +++ b/internal/server/commands/link.go @@ -41,6 +41,8 @@ func (l *link) ValidArgs() map[string]string { "stdio": "Use stdin and stdout as transport, will disable logging, destination after stdio:// is ignored", "http": "Use http polling as the underlying transport", "https": "Use https polling as the underlying transport", + "ws-path": "Set WebSocket transport path to bake into the client", + "push-path": "Set HTTP(S) polling base path to bake into the client", "use-host-header": "Use HTTP Host header as callback address when generating download template (add .sh to your download urls and find out)", "shared-object": "Generate shared object file", "fingerprint": "Set RSSH server fingerprint will default to server public key", @@ -212,6 +214,16 @@ func (l *link) Run(user *users.User, tty io.ReadWriter, line terminal.ParsedLine return err } + buildConfig.WSPath, err = line.GetArgString("ws-path") + if err != nil && err != terminal.ErrFlagNotSet { + return err + } + + buildConfig.PushPath, err = line.GetArgString("push-path") + if err != nil && err != terminal.ErrFlagNotSet { + return err + } + buildConfig.LogLevel, err = line.GetArgString("log-level") if err != nil { if err != terminal.ErrFlagNotSet { diff --git a/internal/server/server.go b/internal/server/server.go index 9a8626f7..b60be23d 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -47,7 +47,7 @@ func CreateOrLoadServerKeys(privateKeyPath string) (ssh.Signer, error) { return private, nil } -func Run(addr, dataDir, connectBackAddress string, autogeneratedConnectBack bool, TLSCertPath, TLSKeyPath string, insecure, enabledDownloads, enabletTLS, openproxy bool, timeout int) { +func Run(addr, dataDir, connectBackAddress string, autogeneratedConnectBack bool, TLSCertPath, TLSKeyPath string, insecure, enabledDownloads, enabletTLS, openproxy bool, timeout int, wsPath, pushPath string) { c := mux.MultiplexerConfig{ Control: true, Downloads: enabledDownloads, @@ -56,6 +56,8 @@ func Run(addr, dataDir, connectBackAddress string, autogeneratedConnectBack bool TLSKeyPath: TLSKeyPath, AutoTLSCommonName: connectBackAddress, TcpKeepAlive: timeout, + WSPath: wsPath, + PushPath: pushPath, PollingAuthChecker: func(key string, addr net.Addr) bool { authorizedKey, err := hex.DecodeString(key) diff --git a/internal/server/webserver/buildmanager.go b/internal/server/webserver/buildmanager.go index 42cf307b..14b6d489 100644 --- a/internal/server/webserver/buildmanager.go +++ b/internal/server/webserver/buildmanager.go @@ -45,7 +45,7 @@ type BuildConfig struct { ConnectBackAdress, Fingerprint string - Proxy, SNI, LogLevel string + Proxy, SNI, WSPath, PushPath, LogLevel string UseKerberosAuth bool @@ -183,7 +183,7 @@ func Build(config BuildConfig) (string, error) { return "", err } - buildArguments = append(buildArguments, fmt.Sprintf("-ldflags=-s -w -X main.logLevel=%s -X main.destination=%s -X main.fingerprint=%s -X main.proxy=%s -X main.customSNI=%s -X main.useHostKerberos=%t -X main.ntlmProxyCreds=%s -X main.versionString=%s -X github.com/NHAS/reverse_ssh/internal.Version=%s", config.LogLevel, config.ConnectBackAdress, config.Fingerprint, config.Proxy, config.SNI, config.UseKerberosAuth, config.NTLMProxyCreds, strings.TrimSpace(config.VersionString), strings.TrimSpace(f.Version))) + buildArguments = append(buildArguments, fmt.Sprintf("-ldflags=-s -w -X main.logLevel=%s -X main.destination=%s -X main.fingerprint=%s -X main.proxy=%s -X main.customSNI=%s -X main.wsPath=%s -X main.pushPath=%s -X main.useHostKerberos=%t -X main.ntlmProxyCreds=%s -X main.versionString=%s -X github.com/NHAS/reverse_ssh/internal.Version=%s", config.LogLevel, config.ConnectBackAdress, config.Fingerprint, config.Proxy, config.SNI, config.WSPath, config.PushPath, config.UseKerberosAuth, config.NTLMProxyCreds, strings.TrimSpace(config.VersionString), strings.TrimSpace(f.Version))) buildArguments = append(buildArguments, "-o", f.FilePath, filepath.Join(projectRoot, "/cmd/client")) cmd := exec.Command(buildTool, buildArguments...) diff --git a/pkg/mux/multiplexer.go b/pkg/mux/multiplexer.go index d0c97eb3..b3c43d95 100644 --- a/pkg/mux/multiplexer.go +++ b/pkg/mux/multiplexer.go @@ -15,6 +15,7 @@ import ( "math/big" "net" "net/http" + "net/url" "sort" "strings" "sync" @@ -22,12 +23,15 @@ import ( "time" "github.com/NHAS/reverse_ssh/pkg/mux/protocols" + "github.com/NHAS/reverse_ssh/pkg/transport" "golang.org/x/net/websocket" ) type MultiplexerConfig struct { Control bool Downloads bool + WSPath string + PushPath string TLS bool AutoTLSCommonName string @@ -342,6 +346,8 @@ func ListenWithConfig(network, address string, _c MultiplexerConfig) (*Multiplex m.newConnections = make(chan net.Conn) m.listeners = make(map[string]net.Listener) m.result = map[protocols.Type]*multiplexerListener{} + _c.WSPath = transport.NormalizePath(_c.WSPath, transport.DefaultWSPath) + _c.PushPath = transport.NormalizePath(_c.PushPath, transport.DefaultPushPath) m.config = _c if _c.PollingAuthChecker == nil { @@ -456,7 +462,7 @@ func isHttp(b []byte) bool { func (m *Multiplexer) determineProtocol(conn net.Conn) (net.Conn, protocols.Type, error) { - header := make([]byte, 14) + header := make([]byte, 1024) n, err := conn.Read(header) if err != nil { conn.Close() @@ -479,11 +485,12 @@ func (m *Multiplexer) determineProtocol(conn net.Conn) (net.Conn, protocols.Type if isHttp(header) { - if bytes.HasPrefix(header, []byte("GET /ws")) { + requestType := classifyHTTPRequest(header[:n], m.config.WSPath, m.config.PushPath) + if requestType == protocols.Websockets { return c, protocols.Websockets, nil } - if bytes.HasPrefix(header, []byte("HEAD /push")) || bytes.HasPrefix(header, []byte("GET /push")) || bytes.HasPrefix(header, []byte("POST /push")) { + if requestType == protocols.HTTP { return c, protocols.HTTP, nil } @@ -494,6 +501,56 @@ func (m *Multiplexer) determineProtocol(conn net.Conn) (net.Conn, protocols.Type return nil, "", errors.New("unknown protocol: " + string(header[:n])) } +func classifyHTTPRequest(header []byte, wsPath, pushPath string) protocols.Type { + method, requestPath, ok := parseHTTPRequestLine(header) + if !ok { + return protocols.Invalid + } + + wsPath = transport.NormalizePath(wsPath, transport.DefaultWSPath) + pushPath = transport.NormalizePath(pushPath, transport.DefaultPushPath) + + switch { + case method == http.MethodGet && requestPath == wsPath: + return protocols.Websockets + case isPollingRequest(method, requestPath, pushPath): + return protocols.HTTP + default: + return protocols.HTTPDownload + } +} + +func parseHTTPRequestLine(header []byte) (method string, requestPath string, ok bool) { + line := string(header) + if idx := strings.Index(line, "\r\n"); idx >= 0 { + line = line[:idx] + } + + fields := strings.Fields(line) + if len(fields) < 2 { + return "", "", false + } + + parsed, err := url.ParseRequestURI(fields[1]) + if err != nil { + return "", "", false + } + + return fields[0], transport.NormalizePath(parsed.Path, "/"), true +} + +func isPollingRequest(method, requestPath, pushPath string) bool { + pushPath = transport.NormalizePath(pushPath, transport.DefaultPushPath) + switch method { + case http.MethodHead, http.MethodPost: + return requestPath == pushPath + case http.MethodGet: + return strings.HasPrefix(requestPath, pushPath+"/") + default: + return false + } +} + func (m *Multiplexer) getProtoListener(proto protocols.Type) net.Listener { ml, ok := m.result[proto] if !ok { @@ -604,7 +661,7 @@ func (m *Multiplexer) unwrapWebsockets(conn net.Conn) (net.Conn, protocols.Type, }, } - wsHttp.Handle("/ws", wsServer) + wsHttp.Handle(transport.NormalizePath(m.config.WSPath, transport.DefaultWSPath), wsServer) go http.Serve(&singleConnListener{conn: conn}, wsHttp) diff --git a/pkg/mux/multiplexer_test.go b/pkg/mux/multiplexer_test.go new file mode 100644 index 00000000..5da81c61 --- /dev/null +++ b/pkg/mux/multiplexer_test.go @@ -0,0 +1,59 @@ +package mux + +import ( + "testing" + + "github.com/NHAS/reverse_ssh/pkg/mux/protocols" +) + +func TestClassifyHTTPRequestDefaultPaths(t *testing.T) { + tests := []struct { + name string + header string + want protocols.Type + }{ + {name: "websocket", header: "GET /ws HTTP/1.1\r\nHost: example\r\n\r\n", want: protocols.Websockets}, + {name: "poll head", header: "HEAD /push?key=abc HTTP/1.1\r\nHost: example\r\n\r\n", want: protocols.HTTP}, + {name: "poll get", header: "GET /push/123?id=abc HTTP/1.1\r\nHost: example\r\n\r\n", want: protocols.HTTP}, + {name: "poll post", header: "POST /push?id=abc HTTP/1.1\r\nHost: example\r\n\r\n", want: protocols.HTTP}, + {name: "download", header: "GET /payload.exe HTTP/1.1\r\nHost: example\r\n\r\n", want: protocols.HTTPDownload}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := classifyHTTPRequest([]byte(tt.header), "", ""); got != tt.want { + t.Fatalf("classifyHTTPRequest = %q, want %q", got, tt.want) + } + }) + } +} + +func TestClassifyHTTPRequestCustomPaths(t *testing.T) { + tests := []struct { + name string + header string + want protocols.Type + }{ + {name: "websocket", header: "GET /socket HTTP/1.1\r\nHost: example\r\n\r\n", want: protocols.Websockets}, + {name: "poll head", header: "HEAD /api/poll?key=abc HTTP/1.1\r\nHost: example\r\n\r\n", want: protocols.HTTP}, + {name: "poll get", header: "GET /api/poll/123?id=abc HTTP/1.1\r\nHost: example\r\n\r\n", want: protocols.HTTP}, + {name: "poll post", header: "POST /api/poll?id=abc HTTP/1.1\r\nHost: example\r\n\r\n", want: protocols.HTTP}, + {name: "default websocket becomes download", header: "GET /ws HTTP/1.1\r\nHost: example\r\n\r\n", want: protocols.HTTPDownload}, + {name: "default push becomes download", header: "HEAD /push?key=abc HTTP/1.1\r\nHost: example\r\n\r\n", want: protocols.HTTPDownload}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := classifyHTTPRequest([]byte(tt.header), "/socket", "/api/poll"); got != tt.want { + t.Fatalf("classifyHTTPRequest = %q, want %q", got, tt.want) + } + }) + } +} + +func TestClassifyHTTPRequestIgnoresQueryString(t *testing.T) { + got := classifyHTTPRequest([]byte("HEAD /api/poll?key=abc HTTP/1.1\r\n"), "/socket", "/api/poll") + if got != protocols.HTTP { + t.Fatalf("classifyHTTPRequest = %q, want %q", got, protocols.HTTP) + } +} diff --git a/pkg/transport/paths.go b/pkg/transport/paths.go new file mode 100644 index 00000000..29111bd2 --- /dev/null +++ b/pkg/transport/paths.go @@ -0,0 +1,31 @@ +package transport + +import "strings" + +const ( + DefaultWSPath = "/ws" + DefaultPushPath = "/push" +) + +func NormalizePath(path, fallback string) string { + path = strings.TrimSpace(path) + if path == "" { + path = fallback + } + if !strings.HasPrefix(path, "/") { + path = "/" + path + } + if len(path) > 1 { + path = strings.TrimRight(path, "/") + } + return path +} + +func JoinPushPath(pushPath, suffix string) string { + pushPath = NormalizePath(pushPath, DefaultPushPath) + suffix = strings.TrimLeft(suffix, "/") + if suffix == "" { + return pushPath + } + return pushPath + "/" + suffix +} diff --git a/pkg/transport/paths_test.go b/pkg/transport/paths_test.go new file mode 100644 index 00000000..65d7506b --- /dev/null +++ b/pkg/transport/paths_test.go @@ -0,0 +1,47 @@ +package transport + +import "testing" + +func TestNormalizePath(t *testing.T) { + tests := []struct { + name string + input string + fallback string + want string + }{ + {name: "fallback", fallback: "/ws", want: "/ws"}, + {name: "adds leading slash", input: "push", fallback: "/ws", want: "/push"}, + {name: "trims trailing slash", input: "/push/", fallback: "/ws", want: "/push"}, + {name: "keeps root", input: "/", fallback: "/ws", want: "/"}, + {name: "trims whitespace", input: " api/push/ ", fallback: "/ws", want: "/api/push"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := NormalizePath(tt.input, tt.fallback); got != tt.want { + t.Fatalf("NormalizePath(%q, %q) = %q, want %q", tt.input, tt.fallback, got, tt.want) + } + }) + } +} + +func TestJoinPushPath(t *testing.T) { + tests := []struct { + name string + path string + suffix string + want string + }{ + {name: "joins suffix", path: "/custom/", suffix: "123", want: "/custom/123"}, + {name: "trims suffix slash", path: "/custom", suffix: "/123", want: "/custom/123"}, + {name: "empty suffix returns path", path: "/custom", want: "/custom"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := JoinPushPath(tt.path, tt.suffix); got != tt.want { + t.Fatalf("JoinPushPath(%q, %q) = %q, want %q", tt.path, tt.suffix, got, tt.want) + } + }) + } +}