client/client.go

@@ -2,15 +2,17 @@ package client
 
 import (
 	"fmt"
+	"net"
 	"net/url"
+	"sync"
 
 	"github.com/gorilla/websocket"
 	log "github.com/sirupsen/logrus"
 	"reichard.io/conduit/config"
-	"reichard.io/conduit/tunnel"
+	"reichard.io/conduit/types"
 )
 
-func NewTunnel(cfg *config.ClientConfig) (*tunnel.Tunnel, error) {
+func NewTunnel(cfg *config.ClientConfig) (*Tunnel, error) {
 	// Parse Server URL
 	serverURL, err := url.Parse(cfg.ServerAddress)
 	if err != nil {
@@ -41,5 +43,117 @@ func NewTunnel(cfg *config.ClientConfig) (*tunnel.Tunnel, error) {
 		return nil, fmt.Errorf("failed to connect: %v", err)
 	}
 
-	return tunnel.NewClientTunnel(cfg.TunnelName, cfg.TunnelTarget, serverConn), nil
+	return &Tunnel{
+		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

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

server/reconstructed_conn.go

@@ -1,30 +0,0 @@
-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

@@ -17,7 +17,7 @@ import (
 	"github.com/gorilla/websocket"
 	log "github.com/sirupsen/logrus"
 	"reichard.io/conduit/config"
-	"reichard.io/conduit/tunnel"
+	"reichard.io/conduit/types"
 )
 
 type InfoResponse struct {
@@ -30,13 +30,19 @@ 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]*tunnel.Tunnel
+	tunnels  map[string]*TunnelConnection
 }
 
 func NewServer(cfg *config.ServerConfig) (*Server, error) {
@@ -50,7 +56,7 @@ func NewServer(cfg *config.ServerConfig) (*Server, error) {
 	return &Server{
 		cfg:     cfg,
 		host:    serverURL.Host,
-		tunnels: make(map[string]*tunnel.Tunnel),
+		tunnels: make(map[string]*TunnelConnection),
 		upgrader: websocket.Upgrader{
 			CheckOrigin: func(r *http.Request) bool {
 				return true
@@ -88,7 +94,7 @@ func (s *Server) getInfo(w http.ResponseWriter, _ *http.Request) {
 	for t, c := range s.tunnels {
 		allTunnels = append(allTunnels, TunnelInfo{
 			Name:   t,
-			Target: c.Source(),
+			Target: c.RemoteAddr().String(),
 		})
 	}
 	s.mu.RUnlock()
@@ -108,6 +114,63 @@ func (s *Server) getInfo(w http.ResponseWriter, _ *http.Request) {
 	_, _ = 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()
 
@@ -120,7 +183,7 @@ func (s *Server) handleRawConnection(conn net.Conn) {
 	bufReader := bufio.NewReader(teeReader)
 
 	// Create HTTP Request & Writer
-	w := &rawHTTPResponseWriter{conn: conn}
+	w := &connResponseWriter{conn: conn}
 	r, err := http.ReadRequest(bufReader)
 	if err != nil {
 		w.WriteHeader(http.StatusBadRequest)
@@ -151,25 +214,15 @@ func (s *Server) handleRawConnection(conn net.Conn) {
 
 	// Handle Tunnels
 	s.mu.RLock()
-	conduitTunnel, exists := s.tunnels[subdomain]
+	tunnelConn, exists := s.tunnels[subdomain]
 	s.mu.RUnlock()
-	if !exists {
-		w.WriteHeader(http.StatusNotFound)
-		_, _ = fmt.Fprintf(w, "unknown tunnel: %s", subdomain)
-		return
-	}
+	if exists {
+		log.Infof("relaying %s to tunnel", subdomain)
 
-	// 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
+		// Reconstruct Data & Proxy Connection
+		allReader := io.MultiReader(&capturedData, r.Body)
+		s.proxyRawConnection(conn, tunnelConn, allReader)
 	}
-
-	log.Infof("relaying %s to tunnel", subdomain)
-	_ = conduitTunnel.StartStream(streamID)
 }
 
 func (s *Server) handleAsHTTP(w http.ResponseWriter, r *http.Request) {
@@ -192,6 +245,40 @@ 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")
@@ -215,20 +302,26 @@ func (s *Server) createTunnel(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Create Tunnel
-	conduitTunnel := tunnel.NewServerTunnel(tunnelName, wsConn)
+	// Create & Cache TunnelConnection
+	tunnel := &TunnelConnection{
+		Conn:    wsConn,
+		name:    tunnelName,
+		streams: make(map[string]chan []byte),
+	}
 	s.mu.Lock()
-	s.tunnels[tunnelName] = conduitTunnel
+	s.tunnels[tunnelName] = tunnel
 	s.mu.Unlock()
 	log.Infof("tunnel established: %s", tunnelName)
 
-	// Start Tunnel - This is blocking
-	conduitTunnel.Start()
+	// Keep connection alive and handle cleanup
+	defer func() {
+		s.mu.Lock()
+		delete(s.tunnels, tunnelName)
+		s.mu.Unlock()
+		_ = wsConn.Close()
+		log.Infof("tunnel closed: %s", tunnelName)
+	}()
 
-	// Cleanup Tunnel
-	s.mu.Lock()
-	delete(s.tunnels, tunnelName)
-	s.mu.Unlock()
-	_ = wsConn.Close()
-	log.Infof("tunnel closed: %s", tunnelName)
+	// Handle tunnel messages
+	s.handleTunnelMessages(tunnel)
 }

server/writer.go

@@ -7,25 +7,25 @@ import (
 	"net/http"
 )
 
-var _ http.ResponseWriter = (*rawHTTPResponseWriter)(nil)
+var _ http.ResponseWriter = (*connResponseWriter)(nil)
 
-type rawHTTPResponseWriter struct {
+type connResponseWriter struct {
 	conn   net.Conn
 	header http.Header
 }
 
-func (f *rawHTTPResponseWriter) Header() http.Header {
+func (f *connResponseWriter) Header() http.Header {
 	if f.header == nil {
 		f.header = make(http.Header)
 	}
 	return f.header
 }
 
-func (f *rawHTTPResponseWriter) Write(data []byte) (int, error) {
+func (f *connResponseWriter) Write(data []byte) (int, error) {
 	return f.conn.Write(data)
 }
 
-func (f *rawHTTPResponseWriter) WriteHeader(statusCode int) {
+func (f *connResponseWriter) 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 *rawHTTPResponseWriter) WriteHeader(statusCode int) {
 	_, _ = f.conn.Write([]byte("\r\n"))
 }
 
-func (f *rawHTTPResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+func (f *connResponseWriter) 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

tunnel/tunnel.go

@@ -1,171 +0,0 @@
-package tunnel
-
-import (
-	"fmt"
-	"io"
-	"net"
-	"sync"
-
-	"github.com/gorilla/websocket"
-	log "github.com/sirupsen/logrus"
-	"reichard.io/conduit/types"
-)
-
-func NewServerTunnel(name string, wsConn *websocket.Conn) *Tunnel {
-	return &Tunnel{
-		name:    name,
-		wsConn:  wsConn,
-		streams: make(map[string]io.ReadWriteCloser),
-	}
-}
-
-func NewClientTunnel(name, target string, wsConn *websocket.Conn) *Tunnel {
-	return &Tunnel{
-		name:    name,
-		wsConn:  wsConn,
-		streams: make(map[string]io.ReadWriteCloser),
-		connBuilder: func() (io.ReadWriteCloser, error) {
-			return net.Dial("tcp", target)
-		},
-	}
-}
-
-type Tunnel struct {
-	name        string
-	wsConn      *websocket.Conn
-	streams     map[string]io.ReadWriteCloser
-	connBuilder func() (io.ReadWriteCloser, error)
-
-	wsMu, streamsMu 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 _, found := t.getStream(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 {
-	t.streamsMu.Lock()
-	defer t.streamsMu.Unlock()
-
-	if _, found := t.streams[streamID]; found {
-		return fmt.Errorf("stream %s already exists", streamID)
-	}
-	t.streams[streamID] = conn
-	return nil
-}
-
-func (t *Tunnel) StartStream(streamID string) error {
-	// Get Stream
-	conn, found := t.getStream(streamID)
-	if !found {
-		return fmt.Errorf("stream %s does not exist", streamID)
-	}
-
-	// Start Stream
-	defer t.CloseStream(streamID)
-	buffer := make([]byte, 4096)
-	for {
-		n, err := conn.Read(buffer)
-		if err != nil {
-			return err
-		}
-
-		if err := t.sendWS(&types.Message{
-			Type:     types.MessageTypeData,
-			Data:     buffer[:n],
-			StreamID: streamID,
-		}); err != nil {
-			return err
-		}
-	}
-}
-
-func (t *Tunnel) WriteStream(streamID string, data []byte) error {
-	// Get Stream
-	conn, found := t.getStream(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 {
-	_ = t.sendWS(&types.Message{
-		Type:     types.MessageTypeClose,
-		StreamID: streamID,
-	})
-
-	t.streamsMu.Lock()
-	defer t.streamsMu.Unlock()
-	if conn, ok := t.streams[streamID]; ok {
-		delete(t.streams, 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.wsMu.Lock()
-	defer t.wsMu.Unlock()
-	return t.wsConn.WriteJSON(msg)
-}
-
-func (t *Tunnel) getStream(streamID string) (io.ReadWriteCloser, bool) {
-	t.streamsMu.Lock()
-	defer t.streamsMu.Unlock()
-
-	if conn, ok := t.streams[streamID]; ok {
-		return conn, true
-	}
-	return nil, false
-}