diff --git a/README.md b/README.md index 822451e..cc095ed 100644 --- a/README.md +++ b/README.md @@ -216,6 +216,9 @@ $ chisel client --help the chisel server. Authentication can be specified inside the URL. For example, http://admin:password@my-server.com:8081 + --header, Set a custom header in the form "HeaderName: HeaderContent". + Can be used multiple times. (e.g --header "Foo: Bar" --header "Hello: World") + --hostname, Optionally set the 'Host' header (defaults to the host defined in the endpoint url). diff --git a/client/client.go b/client/client.go index e278e99..1fc5d03 100644 --- a/client/client.go +++ b/client/client.go @@ -13,7 +13,7 @@ import ( "github.com/gorilla/websocket" "github.com/jpillora/backoff" - "github.com/jpillora/chisel/share" + chshare "github.com/jpillora/chisel/share" "golang.org/x/crypto/ssh" ) @@ -28,7 +28,7 @@ type Config struct { Server string HTTPProxy string Remotes []string - HostHeader string + Headers http.Header } //Client represents a client instance @@ -199,13 +199,7 @@ func (c *Client) connectionLoop() { return c.httpProxyURL, nil } } - wsHeaders := http.Header{} - if c.config.HostHeader != "" { - wsHeaders = http.Header{ - "Host": {c.config.HostHeader}, - } - } - wsConn, _, err := d.Dial(c.server, wsHeaders) + wsConn, _, err := d.Dial(c.server, c.config.Headers) if err != nil { connerr = err continue diff --git a/client/client_test.go b/client/client_test.go new file mode 100644 index 0000000..d58f8e4 --- /dev/null +++ b/client/client_test.go @@ -0,0 +1,40 @@ +package chclient + +import ( + "log" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestCustomHeaders(t *testing.T) { + assert := assert.New(t) + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + assert.Equal(req.Header.Get("Foo"), "Bar") + })) + // Close the server when test finishes + defer server.Close() + headers := http.Header{} + headers.Set("Foo", "Bar") + config := Config{ + Fingerprint: "", + Auth: "", + KeepAlive: time.Second, + MaxRetryCount: 0, + MaxRetryInterval: time.Second, + HTTPProxy: "", + Server: server.URL, + Remotes: []string{"socks"}, + Headers: headers, + } + c, err := NewClient(&config) + if err != nil { + log.Fatal(err) + } + if err = c.Run(); err != nil { + log.Fatal(err) + } +} diff --git a/go.mod b/go.mod index 72fdfab..364d16d 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,7 @@ require ( github.com/jpillora/backoff v0.0.0-20180909062703-3050d21c67d7 github.com/jpillora/requestlog v0.0.0-20181015073026-df8817be5f82 github.com/jpillora/sizestr v0.0.0-20160130011556-e2ea2fa42fb9 + github.com/stretchr/testify v1.3.0 github.com/tomasen/realip v0.0.0-20180522021738-f0c99a92ddce // indirect golang.org/x/crypto v0.0.0-20181015023909-0c41d7ab0a0e golang.org/x/net v0.0.0-20181017193950-04a2e542c03f // indirect diff --git a/go.sum b/go.sum index e500886..8b7e51e 100644 --- a/go.sum +++ b/go.sum @@ -2,6 +2,8 @@ github.com/andrew-d/go-termutil v0.0.0-20150726205930-009166a695a2 h1:axBiC50cNZ github.com/andrew-d/go-termutil v0.0.0-20150726205930-009166a695a2/go.mod h1:jnzFpU88PccN/tPPhCpnNU8mZphvKxYM9lLNkd8e+os= github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio= github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs= +github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I= github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= github.com/gorilla/websocket v1.4.0 h1:WDFjx/TMzVgy9VdMMQi2K2Emtwi2QcUQsztZ/zLaH/Q= @@ -14,6 +16,11 @@ github.com/jpillora/requestlog v0.0.0-20181015073026-df8817be5f82 h1:7ufdyC3aMxF github.com/jpillora/requestlog v0.0.0-20181015073026-df8817be5f82/go.mod h1:w8buj+yNfmLEP0ENlbG/FRnK6bVmuhqXnukYCs9sDvY= github.com/jpillora/sizestr v0.0.0-20160130011556-e2ea2fa42fb9 h1:0c9jcgBtHRtDU//jTrcCgWG6UHjMZytiq/3WhraNgUM= github.com/jpillora/sizestr v0.0.0-20160130011556-e2ea2fa42fb9/go.mod h1:1ffp+CRe0eAwwRb0/BownUAjMBsmTLwgAvRbfj9dRwE= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/tomasen/realip v0.0.0-20180522021738-f0c99a92ddce h1:fb190+cK2Xz/dvi9Hv8eCYJYvIGUTN2/KLq1pT6CjEc= github.com/tomasen/realip v0.0.0-20180522021738-f0c99a92ddce/go.mod h1:o8v6yHRoik09Xen7gje4m9ERNah1d1PPsVq1VEx9vE4= golang.org/x/crypto v0.0.0-20181015023909-0c41d7ab0a0e h1:IzypfodbhbnViNUO/MEh0FzCUooG97cIGfdggUrUSyU= diff --git a/main.go b/main.go index 570b508..d83f869 100644 --- a/main.go +++ b/main.go @@ -5,11 +5,13 @@ import ( "fmt" "io/ioutil" "log" + "net/http" "os" "strconv" + "strings" - "github.com/jpillora/chisel/client" - "github.com/jpillora/chisel/server" + chclient "github.com/jpillora/chisel/client" + chserver "github.com/jpillora/chisel/server" chshare "github.com/jpillora/chisel/share" ) @@ -192,6 +194,32 @@ func server(args []string) { } } +type headerFlags struct { + http.Header +} + +func (flag *headerFlags) String() string { + out := "" + for k, v := range flag.Header { + out += fmt.Sprintf("%s: %s\n", k, v) + } + return out +} + +func (flag *headerFlags) Set(arg string) error { + index := strings.Index(arg, ":") + if index < 0 { + return fmt.Errorf(`Invalid header (%s). Should be in the format "HeaderName: HeaderContent"`, arg) + } + if flag.Header == nil { + flag.Header = http.Header{} + } + key := arg[0:index] + value := arg[index+1 : len(arg)] + flag.Header.Set(key, strings.TrimSpace(value)) + return nil +} + var clientHelp = ` Usage: chisel client [options] [remote] [remote] ... @@ -263,13 +291,15 @@ var clientHelp = ` --proxy, An optional HTTP CONNECT proxy which will be used reach the chisel server. Authentication can be specified inside the URL. For example, http://admin:password@my-server.com:8081 + + --header, Set a custom header in the form "HeaderName: HeaderContent". + Can be used multiple times. (e.g --header "Foo: Bar" --header "Hello: World") --hostname, Optionally set the 'Host' header (defaults to the host found in the server url). ` + commonHelp -func client(args []string) { - +func parseClientFlags(args []string) (config chclient.Config, pid *bool, verbose *bool) { flags := flag.NewFlagSet("client", flag.ContinueOnError) fingerprint := flags.String("fingerprint", "", "") @@ -278,9 +308,13 @@ func client(args []string) { maxRetryCount := flags.Int("max-retry-count", -1, "") maxRetryInterval := flags.Duration("max-retry-interval", 0, "") proxy := flags.String("proxy", "", "") - pid := flags.Bool("pid", false, "") + pid = flags.Bool("pid", false, "") hostname := flags.String("hostname", "", "") - verbose := flags.Bool("v", false, "") + headers := headerFlags{ + Header: http.Header{}, + } + flags.Var(&headers, "header", "") + verbose = flags.Bool("v", false, "") flags.Usage = func() { fmt.Print(clientHelp) os.Exit(1) @@ -294,7 +328,11 @@ func client(args []string) { if *auth == "" { *auth = os.Getenv("AUTH") } - c, err := chclient.NewClient(&chclient.Config{ + hostHeader := *hostname + if hostHeader != "" { + headers.Header.Set("Host", hostHeader) + } + config = chclient.Config{ Fingerprint: *fingerprint, Auth: *auth, KeepAlive: *keepalive, @@ -303,8 +341,15 @@ func client(args []string) { HTTPProxy: *proxy, Server: args[0], Remotes: args[1:], - HostHeader: *hostname, - }) + Headers: headers.Header, + } + return +} + +func client(args []string) { + config, pid, verbose := parseClientFlags(args) + c, err := chclient.NewClient(&config) + if err != nil { log.Fatal(err) } diff --git a/main_test.go b/main_test.go new file mode 100644 index 0000000..f57e4fd --- /dev/null +++ b/main_test.go @@ -0,0 +1,37 @@ +package main + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestParseClientFlag(t *testing.T) { + assert := assert.New(t) + args := []string{ + "-fingerprint", "FINGERPRINT", + "-auth", "AUTH-VALUE", + "-hostname", "HOSTNAME", + "-keepalive", "30s", + "-header", "Header1: Foo", + "-header", "Header2: Bar", + "-max-retry-count", "2", + "-max-retry-interval", "12s", + "SERVER", + "REMOTE", + } + config, pid, verbose := parseClientFlags(args) + assert.Equal(config.Fingerprint, "FINGERPRINT") + assert.Equal(config.Headers.Get("Header1"), "Foo") + assert.Equal(config.Headers.Get("Header2"), "Bar") + assert.Equal(config.Headers.Get("Host"), "HOSTNAME") + assert.Equal(config.KeepAlive, 30*time.Second) + assert.Equal(config.MaxRetryInterval, 12*time.Second) + assert.Equal(config.MaxRetryCount, 2) + assert.Equal(config.Auth, "AUTH-VALUE") + assert.Equal(config.Server, "SERVER") + assert.Equal(config.Remotes, []string{"REMOTE"}) + assert.Equal(*pid, false) + assert.Equal(*verbose, false) +}