fix: Properly handle invalid resp from dest server

This fix properly handles case when destination server returns an invalid
or malformed http response, preventing mmar client from crashing and informing
the user to check the destination server's logs.

Author: yusuf-musleh

constants/main.go

@@ -23,7 +23,7 @@ const (
 	CLIENT_TCP_PORT_HELP   = "Define port of mmar TCP server for client to connect to, creating a tunnel."
 	TUNNEL_HOST_HELP       = "Define host domain of mmar server for client to connect to."
 
-	TUNNEL_MESSAGE_PROTOCOL_VERSION = 2
+	TUNNEL_MESSAGE_PROTOCOL_VERSION = 3
 	TUNNEL_MESSAGE_DATA_DELIMITER   = '\n'
 	ID_CHARSET                      = "abcdefghijklmnopqrstuvwxyz0123456789"
 	ID_LENGTH                       = 6

internal/client/main.go

@@ -93,7 +93,11 @@ func (mc *MmarClient) handleRequestMessage(tunnelMsg protocol.TunnelMessage) {
 			return
 		}
 
-		log.Fatalf("Failed to forward: %v", fwdErr)
+		invalidRespFromDestMsg := protocol.TunnelMessage{MsgType: protocol.INVALID_RESP_FROM_DEST}
+		if err := mc.SendMessage(invalidRespFromDestMsg); err != nil {
+			log.Fatal(err)
+		}
+		return
 	}
 
 	logger.LogHTTP(req, resp.StatusCode, resp.ContentLength, false, true)
@@ -118,7 +122,7 @@ func (mc *MmarClient) reconnectTunnel(ctx context.Context) {
 		logger.Log(constants.DEFAULT_COLOR, "Attempting to reconnect...")
 		conn, err := net.DialTimeout(
 			"tcp",
-			fmt.Sprintf("%s:%s", mc.ConfigOptions.TunnelHost, mc.ConfigOptions.TunnelTcpPort),
+			net.JoinHostPort(mc.ConfigOptions.TunnelHost, mc.ConfigOptions.TunnelTcpPort),
 			constants.TUNNEL_CREATE_TIMEOUT*time.Second,
 		)
 		if err != nil {
@@ -229,7 +233,7 @@ func Run(config ConfigOptions) {
 
 	conn, err := net.DialTimeout(
 		"tcp",
-		fmt.Sprintf("%s:%s", config.TunnelHost, config.TunnelTcpPort),
+		net.JoinHostPort(config.TunnelHost, config.TunnelTcpPort),
 		constants.TUNNEL_CREATE_TIMEOUT*time.Second,
 	)
 	if err != nil {

internal/protocol/main.go

@@ -28,6 +28,7 @@ const (
 	HEARTBEAT_FROM_CLIENT
 	HEARTBEAT_FROM_SERVER
 	HEARTBEAT_ACK
+	INVALID_RESP_FROM_DEST
 )
 
 var INVALID_MESSAGE_PROTOCOL_VERSION = errors.New("Invalid Message Protocol Version")
@@ -36,7 +37,7 @@ var INVALID_MESSAGE_TYPE = errors.New("Invalid Tunnel Message Type")
 func isValidTunnelMessageType(mt uint8) (uint8, error) {
 	// Iterate through all the message type, from first to last, checking
 	// if the provided message type matches one of them
-	for msgType := REQUEST; msgType <= HEARTBEAT_ACK; msgType++ {
+	for msgType := REQUEST; msgType <= INVALID_RESP_FROM_DEST; msgType++ {
 		if mt == msgType {
 			return msgType, nil
 		}
@@ -48,9 +49,10 @@ func isValidTunnelMessageType(mt uint8) (uint8, error) {
 func TunnelErrState(errState uint8) string {
 	// TODO: Have nicer/more elaborative error messages/pages
 	errStates := map[uint8]string{
-		CLIENT_DISCONNECT:     constants.CLIENT_DISCONNECT_ERR_TEXT,
-		LOCALHOST_NOT_RUNNING: constants.LOCALHOST_NOT_RUNNING_ERR_TEXT,
-		DEST_REQUEST_TIMEDOUT: constants.DEST_REQUEST_TIMEDOUT_ERR_TEXT,
+		CLIENT_DISCONNECT:      constants.CLIENT_DISCONNECT_ERR_TEXT,
+		LOCALHOST_NOT_RUNNING:  constants.LOCALHOST_NOT_RUNNING_ERR_TEXT,
+		DEST_REQUEST_TIMEDOUT:  constants.DEST_REQUEST_TIMEDOUT_ERR_TEXT,
+		INVALID_RESP_FROM_DEST: constants.READ_RESP_BODY_ERR_TEXT,
 	}
 	fallbackErr := "An error occured while attempting to tunnel."
 

internal/server/main.go

@@ -461,29 +461,13 @@ func (ms *MmarServer) processTunnelMessages(ct *ClientTunnel) {
 		case protocol.LOCALHOST_NOT_RUNNING:
 			// Create a response for Tunnel connected but localhost not running
 			errState := protocol.TunnelErrState(protocol.LOCALHOST_NOT_RUNNING)
-			resp := http.Response{
-				Status:     "200 OK",
-				StatusCode: http.StatusOK,
-				Body:       io.NopCloser(bytes.NewBufferString(errState)),
-			}
-
-			// Writing response to buffer to tunnel it back
-			var responseBuff bytes.Buffer
-			resp.Write(&responseBuff)
+			responseBuff := createSerializedServerResp("200 OK", http.StatusOK, errState)
 			notRunningMsg := protocol.TunnelMessage{MsgType: protocol.RESPONSE, MsgData: responseBuff.Bytes()}
 			ct.outgoingChannel <- notRunningMsg
 		case protocol.DEST_REQUEST_TIMEDOUT:
 			// Create a response for Tunnel connected but localhost took too long to respond
 			errState := protocol.TunnelErrState(protocol.DEST_REQUEST_TIMEDOUT)
-			resp := http.Response{
-				Status:     "200 OK",
-				StatusCode: http.StatusOK,
-				Body:       io.NopCloser(bytes.NewBufferString(errState)),
-			}
-
-			// Writing response to buffer to tunnel it back
-			var responseBuff bytes.Buffer
-			resp.Write(&responseBuff)
+			responseBuff := createSerializedServerResp("200 OK", http.StatusOK, errState)
 			destTimedoutMsg := protocol.TunnelMessage{MsgType: protocol.RESPONSE, MsgData: responseBuff.Bytes()}
 			ct.outgoingChannel <- destTimedoutMsg
 		case protocol.CLIENT_DISCONNECT:
@@ -545,6 +529,12 @@ func (ms *MmarServer) processTunnelMessages(ct *ClientTunnel) {
 					existingId,
 				),
 			)
+		case protocol.INVALID_RESP_FROM_DEST:
+			// Create a response for receiving invalid response from destination server
+			errState := protocol.TunnelErrState(protocol.INVALID_RESP_FROM_DEST)
+			responseBuff := createSerializedServerResp("500 Internal Server Error", http.StatusInternalServerError, errState)
+			invalidRespFromDestMsg := protocol.TunnelMessage{MsgType: protocol.RESPONSE, MsgData: responseBuff.Bytes()}
+			ct.outgoingChannel <- invalidRespFromDestMsg
 		}
 	}
 }

internal/server/utils.go

@@ -21,7 +21,7 @@ var MAX_REQ_BODY_SIZE_ERR error = errors.New(constants.MAX_REQ_BODY_SIZE_ERR_TEX
 var FAILED_TO_FORWARD_TO_MMAR_CLIENT_ERR error = errors.New(constants.FAILED_TO_FORWARD_TO_MMAR_CLIENT_ERR_TEXT)
 var FAILED_TO_READ_RESP_FROM_MMAR_CLIENT_ERR error = errors.New(constants.FAILED_TO_READ_RESP_FROM_MMAR_CLIENT_ERR_TEXT)
 
-func responseWith(respText string, w http.ResponseWriter, statusCode int) {
+func respondWith(respText string, w http.ResponseWriter, statusCode int) {
 	w.Header().Set("Content-Length", strconv.Itoa(len(respText)))
 	w.Header().Set("Connection", "close")
 	w.WriteHeader(statusCode)
@@ -34,15 +34,15 @@ func handleCancel(cause error, w http.ResponseWriter) {
 		// Cancelled, do nothing
 		return
 	case READ_BODY_CHUNK_TIMEOUT_ERR:
-		responseWith(cause.Error(), w, http.StatusRequestTimeout)
+		respondWith(cause.Error(), w, http.StatusRequestTimeout)
 	case READ_BODY_CHUNK_ERR, CLIENT_DISCONNECTED_ERR:
-		responseWith(cause.Error(), w, http.StatusBadRequest)
+		respondWith(cause.Error(), w, http.StatusBadRequest)
 	case READ_RESP_BODY_ERR:
-		responseWith(cause.Error(), w, http.StatusInternalServerError)
+		respondWith(cause.Error(), w, http.StatusInternalServerError)
 	case MAX_REQ_BODY_SIZE_ERR:
-		responseWith(cause.Error(), w, http.StatusRequestEntityTooLarge)
+		respondWith(cause.Error(), w, http.StatusRequestEntityTooLarge)
 	case FAILED_TO_FORWARD_TO_MMAR_CLIENT_ERR, FAILED_TO_READ_RESP_FROM_MMAR_CLIENT_ERR:
-		responseWith(cause.Error(), w, http.StatusServiceUnavailable)
+		respondWith(cause.Error(), w, http.StatusServiceUnavailable)
 	}
 }
 
@@ -119,3 +119,18 @@ func serializeRequest(ctx context.Context, r *http.Request, cancel context.Cance
 	// Send serialized request through channel
 	serializedRequestChannel <- requestBuff.Bytes()
 }
+
+// Create HTTP response sent from mmar server to the end-user client
+func createSerializedServerResp(status string, statusCode int, body string) bytes.Buffer {
+	resp := http.Response{
+		Status:     status,
+		StatusCode: statusCode,
+		Body:       io.NopCloser(bytes.NewBufferString(body)),
+	}
+
+	// Writing response to buffer to tunnel it back
+	var responseBuff bytes.Buffer
+	resp.Write(&responseBuff)
+
+	return responseBuff
+}

simulations/devserver/main.go

@@ -170,19 +170,27 @@ func handleRedirect(w http.ResponseWriter, r *http.Request) {
 
 // Request handler that returns an invalid HTTP response
 func handleBadResp(w http.ResponseWriter, r *http.Request) {
-	// Return a response with Content-Length headers that do not match the actual data
-	respBody, err := json.Marshal(map[string]interface{}{
-		"data": "some data",
-	})
+	// Get the underlying connection object
+	// Assert that w supports Hijacking
+	hijacker, ok := w.(http.Hijacker)
+	if !ok {
+		http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
+		return
+	}
 
+	// Hijack the connection
+	conn, buf, err := hijacker.Hijack()
 	if err != nil {
-		log.Fatalf("Failed to marshal response for GET: %v", err)
+		http.Error(w, "Hijacking failed", http.StatusInternalServerError)
+		return
 	}
+	defer conn.Close()
 
-	w.Header().Set("Content-Type", "application/json")
-	w.Header().Set("Content-Length", "123") // Content length much larger than actual content
-	w.WriteHeader(http.StatusOK)
-	w.Write(respBody)
+	// Send back an invalid HTTP response
+	buf.WriteString("some random string\r\n" +
+		"\r\n" +
+		"that is not a valid http resp")
+	buf.Flush()
 }
 
 // Request handler that takes a long time before returning response