fix: Properly handle invalid resp from dest server

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