From b592d63a083cd340910aedecae9c2e112ad779fb Mon Sep 17 00:00:00 2001
From: Andrew Dunstall <andydunstall@hotmail.co.uk>
Date: Tue, 16 Jul 2024 18:18:29 +0100
Subject: [PATCH] forward: add client authentication

Adds a '--connect.token' option to 'piko forward' to authenticate client
connections.
---
 cli/forward/command.go   |  1 +
 client/dialer.go         | 12 +++++++++++-
 forward/config/config.go | 11 +++++++++++
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/cli/forward/command.go b/cli/forward/command.go
index a0c8cda..23fb793 100644
--- a/cli/forward/command.go
+++ b/cli/forward/command.go
@@ -90,6 +90,7 @@ func runForward(conf *config.Config, logger log.Logger) error {
 	}
 	dialer := &client.Dialer{
 		URL:       connectURL,
+		Token:     conf.Connect.Token,
 		TLSConfig: connectTLSConfig,
 	}
 
diff --git a/client/dialer.go b/client/dialer.go
index 7206e64..dc24bbb 100644
--- a/client/dialer.go
+++ b/client/dialer.go
@@ -22,6 +22,12 @@ type Dialer struct {
 	// 'http://localhost:8000'.
 	URL *url.URL
 
+	// Token configures the API key token to authenticate the listener with the
+	// Piko server.
+	//
+	// Defaults to no authentication.
+	Token string
+
 	// TLSConfig specifies the TLS configuration to use with the Piko server.
 	//
 	// If nil, the default configuration is used.
@@ -33,7 +39,11 @@ type Dialer struct {
 func (d *Dialer) Dial(ctx context.Context, endpointID string) (net.Conn, error) {
 	// Dialing is simply opening a WebSocket connection to the target endpoint,
 	// then wrapping the WebSocket in a net.Conn.
-	return websocket.Dial(ctx, d.dialURL(endpointID))
+	return websocket.Dial(
+		ctx,
+		d.dialURL(endpointID),
+		websocket.WithToken(d.Token),
+	)
 }
 
 func (d *Dialer) dialURL(endpointID string) string {
diff --git a/forward/config/config.go b/forward/config/config.go
index ddeceb6..622bf73 100644
--- a/forward/config/config.go
+++ b/forward/config/config.go
@@ -103,6 +103,9 @@ type ConnectConfig struct {
 	// URL is the Piko server URL to connect to.
 	URL string
 
+	// Token is a token to authenticate with the Piko server.
+	Token string
+
 	// Timeout is the timeout attempting to connect to the Piko server.
 	Timeout time.Duration `json:"timeout" yaml:"timeout"`
 
@@ -132,6 +135,14 @@ The Piko server URL to connect to. Note this must be configured to use the
 Piko server 'proxy' port.`,
 	)
 
+	fs.StringVar(
+		&c.Token,
+		"connect.token",
+		c.Token,
+		`
+Token is a token to authenticate with the Piko server.`,
+	)
+
 	fs.DurationVar(
 		&c.Timeout,
 		"connect.timeout",