Add support for multiple tunnels simultaneously (#8)

.github/workflows/playwright.yml

@@ -1,51 +1,51 @@
 name: Playwright Tests
 on:
   push:
-    branches: [ main, master ]
+    branches: [main, master]
   pull_request:
-    branches: [ main, master ]
+    branches: [main, master]
 jobs:
   test:
     timeout-minutes: 60
     runs-on: ubuntu-latest
     steps:
-    - name: Install pnpm
-      run: npm i -g pnpm
-    - uses: actions/checkout@v3
-    - uses: actions/setup-node@v3
-      with:
-        node-version: 16
-        cache: 'pnpm'
-    - name: Install dependencies
-      run: pnpm install
-    - name: Install Playwright Browsers
-      run: pnpm dlx playwright install --with-deps
-    - name: Add hosts to /etc/hosts
-      run: |
+      - name: Install pnpm
+        run: npm i -g pnpm
+      - uses: actions/checkout@v3
+      - uses: actions/setup-node@v3
+        with:
+          node-version: 16
+          cache: "pnpm"
+      - name: Install dependencies
+        run: pnpm install
+      - name: Install Playwright Browsers
+        run: pnpm dlx playwright install --with-deps
+      - name: Add hosts to /etc/hosts
+        run: |
           sudo echo "127.0.0.1 test.localhost" | sudo tee -a /etc/hosts
-    - name: Build client and server
-      run: |
+      - name: Build client and server
+        run: |
           go build -ldflags="-s -w" -o beaver ./cmd/beaver_client
           go build -ldflags="-s -w" -o beaver_server ./cmd/beaver_server
           go build -ldflags="-s -w" -o test_server ./e2e/server.go
-    - name: Start test servers
-      run: |
+      - name: Start test servers
+        run: |
           ./test_server &
           sleep 5
           ./beaver_server --config docs/beaver_server.yaml &
           sleep 5
-    - name: Start test client
-      run: |
-          ./beaver --config docs/beaver_client.yaml --port 9999 --subdomain test &
+      - name: Start test client
+        run: |
+          ./beaver --config docs/beaver_client.yaml http 9999 --subdomain test &
           sleep 3
-    - name: Run Playwright tests
-      run: pnpm dlx playwright test
-    - name: Kill test servers
-      run: |
+      - name: Run Playwright tests
+        run: pnpm dlx playwright test
+      - name: Kill test servers
+        run: |
           make kill-test-servers
-    - uses: actions/upload-artifact@v3
-      if: always()
-      with:
-        name: playwright-report
-        path: playwright-report/
-        retention-days: 30
+      - uses: actions/upload-artifact@v3
+        if: always()
+        with:
+          name: playwright-report
+          path: playwright-report/
+          retention-days: 30

.gitignore

@@ -20,3 +20,4 @@ node_modules
 /playwright/.cache/
 
 dist/
+!beaver_server/

Makefile

@@ -17,7 +17,7 @@ start-test-servers:
 	@go run cmd/beaver_server/main.go --config docs/beaver_server.yaml &
 	@sleep 5
 	@echo "Starting beaver client"
-	@go run cmd/beaver_client/main.go --config docs/beaver_client.yaml --port 9999 --subdomain test &
+	@go run cmd/beaver_client/main.go --config docs/beaver_client.yaml http 9999 --subdomain test &
 	@sleep 5
 
 kill-test-servers:

README.md

@@ -10,12 +10,48 @@
 Download the binary from [releases page](https://github.com/amalshaji/beaver/releases).
 
 ```bash
-beaver - tunnel local ports to public URLs:
+➜  beaver --help
+Tunnel local ports to public URLs
 
 Usage:
-      --config string      Config file path (default "/Users/amalshaji/.beaver/beaver_client.yaml")
+  beaver [flags]
+  beaver [command]
+
+Available Commands:
+  completion  Generate the autocompletion script for the specified shell
+  help        Help about any command
+  http        Tunnel local http servers
+  start       Start tunnels defined in the config file
+
+Flags:
+      --config string   Path to the client config file (default "/Users/amalshaji/.beaver/beaver_client.yaml")
+  -h, --help            help for beaver
+
+Use "beaver [command] --help" for more information about a command.
+
+---
+
+➜  beaver http --help
+Tunnel local http servers
+
+Usage:
+  beaver http [PORT] [flags]
+
+Flags:
+  -h, --help               help for http
       --subdomain string   Subdomain to tunnel http requests (default "<random_subdomain>")
-      --port int           Local http server port (required)
+
+---
+
+➜  beaver start --help
+Start tunnels defined in the config file
+
+Usage:
+  beaver start [--all] or [tunnel1 tunnel2] [flags]
+
+Flags:
+      --all    Start all tunnels listed in the config
+  -h, --help   help for start
 ```
 
 #### Example
@@ -26,7 +62,7 @@ Usage:
 2023/02/05 19:46:07 Tunnel running on https://sccrej.tunnel.example.com
 ```
 
-Now, `https://sccrej.tunnel.example.com ⇄ http://localhost:8000`
+Now, `https://sccrej.tunnel.example.com -> http://localhost:8000`
 
 #### Config
 

client/client.go

@@ -33,11 +33,9 @@ func NewClient(config *Config) (c *Client) {
 
 // Start the Proxy
 func (c *Client) Start(ctx context.Context) {
-	for _, target := range c.Config.Targets {
-		pool := NewPool(c, target)
-		c.pools[target] = pool
-		go pool.Start(ctx)
-	}
+	pool := NewPool(c, c.Config.Target)
+	c.pools[c.Config.id] = pool
+	go pool.Start(ctx)
 }
 
 // Shutdown the Proxy

client/config.go

@@ -1,62 +1,76 @@
 package client
 
 import (
-	"fmt"
 	"os"
 	"strings"
 
 	gonanoid "github.com/matoous/go-nanoid/v2"
 	uuid "github.com/nu7hatch/gouuid"
+
 	"gopkg.in/yaml.v3"
 )
 
-// Config configures an Proxy
+type TunnelConfig struct {
+	Name      string
+	Subdomain string
+	Port      int
+}
+
+type ProxyTunnels struct {
+	Tunnels []TunnelConfig
+}
+
 type Config struct {
 	id               string
 	subdomain        string
 	port             int
 	showWsReadErrors bool
 
-	Targets      []string
+	Target       string
 	PoolIdleSize int
 	PoolMaxSize  int
 	SecretKey    string
 }
 
-// NewConfig creates a new ProxyConfig
-func NewConfig() (config *Config) {
-	config = new(Config)
+// ProxyConfig configures an ProxyConfig
 
+func (config *Config) setDefaults() {
 	id, err := uuid.NewV4()
 	if err != nil {
 		panic(err)
 	}
 	config.id = id.String()
 
-	config.Targets = []string{"wss://t.amal.sh"}
-	config.PoolIdleSize = 1
-	config.PoolMaxSize = 100
+	if config.Target == "" {
+		config.Target = "wss://x.amal.sh"
+	}
+	if config.PoolIdleSize == 0 {
+		config.PoolIdleSize = 1
+	}
+	if config.PoolMaxSize == 0 {
+		config.PoolMaxSize = 100
+	}
 
-	return
 }
 
 // LoadConfiguration loads configuration from a YAML file
-func LoadConfiguration(path, subdomain string, port int, showWsReadErrors bool) (config *Config, err error) {
-	config = NewConfig()
+func LoadConfiguration(configFile string, subdomain string, port int, showWsReadErrors bool) (Config, error) {
+	var config Config
 
-	bytes, err := os.ReadFile(path)
+	bytes, err := os.ReadFile(configFile)
 	if err != nil {
-		return
+		return Config{}, err
 	}
 
-	err = yaml.Unmarshal(bytes, config)
+	err = yaml.Unmarshal(bytes, &config)
 	if err != nil {
-		return
+		return Config{}, err
 	}
-	for i, v := range config.Targets {
-		if !strings.HasSuffix(v, "/register") {
-			config.Targets[i] = fmt.Sprintf("%s/register", v)
-		}
+
+	config.setDefaults()
+
+	if !strings.HasSuffix(config.Target, "/register") {
+		config.Target = config.Target + "/register"
 	}
 
 	if subdomain == "" {
@@ -69,5 +83,5 @@ func LoadConfiguration(path, subdomain string, port int, showWsReadErrors bool)
 	config.port = port
 	config.showWsReadErrors = showWsReadErrors
 
-	return
+	return config, nil
 }

client/connection.go

@@ -44,7 +44,7 @@ func NewConnection(pool *Pool) *Connection {
 // Connect to the IsolatorServer using a HTTP websocket
 func (connection *Connection) Connect(ctx context.Context) (err error) {
 	if connection.IsInitialConnection() {
-		log.Println("Creating tunnel connection")
+		log.Printf("Creating tunnel connection for :%d", connection.pool.client.Config.port)
 	}
 
 	var res *http.Response
@@ -92,11 +92,12 @@ func (connection *Connection) Connect(ctx context.Context) (err error) {
 	}
 
 	if connection.IsInitialConnection() {
-		log.Printf("Tunnel running on %s://%s.%s%s",
+		log.Printf("Tunnel connected %s://%s.%s%s -> http://localhost:%d",
 			httpScheme,
 			connection.pool.client.Config.subdomain,
 			URL.Hostname(),
 			httpPort,
+			connection.pool.client.Config.port,
 		)
 	}
 
@@ -182,7 +183,11 @@ func (connection *Connection) serve(ctx context.Context) {
 			urlPath = urlPath + "?" + req.URL.RawQuery
 		}
 
-		log.Printf("[%s] %d %s", req.Method, resp.StatusCode, urlPath)
+		log.Printf("[%s] %d %s",
+			req.Method,
+			resp.StatusCode,
+			urlPath,
+		)
 
 		// Serialize response
 		jsonResponse, err := json.Marshal(beaver.SerializeHTTPResponse(resp))

client/pool.go

@@ -71,7 +71,7 @@ func (pool *Pool) connector(ctx context.Context) {
 	// Try to reach ideal pool size
 	for i := 0; i < toCreate; i++ {
 		conn := NewConnection(pool)
-		pool.connections = append(pool.connections, conn)
+		pool.add(conn)
 
 		go func() {
 			err := conn.Connect(ctx)

cmd/beaver_client/http.go

@@ -0,0 +1,43 @@
+package main
+
+import (
+	"fmt"
+	"strconv"
+
+	"github.com/amalshaji/beaver/client"
+	"github.com/spf13/cobra"
+)
+
+var (
+	port      int
+	subdomain string
+	httpCmd   = &cobra.Command{
+		Use:   "http [PORT]",
+		Short: "Tunnel local http servers",
+		Args: func(cmd *cobra.Command, args []string) error {
+			if len(args) == 0 {
+				return fmt.Errorf("local server port is required")
+			}
+			if len(args) > 1 {
+				return fmt.Errorf("only one port number is allowed")
+			}
+
+			var err error
+			port, err = strconv.Atoi(args[0])
+			if err != nil {
+				return fmt.Errorf("port must be a number")
+			}
+
+			return nil
+		},
+		Run: func(cmd *cobra.Command, args []string) {
+			var tunnels = make([]client.TunnelConfig, 0)
+			tunnels = append(tunnels, client.TunnelConfig{Port: port, Subdomain: subdomain})
+			startTunnels(tunnels)
+		},
+	}
+)
+
+func init() {
+	httpCmd.Flags().StringVar(&subdomain, "subdomain", "", "Subdomain to tunnel http requests (default \"<random_subdomain>\")")
+}