AUTH-1511: Add custom headers for ssh command

This commit is contained in:
Austin Cherry 2019-02-07 10:56:33 -06:00
parent c9f1c1ec70
commit 92defa26d4
5 changed files with 50 additions and 13 deletions

View File

@ -34,13 +34,13 @@ func (c *StdinoutStream) Write(p []byte) (int, error) {
// StartClient will copy the data from stdin/stdout over a WebSocket connection // StartClient will copy the data from stdin/stdout over a WebSocket connection
// to the edge (originURL) // to the edge (originURL)
func StartClient(logger *logrus.Logger, originURL string, stream io.ReadWriter) error { func StartClient(logger *logrus.Logger, originURL string, stream io.ReadWriter, headers http.Header) error {
return serveStream(logger, originURL, stream) return serveStream(logger, originURL, stream, headers)
} }
// StartServer will setup a server on a specified port and copy data over a WebSocket connection // StartServer will setup a server on a specified port and copy data over a WebSocket connection
// to the edge (originURL) // to the edge (originURL)
func StartServer(logger *logrus.Logger, address, originURL string, shutdownC <-chan struct{}) error { func StartServer(logger *logrus.Logger, address, originURL string, shutdownC <-chan struct{}, headers http.Header) error {
listener, err := net.Listen("tcp", address) listener, err := net.Listen("tcp", address)
if err != nil { if err != nil {
logger.WithError(err).Error("failed to start forwarding server") logger.WithError(err).Error("failed to start forwarding server")
@ -56,20 +56,20 @@ func StartServer(logger *logrus.Logger, address, originURL string, shutdownC <-c
if err != nil { if err != nil {
return err return err
} }
go serveConnection(logger, conn, originURL) go serveConnection(logger, conn, originURL, headers)
} }
} }
} }
// serveConnection handles connections for the StartServer call // serveConnection handles connections for the StartServer call
func serveConnection(logger *logrus.Logger, c net.Conn, originURL string) { func serveConnection(logger *logrus.Logger, c net.Conn, originURL string, headers http.Header) {
defer c.Close() defer c.Close()
serveStream(logger, originURL, c) serveStream(logger, originURL, c, headers)
} }
// serveStream will serve the data over the WebSocket stream // serveStream will serve the data over the WebSocket stream
func serveStream(logger *logrus.Logger, originURL string, conn io.ReadWriter) error { func serveStream(logger *logrus.Logger, originURL string, conn io.ReadWriter, headers http.Header) error {
wsConn, err := createWebsocketStream(originURL) wsConn, err := createWebsocketStream(originURL, headers)
if err != nil { if err != nil {
logger.WithError(err).Errorf("failed to connect to %s\n", originURL) logger.WithError(err).Errorf("failed to connect to %s\n", originURL)
return err return err
@ -84,11 +84,12 @@ func serveStream(logger *logrus.Logger, originURL string, conn io.ReadWriter) er
// createWebsocketStream will create a WebSocket connection to stream data over // createWebsocketStream will create a WebSocket connection to stream data over
// It also handles redirects from Access and will present that flow if // It also handles redirects from Access and will present that flow if
// the token is not present on the request // the token is not present on the request
func createWebsocketStream(originURL string) (*websocket.Conn, error) { func createWebsocketStream(originURL string, headers http.Header) (*websocket.Conn, error) {
req, err := http.NewRequest(http.MethodGet, originURL, nil) req, err := http.NewRequest(http.MethodGet, originURL, nil)
if err != nil { if err != nil {
return nil, err return nil, err
} }
req.Header = headers
wsConn, resp, err := websocket.ClientConnect(req, nil) wsConn, resp, err := websocket.ClientConnect(req, nil)
if err != nil && resp != nil && resp.StatusCode > 300 { if err != nil && resp != nil && resp.StatusCode > 300 {

View File

@ -48,7 +48,7 @@ func TestStartClient(t *testing.T) {
defer ts.Close() defer ts.Close()
buf := newTestStream() buf := newTestStream()
err := StartClient(logger, "http://"+ts.Listener.Addr().String(), buf) err := StartClient(logger, "http://"+ts.Listener.Addr().String(), buf, nil)
assert.NoError(t, err) assert.NoError(t, err)
buf.Write([]byte(message)) buf.Write([]byte(message))
@ -66,7 +66,7 @@ func TestStartServer(t *testing.T) {
defer ts.Close() defer ts.Close()
go func() { go func() {
err := StartServer(logger, listenerAddress, "http://"+ts.Listener.Addr().String(), shutdownC) err := StartServer(logger, listenerAddress, "http://"+ts.Listener.Addr().String(), shutdownC, nil)
if err != nil { if err != nil {
t.Fatalf("Error starting server: %v", err) t.Fatalf("Error starting server: %v", err)
} }

View File

@ -1,7 +1,9 @@
package access package access
import ( import (
"net/http"
"net/url" "net/url"
"strings"
"github.com/cloudflare/cloudflared/carrier" "github.com/cloudflare/cloudflared/carrier"
"github.com/cloudflare/cloudflared/cmd/cloudflared/config" "github.com/cloudflare/cloudflared/cmd/cloudflared/config"
@ -19,6 +21,7 @@ func ssh(c *cli.Context) error {
if err != nil || c.String("hostname") == "" { if err != nil || c.String("hostname") == "" {
return cli.ShowCommandHelp(c, "ssh") return cli.ShowCommandHelp(c, "ssh")
} }
headers := buildRequestHeaders(c.StringSlice("header"))
if c.NArg() > 0 || c.IsSet("url") { if c.NArg() > 0 || c.IsSet("url") {
localForwarder, err := config.ValidateUrl(c) localForwarder, err := config.ValidateUrl(c)
@ -31,8 +34,19 @@ func ssh(c *cli.Context) error {
logger.WithError(err).Error("Error validating origin URL") logger.WithError(err).Error("Error validating origin URL")
return errors.Wrap(err, "error validating origin URL") return errors.Wrap(err, "error validating origin URL")
} }
return carrier.StartServer(logger, forwarder.Host, "https://"+hostname, shutdownC) return carrier.StartServer(logger, forwarder.Host, "https://"+hostname, shutdownC, headers)
} }
return carrier.StartClient(logger, "https://"+hostname, &carrier.StdinoutStream{}) return carrier.StartClient(logger, "https://"+hostname, &carrier.StdinoutStream{}, headers)
}
func buildRequestHeaders(values []string) http.Header {
headers := make(http.Header)
for _, valuePair := range values {
split := strings.Split(valuePair, ":")
if len(split) > 1 {
headers.Add(strings.TrimSpace(split[0]), strings.TrimSpace(split[1]))
}
}
return headers
} }

View File

@ -0,0 +1,18 @@
package access
import (
"net/http"
"testing"
"github.com/stretchr/testify/assert"
)
func TestBuildRequestHeaders(t *testing.T) {
headers := make(http.Header)
headers.Add("client", "value")
headers.Add("secret", "safe-value")
values := buildRequestHeaders([]string{"client: value", "secret: safe-value", "trash"})
assert.Equal(t, headers.Get("client"), values.Get("client"))
assert.Equal(t, headers.Get("secret"), values.Get("secret"))
}

View File

@ -96,6 +96,10 @@ func Commands() []*cli.Command {
&cli.StringFlag{ &cli.StringFlag{
Name: "url", Name: "url",
}, },
&cli.StringSliceFlag{
Name: "header",
Aliases: []string{"H"},
},
}, },
}, },
{ {