client cleanup

cmd/link.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"fmt"
 	"net"
+	"net/url"
 	"sync"
 
 	"github.com/gorilla/websocket"
@@ -19,43 +20,88 @@ type TunnelMessage struct {
 var serverAddr string
 
 var linkCmd = &cobra.Command{
-	Use:   "link <vhost_location> <host:port>",
+	Use:   "link <name> <host:port>",
 	Short: "Create a tunnel link",
 	Args:  cobra.ExactArgs(2),
 	Run: func(cmd *cobra.Command, args []string) {
-		vhostLoc := args[0]
+		tunnelName := args[0]
-		hostPort := args[1]
+		tunnelTarget := args[1]
 
-		fmt.Printf("Creating TCP tunnel: %s -> %s\n", vhostLoc, hostPort)
+		// Create Tunnel
+		tunnel, err := NewTunnel(tunnelName, tunnelTarget, serverAddr)
+		if err != nil {
+			log.Fatal("Failed to start tunnel:", err)
+		}
 
-		if err := startTCPTunnel(vhostLoc, hostPort); err != nil {
+		// Start Tunnel
+		log.Infof("Creating TCP tunnel: %s -> %s\n", tunnelName, tunnelTarget)
+		if err := tunnel.Start(); err != nil {
 			log.Fatal("Failed to start tunnel:", err)
 		}
 	},
 }
 
 func init() {
-	linkCmd.Flags().StringVarP(&serverAddr, "server", "s", "localhost:8080", "Conduit server address")
+	linkCmd.Flags().StringVarP(&serverAddr, "server", "s", "http://localhost:8080", "Conduit server address")
 }
 
-func startTCPTunnel(vhost, hostPort string) error {
+type TunnelConfig struct {
-	wsURL := fmt.Sprintf("ws://%s/_conduit/tunnel?vhost=%s", serverAddr, vhost)
+	// The conduit server address, e.g. https://conduit.example.com
-	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	ServerAddress string `default:"http://localhost:8080"`
+}
+
+func NewTunnel(tunnelName, tunnelTarget, serverAddress string) (*Tunnel, error) {
+	// Parse Server URL
+	serverURL, err := url.Parse(serverAddress)
 	if err != nil {
-		return fmt.Errorf("failed to connect: %v", err)
+		return nil, err
 	}
-	defer conn.Close()
 
-	fmt.Printf("TCP Tunnel active! %s.example.com -> %s\n", vhost, hostPort)
+	// Parse Scheme
+	var wsScheme string
+	switch serverURL.Scheme {
+	case "https":
+		wsScheme = "wss"
+	case "http":
+		wsScheme = "ws"
+	default:
+		return nil, fmt.Errorf("unsupported scheme: %s", serverURL.Scheme)
+	}
 
-	// Track active connections
+	// Connect Server WS
-	connections := make(map[string]net.Conn)
+	wsURL := fmt.Sprintf("%s://%s/_conduit/tunnel?vhost=%s", wsScheme, serverURL.Host, tunnelName)
-	var connMutex sync.RWMutex
+	serverConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to connect: %v", err)
+	}
 
-	// Handle messages from server
+	return &Tunnel{
+		name:       tunnelName,
+		target:     tunnelTarget,
+		serverConn: serverConn,
+		localConns: make(map[string]net.Conn),
+	}, nil
+
+}
+
+type Tunnel struct {
+	name   string
+	target string
+
+	serverConn *websocket.Conn
+	localConns map[string]net.Conn
+	mu         sync.RWMutex
+}
+
+func (t *Tunnel) Start() error {
+	log.Infof("TCP Tunnel active! %s.example.com -> %s\n", t.name, t.target)
+	defer t.serverConn.Close()
+
+	// Handle Messages
 	for {
+		// Read Message
 		var msg TunnelMessage
-		err := conn.ReadJSON(&msg)
+		err := t.serverConn.ReadJSON(&msg)
 		if err != nil {
 			log.Errorf("Error reading from tunnel: %v", err)
 			break
@@ -63,28 +109,63 @@ func startTCPTunnel(vhost, hostPort string) error {
 
 		switch msg.Type {
 		case "data":
-			connMutex.RLock()
+			localConn, err := t.getLocalConn(msg.StreamID)
-			localConn, exists := connections[msg.StreamID]
-			connMutex.RUnlock()
-
-			if !exists {
-				// New connection
-				localConn, err = net.Dial("tcp", hostPort)
 			if err != nil {
-					log.Errorf("Failed to connect to %s: %v", hostPort, err)
+				log.Errorf("Failed to get local connection: %v", err)
 				continue
 			}
 
-				connMutex.Lock()
+			// Write data to local connection
-				connections[msg.StreamID] = localConn
+			if _, err := localConn.Write(msg.Data); err != nil {
-				connMutex.Unlock()
+				log.Errorf("Error writing to local connection: %v", err)
+				localConn.Close()
+				t.mu.Lock()
+				delete(t.localConns, msg.StreamID)
+				t.mu.Unlock()
+			}
 
-				// Start reading from local connection
+		case "close":
-				go func(streamID string, lConn net.Conn) {
+			t.mu.Lock()
+			if localConn, exists := t.localConns[msg.StreamID]; exists {
+				localConn.Close()
+				delete(t.localConns, msg.StreamID)
+			}
+			t.mu.Unlock()
+		}
+	}
+
+	return nil
+}
+
+func (t *Tunnel) getLocalConn(streamID string) (net.Conn, error) {
+	// Get Cached Connection
+	t.mu.RLock()
+	localConn, exists := t.localConns[streamID]
+	t.mu.RUnlock()
+	if exists {
+		return localConn, nil
+	}
+
+	// Initiate Connection & Cache
+	localConn, err := net.Dial("tcp", t.target)
+	if err != nil {
+		log.Errorf("Failed to connect to %s: %v", t.target, err)
+		return nil, err
+	}
+	t.mu.Lock()
+	t.localConns[streamID] = localConn
+	t.mu.Unlock()
+
+	// Start Response Relay & Return Connection
+	go t.startResponseRelay(streamID, localConn)
+	return localConn, nil
+}
+
+func (t *Tunnel) startResponseRelay(streamID string, lConn net.Conn) {
 	defer func() {
-						connMutex.Lock()
+		t.mu.Lock()
-						delete(connections, streamID)
+		delete(t.localConns, streamID)
-						connMutex.Unlock()
+		t.mu.Unlock()
 		lConn.Close()
 	}()
 
@@ -101,31 +182,8 @@ func startTCPTunnel(vhost, hostPort string) error {
 			Data:     buffer[:n],
 		}
 
-						if err := conn.WriteJSON(response); err != nil {
+		if err := t.serverConn.WriteJSON(response); err != nil {
 			break
 		}
 	}
-				}(msg.StreamID, localConn)
-			}
-
-			// Write data to local connection
-			if _, err := localConn.Write(msg.Data); err != nil {
-				log.Errorf("Error writing to local connection: %v", err)
-				localConn.Close()
-				connMutex.Lock()
-				delete(connections, msg.StreamID)
-				connMutex.Unlock()
-			}
-
-		case "close":
-			connMutex.Lock()
-			if localConn, exists := connections[msg.StreamID]; exists {
-				localConn.Close()
-				delete(connections, msg.StreamID)
-			}
-			connMutex.Unlock()
-		}
-	}
-
-	return nil
 }

server/server.go

@@ -171,7 +171,7 @@ func (s *Server) proxyRawConnection(clientConn net.Conn, tunnelConn *TunnelConne
 		}
 	}()
 
-	// Handle tunnel -> client (read from response channel)
+	// Return Client Response Data
 	for data := range responseChan {
 		if _, err := clientConn.Write(data); err != nil {
 			break