feat: Reconnecting mmar clients

This change adds the ability for the mmar client to attempt to reconnect
if there is ever a disconnection of the tunnel. Which allows for the
persistance of the subdomain as long as the mmar client is running.
Once the mmar client is shutdown and then restarted, a new subdomain
will be issued.

Author: yusuf-musleh

constants/main.go

@@ -29,6 +29,7 @@ const (
 	ID_LENGTH                       = 6
 
 	MAX_TUNNELS_PER_IP          = 5
+	TUNNEL_RECONNECT_TIMEOUT    = 3
 	GRACEFUL_SHUTDOWN_TIMEOUT   = 3
 	TUNNEL_CREATE_TIMEOUT       = 3
 	REQ_BODY_READ_CHUNK_TIMEOUT = 3

internal/client/main.go

@@ -32,6 +32,7 @@ type MmarClient struct {
 	// Tunnel to Server
 	protocol.Tunnel
 	ConfigOptions
+	subdomain string
 }
 
 func (mc *MmarClient) localizeRequest(request *http.Request) {
@@ -107,6 +108,28 @@ func (mc *MmarClient) handleRequestMessage(tunnelMsg protocol.TunnelMessage) {
 	}
 }
 
+// Keep attempting to reconnect the existing tunnel until successful
+func (mc *MmarClient) reconnectTunnel(ctx context.Context) {
+	for {
+		// If context is cancelled, do not reconnect
+		if errors.Is(ctx.Err(), context.Canceled) {
+			return
+		}
+		logger.Log(constants.DEFAULT_COLOR, "Attempting to reconnect...")
+		conn, err := net.DialTimeout(
+			"tcp",
+			fmt.Sprintf("%s:%s", mc.ConfigOptions.TunnelHost, mc.ConfigOptions.TunnelTcpPort),
+			constants.TUNNEL_CREATE_TIMEOUT*time.Second,
+		)
+		if err != nil {
+			time.Sleep(constants.TUNNEL_RECONNECT_TIMEOUT * time.Second)
+			continue
+		}
+		mc.Tunnel.Conn = conn
+		break
+	}
+}
+
 func (mc *MmarClient) ProcessTunnelMessages(ctx context.Context) {
 	for {
 		select {
@@ -115,21 +138,38 @@ func (mc *MmarClient) ProcessTunnelMessages(ctx context.Context) {
 		default:
 			tunnelMsg, err := mc.ReceiveMessage()
 			if err != nil {
-				if errors.Is(err, io.EOF) {
-					logger.Log(constants.DEFAULT_COLOR, "Tunnel connection closed from Server. Exiting...")
-					os.Exit(0)
-				} else if errors.Is(err, net.ErrClosed) {
-					logger.Log(constants.DEFAULT_COLOR, "Tunnel connection disconnected from Server. Exiting...")
-					os.Exit(0)
+				// If the context was cancelled just return
+				if errors.Is(ctx.Err(), context.Canceled) {
+					return
 				} else if errors.Is(err, os.ErrDeadlineExceeded) {
 					continue
 				}
-				log.Fatalf("Failed to receive message from server tunnel: %v", err)
+
+				logger.Log(constants.DEFAULT_COLOR, "Tunnel connection disconnected.")
+
+				// Keep trying to reconnect
+				mc.reconnectTunnel(ctx)
+
+				continue
 			}
 
 			switch tunnelMsg.MsgType {
 			case protocol.CLIENT_CONNECT:
-				logger.LogTunnelCreated(string(tunnelMsg.MsgData), mc.TunnelHost, mc.TunnelHttpPort, mc.LocalPort)
+				tunnelSubdomain := string(tunnelMsg.MsgData)
+				// If there is an existing subdomain, that means we are reconnecting with an
+				// existing mmar client, try to reclaim the same subdomain
+				if mc.subdomain != "" {
+					reconnectMsg := protocol.TunnelMessage{MsgType: protocol.CLIENT_RECLAIM_SUBDOMAIN, MsgData: []byte(tunnelSubdomain + ":" + mc.subdomain)}
+					mc.subdomain = ""
+					if err := mc.SendMessage(reconnectMsg); err != nil {
+						logger.Log(constants.DEFAULT_COLOR, "Tunnel failed to reconnect. Exiting...")
+						os.Exit(0)
+					}
+					continue
+				} else {
+					mc.subdomain = tunnelSubdomain
+				}
+				logger.LogTunnelCreated(tunnelSubdomain, mc.TunnelHost, mc.TunnelHttpPort, mc.LocalPort)
 			case protocol.CLIENT_TUNNEL_LIMIT:
 				limit := logger.ColorLogStr(
 					constants.RED,
@@ -177,6 +217,7 @@ func Run(config ConfigOptions) {
 	mmarClient := MmarClient{
 		protocol.Tunnel{Conn: conn},
 		config,
+		"",
 	}
 
 	// Create context to cancel running gouroutines when shutting down

internal/protocol/main.go

@@ -20,6 +20,7 @@ const (
 	REQUEST = uint8(iota + 1)
 	RESPONSE
 	CLIENT_CONNECT
+	CLIENT_RECLAIM_SUBDOMAIN
 	CLIENT_DISCONNECT
 	CLIENT_TUNNEL_LIMIT
 	LOCALHOST_NOT_RUNNING

internal/server/main.go

@@ -16,6 +16,7 @@ import (
 	"os"
 	"os/signal"
 	"slices"
+	"strings"
 	"sync"
 	"time"
 
@@ -465,6 +466,52 @@ func (ms *MmarServer) processTunnelMessages(ct *ClientTunnel) {
 		case protocol.CLIENT_DISCONNECT:
 			ms.closeClientTunnel(ct)
 			return
+		case protocol.CLIENT_RECLAIM_SUBDOMAIN:
+			newAndExistingIDs := strings.Split(string(tunnelMsg.MsgData), ":")
+			newId := newAndExistingIDs[0]
+			existingId := newAndExistingIDs[1]
+
+			// Check if the subdomain has already been taken
+			_, ok := ms.clients[existingId]
+			if ok {
+				// if so, close the tunnel, so the user can create a new one
+				ms.closeClientTunnel(ct)
+				return
+			}
+
+			ct.Tunnel.Id = existingId
+
+			// Add existing client tunnel to clients
+			ms.clients[existingId] = *ct
+
+			// Remove newId tunnel from clients
+			delete(ms.clients, newId)
+
+			// Update the tunnels for the IP
+			clientIP := utils.ExtractIP(ct.Conn.RemoteAddr().String())
+			newIdIndex := slices.Index(ms.tunnelsPerIP[clientIP], newId)
+			if newIdIndex == -1 {
+				ms.tunnelsPerIP[clientIP] = append(ms.tunnelsPerIP[clientIP], existingId)
+			} else {
+				ms.tunnelsPerIP[clientIP][newIdIndex] = existingId
+			}
+
+			connMessage := protocol.TunnelMessage{MsgType: protocol.CLIENT_CONNECT, MsgData: []byte(existingId)}
+			if err := ct.SendMessage(connMessage); err != nil {
+				logger.Log(constants.DEFAULT_COLOR, fmt.Sprintf("Failed to send unique ID msg to client: %v", err))
+				ms.closeClientTunnel(ct)
+				return
+			}
+
+			logger.Log(
+				constants.DEFAULT_COLOR,
+				fmt.Sprintf(
+					"[%s] Tunnel reclaimed: %s -> %s",
+					newId,
+					ct.Conn.RemoteAddr().String(),
+					existingId,
+				),
+			)
 		}
 	}
 }