chore: better source tracking

.gitignore

@@ -1 +1,2 @@
 cover.html
+.DS_Store

client/client.go

@@ -2,17 +2,15 @@ package client
 
 import (
 	"fmt"
-	"net"
 	"net/url"
-	"sync"
 
 	"github.com/gorilla/websocket"
 	log "github.com/sirupsen/logrus"
 	"reichard.io/conduit/config"
-	"reichard.io/conduit/types"
+	"reichard.io/conduit/tunnel"
 )
 
-func NewTunnel(cfg *config.ClientConfig) (*Tunnel, error) {
+func NewTunnel(cfg *config.ClientConfig) (*tunnel.Tunnel, error) {
 	// Parse Server URL
 	serverURL, err := url.Parse(cfg.ServerAddress)
 	if err != nil {
@@ -43,117 +41,5 @@ func NewTunnel(cfg *config.ClientConfig) (*Tunnel, error) {
 		return nil, fmt.Errorf("failed to connect: %v", err)
 	}
 
-	return &Tunnel{
+	return tunnel.NewClientTunnel(cfg.TunnelName, cfg.TunnelTarget, serverURL, serverConn)
-		name:       cfg.TunnelName,
-		target:     cfg.TunnelTarget,
-		serverURL:  serverURL,
-		serverConn: serverConn,
-		localConns: make(map[string]net.Conn),
-	}, nil
-
-}
-
-type Tunnel struct {
-	name      string
-	target    string
-	serverURL *url.URL
-
-	serverConn *websocket.Conn
-	localConns map[string]net.Conn
-	mu         sync.RWMutex
-}
-
-func (t *Tunnel) Start() error {
-	log.Infof("starting tunnel: %s.%s -> %s\n", t.name, t.serverURL.Hostname(), t.target)
-	defer t.serverConn.Close()
-
-	// Handle Messages
-	for {
-		// Read Message
-		var msg types.Message
-		err := t.serverConn.ReadJSON(&msg)
-		if err != nil {
-			log.Errorf("error reading from tunnel: %v", err)
-			break
-		}
-
-		switch msg.Type {
-		case types.MessageTypeData:
-			localConn, err := t.getLocalConn(msg.StreamID)
-			if err != nil {
-				log.Errorf("failed to get local connection: %v", err)
-				continue
-			}
-
-			// Write data to local connection
-			if _, err := localConn.Write(msg.Data); err != nil {
-				log.Errorf("error writing to local connection: %v", err)
-				localConn.Close()
-				t.mu.Lock()
-				delete(t.localConns, msg.StreamID)
-				t.mu.Unlock()
-			}
-
-		case types.MessageTypeClose:
-			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, localConn net.Conn) {
-	defer func() {
-		t.mu.Lock()
-		delete(t.localConns, streamID)
-		t.mu.Unlock()
-		localConn.Close()
-	}()
-
-	buffer := make([]byte, 4096)
-	for {
-		n, err := localConn.Read(buffer)
-		if err != nil {
-			break
-		}
-
-		response := types.Message{
-			Type:     types.MessageTypeData,
-			StreamID: streamID,
-			Data:     buffer[:n],
-		}
-
-		if err := t.serverConn.WriteJSON(response); err != nil {
-			break
-		}
-	}
 }

cmd/tunnel.go

@@ -22,12 +22,7 @@ var tunnelCmd = &cobra.Command{
 		if err != nil {
 			log.Fatal("failed to create tunnel:", err)
 		}
-
+		tunnel.Start()
-		// Start Tunnel
-		log.Infof("creating TCP tunnel: %s -> %s", cfg.TunnelName, cfg.TunnelTarget)
-		if err := tunnel.Start(); err != nil {
-			log.Fatal("failed to start tunnel:", err)
-		}
 	},
 }
 

config/config.go

@@ -111,7 +111,7 @@ func GetVersion() string {
 func getConfigValue(cmdFlags *pflag.FlagSet, def ConfigDef) string {
 	// 1. Get Flags First
 	if cmdFlags != nil {
-		if val, err := cmdFlags.GetString(def.Key); err == nil && val != "" {
+		if val, err := cmdFlags.GetString(def.Key); err == nil && val != "" && val != def.Default {
 			return val
 		}
 	}

pkg/maps/map.go

@@ -0,0 +1,51 @@
+package maps
+
+import (
+	"iter"
+	"sync"
+)
+
+type Map[K comparable, V any] struct {
+	items map[K]V
+	mu    sync.RWMutex
+}
+
+func New[K comparable, V any]() *Map[K, V] {
+	return &Map[K, V]{items: make(map[K]V)}
+}
+
+func (m *Map[K, V]) Get(key K) (V, bool) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	v, ok := m.items[key]
+	return v, ok
+}
+
+func (m *Map[K, V]) Set(key K, value V) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.items[key] = value
+}
+
+func (m *Map[K, V]) Delete(key K) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	delete(m.items, key)
+}
+
+func (m *Map[K, V]) HasKey(key K) bool {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	_, ok := m.items[key]
+	return ok
+}
+
+func (m *Map[K, V]) Entries() iter.Seq2[K, V] {
+	return func(yield func(K, V) bool) {
+		for k, v := range m.items {
+			if !yield(k, v) {
+				return
+			}
+		}
+	}
+}

server/raw_http_response_writer.go

@@ -7,25 +7,25 @@ import (
 	"net/http"
 )
 
-var _ http.ResponseWriter = (*connResponseWriter)(nil)
+var _ http.ResponseWriter = (*rawHTTPResponseWriter)(nil)
 
-type connResponseWriter struct {
+type rawHTTPResponseWriter struct {
 	conn   net.Conn
 	header http.Header
 }
 
-func (f *connResponseWriter) Header() http.Header {
+func (f *rawHTTPResponseWriter) Header() http.Header {
 	if f.header == nil {
 		f.header = make(http.Header)
 	}
 	return f.header
 }
 
-func (f *connResponseWriter) Write(data []byte) (int, error) {
+func (f *rawHTTPResponseWriter) Write(data []byte) (int, error) {
 	return f.conn.Write(data)
 }
 
-func (f *connResponseWriter) WriteHeader(statusCode int) {
+func (f *rawHTTPResponseWriter) WriteHeader(statusCode int) {
 	// Write Status
 	status := fmt.Sprintf("HTTP/1.1 %d %s\r\n", statusCode, http.StatusText(statusCode))
 	_, _ = f.conn.Write([]byte(status))
@@ -41,7 +41,7 @@ func (f *connResponseWriter) WriteHeader(statusCode int) {
 	_, _ = f.conn.Write([]byte("\r\n"))
 }
 
-func (f *connResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+func (f *rawHTTPResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	// Return Raw Connection & ReadWriter
 	rw := bufio.NewReadWriter(bufio.NewReader(f.conn), bufio.NewWriter(f.conn))
 	return f.conn, rw, nil

server/reconstructed_conn.go

@@ -0,0 +1,30 @@
+package server
+
+import (
+	"bytes"
+	"io"
+	"net"
+)
+
+var _ io.ReadWriteCloser = (*reconstructedConn)(nil)
+
+// reconstructedConn wraps a net.Conn and overrides Read to handle captured data.
+type reconstructedConn struct {
+	net.Conn
+	reader io.Reader
+}
+
+// Read reads from the reconstructed reader (captured data + original conn).
+func (rc *reconstructedConn) Read(p []byte) (n int, err error) {
+	return rc.reader.Read(p)
+}
+
+// newReconstructedConn creates a reconstructed connection that replays captured data
+// before reading from the original connection.
+func newReconstructedConn(conn net.Conn, capturedData *bytes.Buffer) net.Conn {
+	allReader := io.MultiReader(capturedData, conn)
+	return &reconstructedConn{
+		Conn:   conn,
+		reader: allReader,
+	}
+}

server/server.go

@@ -11,13 +11,13 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
 	log "github.com/sirupsen/logrus"
 	"reichard.io/conduit/config"
-	"reichard.io/conduit/types"
+	"reichard.io/conduit/pkg/maps"
+	"reichard.io/conduit/tunnel"
 )
 
 type InfoResponse struct {
@@ -30,19 +30,12 @@ type TunnelInfo struct {
 	Target string `json:"target"`
 }
 
-type TunnelConnection struct {
-	*websocket.Conn
-	name    string
-	streams map[string]chan []byte
-}
-
 type Server struct {
 	host string
 	cfg  *config.ServerConfig
-	mu   sync.RWMutex
 
 	upgrader websocket.Upgrader
-	tunnels  map[string]*TunnelConnection
+	tunnels  *maps.Map[string, *tunnel.Tunnel]
 }
 
 func NewServer(cfg *config.ServerConfig) (*Server, error) {
@@ -56,7 +49,7 @@ func NewServer(cfg *config.ServerConfig) (*Server, error) {
 	return &Server{
 		cfg:     cfg,
 		host:    serverURL.Host,
-		tunnels: make(map[string]*TunnelConnection),
+		tunnels: maps.New[string, *tunnel.Tunnel](),
 		upgrader: websocket.Upgrader{
 			CheckOrigin: func(r *http.Request) bool {
 				return true
@@ -79,7 +72,7 @@ func (s *Server) Start() error {
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
-			log.Printf("error accepting connection: %v", err)
+			log.WithError(err).Error("error accepting connection")
 			continue
 		}
 
@@ -90,14 +83,12 @@ func (s *Server) Start() error {
 func (s *Server) getInfo(w http.ResponseWriter, _ *http.Request) {
 	// Get Tunnels
 	var allTunnels []TunnelInfo
-	s.mu.RLock()
+	for t, c := range s.tunnels.Entries() {
-	for t, c := range s.tunnels {
 		allTunnels = append(allTunnels, TunnelInfo{
 			Name:   t,
-			Target: c.RemoteAddr().String(),
+			Target: c.Source(),
 		})
 	}
-	s.mu.RUnlock()
 
 	// Create Response
 	d, err := json.MarshalIndent(InfoResponse{
@@ -105,72 +96,17 @@ func (s *Server) getInfo(w http.ResponseWriter, _ *http.Request) {
 		Version: config.GetVersion(),
 	}, "", "  ")
 	if err != nil {
+		log.WithError(err).Error("failed to marshal info")
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}
 
 	// Send Response
 	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(http.StatusOK)
 	_, _ = w.Write(d)
 }
 
-func (s *Server) proxyRawConnection(clientConn net.Conn, tunnelConn *TunnelConnection, dataReader io.Reader) {
-	defer clientConn.Close()
-
-	// Create Identifiers
-	streamID := fmt.Sprintf("stream_%d", time.Now().UnixNano())
-	responseChan := make(chan []byte, 100)
-
-	// Register Stream
-	s.mu.Lock()
-	if tunnelConn.streams == nil {
-		tunnelConn.streams = make(map[string]chan []byte)
-	}
-	tunnelConn.streams[streamID] = responseChan
-	s.mu.Unlock()
-
-	// Clean Up
-	defer func() {
-		s.mu.Lock()
-		delete(tunnelConn.streams, streamID)
-		close(responseChan)
-		s.mu.Unlock()
-
-		// Send Close
-		closeMsg := types.Message{
-			Type:     types.MessageTypeClose,
-			StreamID: streamID,
-		}
-		_ = tunnelConn.WriteJSON(closeMsg)
-	}()
-
-	// Read & Send Chunks
-	go func() {
-		buffer := make([]byte, 4096)
-		for {
-			n, err := dataReader.Read(buffer)
-			if err != nil {
-				return
-			}
-
-			if err := tunnelConn.WriteJSON(types.Message{
-				Type:     types.MessageTypeData,
-				StreamID: streamID,
-				Data:     buffer[:n],
-			}); err != nil {
-				return
-			}
-		}
-	}()
-
-	// Return Response Data
-	for data := range responseChan {
-		if _, err := clientConn.Write(data); err != nil {
-			break
-		}
-	}
-}
-
 func (s *Server) handleRawConnection(conn net.Conn) {
 	defer conn.Close()
 
@@ -183,7 +119,7 @@ func (s *Server) handleRawConnection(conn net.Conn) {
 	bufReader := bufio.NewReader(teeReader)
 
 	// Create HTTP Request & Writer
-	w := &connResponseWriter{conn: conn}
+	w := &rawHTTPResponseWriter{conn: conn}
 	r, err := http.ReadRequest(bufReader)
 	if err != nil {
 		w.WriteHeader(http.StatusBadRequest)
@@ -199,30 +135,45 @@ func (s *Server) handleRawConnection(conn net.Conn) {
 	}
 
 	// Extract Subdomain
-	subdomain := strings.TrimSuffix(strings.Replace(r.Host, s.host, "", 1), ".")
+	tunnelName := strings.TrimSuffix(strings.Replace(r.Host, s.host, "", 1), ".")
-	if strings.Count(subdomain, ".") != 0 {
+	if strings.Count(tunnelName, ".") != 0 {
 		w.WriteHeader(http.StatusBadRequest)
 		_, _ = fmt.Fprintf(w, "cannot tunnel nested subdomains: %s", r.Host)
 		return
 	}
 
+	// Get True Host
+	remoteHost := conn.RemoteAddr().String()
+	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
+		remoteHost = xff
+	}
+	r.RemoteAddr = remoteHost
+
 	// Handle Control Endpoints
-	if subdomain == "" {
+	if tunnelName == "" {
 		s.handleAsHTTP(w, r)
 		return
 	}
 
 	// Handle Tunnels
-	s.mu.RLock()
+	conduitTunnel, exists := s.tunnels.Get(tunnelName)
-	tunnelConn, exists := s.tunnels[subdomain]
+	if !exists {
-	s.mu.RUnlock()
+		w.WriteHeader(http.StatusNotFound)
-	if exists {
+		_, _ = fmt.Fprintf(w, "unknown tunnel: %s", tunnelName)
-		log.Infof("relaying %s to tunnel", subdomain)
+		return
-
-		// Reconstruct Data & Proxy Connection
-		allReader := io.MultiReader(&capturedData, r.Body)
-		s.proxyRawConnection(conn, tunnelConn, allReader)
 	}
+
+	// Add & Start Stream
+	reconstructedConn := newReconstructedConn(conn, &capturedData)
+	streamID := fmt.Sprintf("stream_%d", time.Now().UnixNano())
+	if err := conduitTunnel.AddStream(streamID, reconstructedConn); err != nil {
+		w.WriteHeader(http.StatusInternalServerError)
+		_, _ = fmt.Fprintf(w, "failed to add stream: %v", err)
+		return
+	}
+
+	log.Infof("tunnel %q connection from %s", tunnelName, r.RemoteAddr)
+	_ = conduitTunnel.StartStream(streamID, r.RemoteAddr)
 }
 
 func (s *Server) handleAsHTTP(w http.ResponseWriter, r *http.Request) {
@@ -245,40 +196,6 @@ func (s *Server) handleAsHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-func (s *Server) handleTunnelMessages(tunnel *TunnelConnection) {
-	for {
-		var msg types.Message
-		err := tunnel.ReadJSON(&msg)
-		if err != nil {
-			return
-		}
-
-		if msg.StreamID == "" {
-			log.Infof("tunnel %s missing streamID", tunnel.name)
-			continue
-		}
-
-		switch msg.Type {
-		case types.MessageTypeClose:
-			return
-		case types.MessageTypeData:
-			s.mu.RLock()
-			streamChan, exists := tunnel.streams[msg.StreamID]
-			if !exists {
-				log.Infof("stream %s does not exist", msg.StreamID)
-				s.mu.RUnlock()
-				continue
-			}
-
-			select {
-			case streamChan <- msg.Data:
-			case <-time.After(time.Second):
-				log.Warnf("stream %s channel full, dropping data", msg.StreamID)
-			}
-			s.mu.RUnlock()
-		}
-	}
-}
 func (s *Server) createTunnel(w http.ResponseWriter, r *http.Request) {
 	// Get Tunnel Name
 	tunnelName := r.URL.Query().Get("tunnelName")
@@ -289,7 +206,7 @@ func (s *Server) createTunnel(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Validate Unique
-	if _, exists := s.tunnels[tunnelName]; exists {
+	if _, exists := s.tunnels.Get(tunnelName); exists {
 		w.WriteHeader(http.StatusConflict)
 		_, _ = w.Write([]byte("Tunnel already registered"))
 		return
@@ -302,26 +219,16 @@ func (s *Server) createTunnel(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Create & Cache TunnelConnection
+	// Create Tunnel
-	tunnel := &TunnelConnection{
+	conduitTunnel := tunnel.NewServerTunnel(tunnelName, wsConn)
-		Conn:    wsConn,
+	s.tunnels.Set(tunnelName, conduitTunnel)
-		name:    tunnelName,
+	log.Infof("tunnel %q created from %s", tunnelName, r.RemoteAddr)
-		streams: make(map[string]chan []byte),
-	}
-	s.mu.Lock()
-	s.tunnels[tunnelName] = tunnel
-	s.mu.Unlock()
-	log.Infof("tunnel established: %s", tunnelName)
 
-	// Keep connection alive and handle cleanup
+	// Start Tunnel - This is blocking
-	defer func() {
+	conduitTunnel.Start()
-		s.mu.Lock()
-		delete(s.tunnels, tunnelName)
-		s.mu.Unlock()
-		_ = wsConn.Close()
-		log.Infof("tunnel closed: %s", tunnelName)
-	}()
 
-	// Handle tunnel messages
+	// Cleanup Tunnel
-	s.handleTunnelMessages(tunnel)
+	s.tunnels.Delete(tunnelName)
+	_ = wsConn.Close()
+	log.Infof("tunnel %q closed from %s", tunnelName, r.RemoteAddr)
 }

tunnel/http.go

@@ -0,0 +1,104 @@
+package tunnel
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"sync"
+)
+
+func HTTPConnectionBuilder(targetURL *url.URL) (ConnBuilder, error) {
+	multiConnListener := newMultiConnListener()
+
+	// Create Reverse Proxy
+	proxy := &httputil.ReverseProxy{
+		Director: func(req *http.Request) {
+			req.Host = targetURL.Host
+			req.URL.Host = targetURL.Host
+			req.URL.Scheme = targetURL.Scheme
+		},
+		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
+			http.Error(w, fmt.Sprintf("Proxy error: %v", err), http.StatusBadGateway)
+		},
+	}
+
+	// Start HTTP Proxy
+	go func() {
+		defer multiConnListener.Close()
+		_ = http.Serve(multiConnListener, proxy)
+	}()
+
+	// Return Connection Builder
+	return func() (conn io.ReadWriteCloser, err error) {
+		clientConn, serverConn := net.Pipe()
+
+		if err := multiConnListener.addConn(serverConn); err != nil {
+			_ = clientConn.Close()
+			_ = serverConn.Close()
+			return nil, err
+		}
+
+		return clientConn, nil
+	}, nil
+}
+
+type multiConnListener struct {
+	connCh chan net.Conn
+	closed chan struct{}
+	once   sync.Once
+}
+
+func newMultiConnListener() *multiConnListener {
+	return &multiConnListener{
+		connCh: make(chan net.Conn, 100),
+		closed: make(chan struct{}),
+	}
+}
+
+func (l *multiConnListener) Accept() (net.Conn, error) {
+	select {
+	case conn := <-l.connCh:
+		if conn == nil {
+			return nil, fmt.Errorf("listener closed")
+		}
+		return conn, nil
+	case <-l.closed:
+		return nil, fmt.Errorf("listener closed")
+	}
+}
+
+func (l *multiConnListener) Close() error {
+	l.once.Do(func() {
+		close(l.closed)
+		// Drain any remaining connections
+		go func() {
+			for conn := range l.connCh {
+				if conn != nil {
+					conn.Close()
+				}
+			}
+		}()
+		close(l.connCh)
+	})
+	return nil
+}
+
+func (l *multiConnListener) Addr() net.Addr {
+	return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}
+}
+
+func (l *multiConnListener) addConn(conn net.Conn) error {
+	select {
+	case l.connCh <- conn:
+		return nil
+	case <-l.closed:
+		conn.Close()
+		return fmt.Errorf("listener is closed")
+	default:
+		conn.Close()
+		return fmt.Errorf("connection queue full")
+	}
+}

tunnel/tunnel.go

@@ -0,0 +1,194 @@
+package tunnel
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"net/url"
+	"sync"
+
+	"github.com/gorilla/websocket"
+	log "github.com/sirupsen/logrus"
+	"reichard.io/conduit/pkg/maps"
+	"reichard.io/conduit/types"
+)
+
+type ConnBuilder func() (conn io.ReadWriteCloser, err error)
+
+func NewServerTunnel(name string, wsConn *websocket.Conn) *Tunnel {
+	return &Tunnel{
+		name:    name,
+		streams: maps.New[string, io.ReadWriteCloser](),
+		wsConn:  wsConn,
+	}
+}
+
+func NewClientTunnel(name, target string, serverURL *url.URL, wsConn *websocket.Conn) (*Tunnel, error) {
+	// Get Target URL
+	targetURL, err := url.Parse(target)
+	if err != nil {
+		return nil, err
+	}
+
+	// Derive Conduit URL
+	conduitURL := *serverURL
+	conduitURL.Host = name + "." + conduitURL.Host
+
+	// Get Connection Builder
+	var connBuilder ConnBuilder
+	switch targetURL.Scheme {
+	case "http", "https":
+		log.Infof("creating HTTP tunnel: %s -> %s", conduitURL.String(), target)
+		connBuilder, err = HTTPConnectionBuilder(targetURL)
+		if err != nil {
+			return nil, err
+		}
+	default:
+		log.Infof("creating TCP tunnel: %s -> %s", conduitURL.String(), target)
+		connBuilder = func() (conn io.ReadWriteCloser, err error) {
+			return net.Dial("tcp", target)
+		}
+	}
+
+	return &Tunnel{
+		name:        name,
+		wsConn:      wsConn,
+		streams:     maps.New[string, io.ReadWriteCloser](),
+		connBuilder: connBuilder,
+	}, nil
+}
+
+type Tunnel struct {
+	name        string
+	wsConn      *websocket.Conn
+	streams     *maps.Map[string, io.ReadWriteCloser]
+	connBuilder ConnBuilder
+
+	mu sync.Mutex
+}
+
+func (t *Tunnel) Start() {
+	for {
+		var msg types.Message
+		err := t.wsConn.ReadJSON(&msg)
+		if err != nil {
+			return
+		}
+
+		// Validate Stream
+		if msg.StreamID == "" {
+			log.Warnf("tunnel %s missing streamID", t.name)
+			continue
+		}
+
+		// Ensure Stream
+		if err := t.initStreamConnection(msg.StreamID); err != nil {
+			log.WithError(err).Errorf("failed to initialize stream %s connection", t.name)
+			continue
+		}
+
+		// Handle Messages
+		switch msg.Type {
+		case types.MessageTypeClose:
+			_ = t.CloseStream(msg.StreamID)
+		case types.MessageTypeData:
+			_ = t.WriteStream(msg.StreamID, msg.Data)
+		}
+	}
+}
+
+func (t *Tunnel) initStreamConnection(streamID string) error {
+	if t.connBuilder == nil {
+		return nil
+	}
+
+	if _, found := t.streams.Get(streamID); found {
+		return nil
+	}
+
+	conn, err := t.connBuilder()
+	if err != nil {
+		return err
+	}
+
+	if err := t.AddStream(streamID, conn); err != nil {
+		return err
+	}
+
+	go t.StartStream(streamID, "")
+	return nil
+}
+
+func (t *Tunnel) AddStream(streamID string, conn io.ReadWriteCloser) error {
+	if t.streams.HasKey(streamID) {
+		return fmt.Errorf("stream %s already exists", streamID)
+	}
+	t.streams.Set(streamID, conn)
+	return nil
+}
+
+func (t *Tunnel) StartStream(streamID string, sourceAddr string) error {
+	// Get Stream
+	conn, found := t.streams.Get(streamID)
+	if !found {
+		return fmt.Errorf("stream %s does not exist", streamID)
+	}
+
+	// Close Stream
+	defer func() {
+		_ = t.sendWS(&types.Message{
+			Type:       types.MessageTypeClose,
+			StreamID:   streamID,
+			SourceAddr: sourceAddr,
+		})
+
+		t.CloseStream(streamID)
+	}()
+
+	// Start Stream
+	buffer := make([]byte, 4096)
+	for {
+		n, err := conn.Read(buffer)
+		if err != nil {
+			return err
+		}
+
+		if err := t.sendWS(&types.Message{
+			Type:       types.MessageTypeData,
+			StreamID:   streamID,
+			Data:       buffer[:n],
+			SourceAddr: sourceAddr,
+		}); err != nil {
+			return err
+		}
+	}
+}
+
+func (t *Tunnel) WriteStream(streamID string, data []byte) error {
+	// Get Stream
+	conn, found := t.streams.Get(streamID)
+	if !found {
+		return fmt.Errorf("stream %s does not exist", streamID)
+	}
+
+	_, err := conn.Write(data)
+	return err
+}
+
+func (t *Tunnel) CloseStream(streamID string) error {
+	if conn, ok := t.streams.Get(streamID); ok {
+		t.streams.Delete(streamID)
+		return conn.Close()
+	}
+	return nil
+}
+
+func (t *Tunnel) Source() string {
+	return t.wsConn.RemoteAddr().String()
+}
+
+func (t *Tunnel) sendWS(msg *types.Message) error {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return t.wsConn.WriteJSON(msg)
+}

types/message.go

@@ -8,7 +8,8 @@ const (
 )
 
 type Message struct {
-	Type     MessageType `json:"type"`
+	Type       MessageType `json:"type"`
-	StreamID string      `json:"stream_id"`
+	StreamID   string      `json:"stream_id"`
-	Data     []byte      `json:"data,omitempty"`
+	SourceAddr string      `json:"source_addr"`
+	Data       []byte      `json:"data,omitempty"`
 }