clean

cmd/link.go

@@ -9,14 +9,9 @@ import (
 	"github.com/gorilla/websocket"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
+	"reichard.io/conduit/types"
 )
 
-type TunnelMessage struct {
-	Type     string `json:"type"`
-	StreamID string `json:"stream_id"`
-	Data     []byte `json:"data,omitempty"`
-}
-
 var serverAddr string
 
 var linkCmd = &cobra.Command{
@@ -69,7 +64,7 @@ func NewTunnel(tunnelName, tunnelTarget, serverAddress string) (*Tunnel, error)
 	}
 
 	// Connect Server WS
-	wsURL := fmt.Sprintf("%s://%s/_conduit/tunnel?vhost=%s", wsScheme, serverURL.Host, tunnelName)
+	wsURL := fmt.Sprintf("%s://%s/_conduit/tunnel?tunnelName=%s", wsScheme, serverURL.Host, tunnelName)
 	serverConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
 	if err != nil {
 		return nil, fmt.Errorf("failed to connect: %v", err)
@@ -100,7 +95,7 @@ func (t *Tunnel) Start() error {
 	// Handle Messages
 	for {
 		// Read Message
-		var msg TunnelMessage
+		var msg types.Message
 		err := t.serverConn.ReadJSON(&msg)
 		if err != nil {
 			log.Errorf("Error reading from tunnel: %v", err)
@@ -108,7 +103,7 @@ func (t *Tunnel) Start() error {
 		}
 
 		switch msg.Type {
-		case "data":
+		case types.MessageTypeData:
 			localConn, err := t.getLocalConn(msg.StreamID)
 			if err != nil {
 				log.Errorf("Failed to get local connection: %v", err)
@@ -124,7 +119,7 @@ func (t *Tunnel) Start() error {
 				t.mu.Unlock()
 			}
 
-		case "close":
+		case types.MessageTypeClose:
 			t.mu.Lock()
 			if localConn, exists := t.localConns[msg.StreamID]; exists {
 				localConn.Close()
@@ -176,8 +171,8 @@ func (t *Tunnel) startResponseRelay(streamID string, lConn net.Conn) {
 			break
 		}
 
-		response := TunnelMessage{
-			Type:     "data",
+		response := types.Message{
+			Type:     types.MessageTypeData,
 			StreamID: streamID,
 			Data:     buffer[:n],
 		}