refactor: Unify the edgeMarker adding logic

packages/mermaid/src/dagre-wrapper/edgeMarker.spec.ts

@@ -0,0 +1,79 @@
+import type { Mocked } from 'vitest';
+import type { SVG } from '../diagram-api/types.js';
+import { addEdgeMarkers } from './edgeMarker.js';
+
+describe('addEdgeMarker', () => {
+  const svgPath = {
+    attr: vitest.fn(),
+  } as unknown as Mocked<SVG>;
+  const url = 'http://example.com';
+  const id = 'test';
+  const diagramType = 'test';
+
+  beforeEach(() => {
+    svgPath.attr.mockReset();
+  });
+
+  it('should add markers for arrow_cross:arrow_point', () => {
+    const arrowTypeStart = 'arrow_cross';
+    const arrowTypeEnd = 'arrow_point';
+    addEdgeMarkers(svgPath, { arrowTypeStart, arrowTypeEnd }, url, id, diagramType);
+    expect(svgPath.attr).toHaveBeenCalledWith(
+      'marker-start',
+      `url(${url}#${id}_${diagramType}-crossStart)`
+    );
+    expect(svgPath.attr).toHaveBeenCalledWith(
+      'marker-end',
+      `url(${url}#${id}_${diagramType}-pointEnd)`
+    );
+  });
+
+  it('should add markers for aggregation:arrow_point', () => {
+    const arrowTypeStart = 'aggregation';
+    const arrowTypeEnd = 'arrow_point';
+    addEdgeMarkers(svgPath, { arrowTypeStart, arrowTypeEnd }, url, id, diagramType);
+    expect(svgPath.attr).toHaveBeenCalledWith(
+      'marker-start',
+      `url(${url}#${id}_${diagramType}-aggregationStart)`
+    );
+    expect(svgPath.attr).toHaveBeenCalledWith(
+      'marker-end',
+      `url(${url}#${id}_${diagramType}-pointEnd)`
+    );
+  });
+
+  it('should add markers for arrow_point:aggregation', () => {
+    const arrowTypeStart = 'arrow_point';
+    const arrowTypeEnd = 'aggregation';
+    addEdgeMarkers(svgPath, { arrowTypeStart, arrowTypeEnd }, url, id, diagramType);
+    expect(svgPath.attr).toHaveBeenCalledWith(
+      'marker-start',
+      `url(${url}#${id}_${diagramType}-pointStart)`
+    );
+    expect(svgPath.attr).toHaveBeenCalledWith(
+      'marker-end',
+      `url(${url}#${id}_${diagramType}-aggregationEnd)`
+    );
+  });
+
+  it('should add markers for aggregation:composition', () => {
+    const arrowTypeStart = 'aggregation';
+    const arrowTypeEnd = 'composition';
+    addEdgeMarkers(svgPath, { arrowTypeStart, arrowTypeEnd }, url, id, diagramType);
+    expect(svgPath.attr).toHaveBeenCalledWith(
+      'marker-start',
+      `url(${url}#${id}_${diagramType}-aggregationStart)`
+    );
+    expect(svgPath.attr).toHaveBeenCalledWith(
+      'marker-end',
+      `url(${url}#${id}_${diagramType}-compositionEnd)`
+    );
+  });
+
+  it('should not add invalid markers', () => {
+    const arrowTypeStart = 'this is an invalid marker';
+    const arrowTypeEnd = ') url(https://my-malicious-site.example)';
+    addEdgeMarkers(svgPath, { arrowTypeStart, arrowTypeEnd }, url, id, diagramType);
+    expect(svgPath.attr).not.toHaveBeenCalled();
+  });
+});

packages/mermaid/src/dagre-wrapper/edgeMarker.ts

@@ -0,0 +1,57 @@
+import type { SVG } from '../diagram-api/types.js';
+import { log } from '../logger.js';
+import type { EdgeData } from '../types.js';
+/**
+ * Adds SVG markers to a path element based on the arrow types specified in the edge.
+ *
+ * @param svgPath - The SVG path element to add markers to.
+ * @param edge - The edge data object containing the arrow types.
+ * @param url - The URL of the SVG marker definitions.
+ * @param id - The ID prefix for the SVG marker definitions.
+ * @param diagramType - The type of diagram being rendered.
+ */
+export const addEdgeMarkers = (
+  svgPath: SVG,
+  edge: Pick<EdgeData, 'arrowTypeStart' | 'arrowTypeEnd'>,
+  url: string,
+  id: string,
+  diagramType: string
+) => {
+  if (edge.arrowTypeStart) {
+    addEdgeMarker(svgPath, 'start', edge.arrowTypeStart, url, id, diagramType);
+  }
+  if (edge.arrowTypeEnd) {
+    addEdgeMarker(svgPath, 'end', edge.arrowTypeEnd, url, id, diagramType);
+  }
+};
+
+const arrowTypesMap = {
+  arrow_cross: 'cross',
+  arrow_point: 'point',
+  arrow_barb: 'barb',
+  arrow_circle: 'circle',
+  aggregation: 'aggregation',
+  extension: 'extension',
+  composition: 'composition',
+  dependency: 'dependency',
+  lollipop: 'lollipop',
+} as const;
+
+const addEdgeMarker = (
+  svgPath: SVG,
+  position: 'start' | 'end',
+  arrowType: string,
+  url: string,
+  id: string,
+  diagramType: string
+) => {
+  const endMarkerType = arrowTypesMap[arrowType as keyof typeof arrowTypesMap];
+
+  if (!endMarkerType) {
+    log.warn(`Unknown arrow type: ${arrowType}`);
+    return; // unknown arrow type, ignore
+  }
+
+  const suffix = position === 'start' ? 'Start' : 'End';
+  svgPath.attr(`marker-${position}`, `url(${url}#${id}_${diagramType}-${endMarkerType}${suffix})`);
+};

packages/mermaid/src/dagre-wrapper/edges.js

@@ -6,6 +6,7 @@ import { getConfig } from '../diagram-api/diagramAPI.js';
 import utils from '../utils.js';
 import { evaluate } from '../diagrams/common/common.js';
 import { getLineFunctionsWithOffset } from '../utils/lineWithOffset.js';
+import { addEdgeMarkers } from './edgeMarker.js';
 
 let edgeLabels = {};
 let terminalLabels = {};
@@ -506,108 +507,8 @@ export const insertEdge = function (elem, e, edge, clusterDb, diagramType, graph
   log.info('arrowTypeStart', edge.arrowTypeStart);
   log.info('arrowTypeEnd', edge.arrowTypeEnd);
 
-  switch (edge.arrowTypeStart) {
-    case 'arrow_cross':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-crossStart' + ')'
-      );
-      break;
-    case 'arrow_point':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-pointStart' + ')'
-      );
-      break;
-    case 'arrow_barb':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-barbStart' + ')'
-      );
-      break;
-    case 'arrow_circle':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-circleStart' + ')'
-      );
-      break;
-    case 'aggregation':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-aggregationStart' + ')'
-      );
-      break;
-    case 'extension':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-extensionStart' + ')'
-      );
-      break;
-    case 'composition':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-compositionStart' + ')'
-      );
-      break;
-    case 'dependency':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-dependencyStart' + ')'
-      );
-      break;
-    case 'lollipop':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-lollipopStart' + ')'
-      );
-      break;
-    default:
-  }
-  switch (edge.arrowTypeEnd) {
-    case 'arrow_cross':
-      svgPath.attr('marker-end', 'url(' + url + '#' + id + '_' + diagramType + '-crossEnd' + ')');
-      break;
-    case 'arrow_point':
-      svgPath.attr('marker-end', 'url(' + url + '#' + id + '_' + diagramType + '-pointEnd' + ')');
-      break;
-    case 'arrow_barb':
-      svgPath.attr('marker-end', 'url(' + url + '#' + id + '_' + diagramType + '-barbEnd' + ')');
-      break;
-    case 'arrow_circle':
-      svgPath.attr('marker-end', 'url(' + url + '#' + id + '_' + diagramType + '-circleEnd' + ')');
-      break;
-    case 'aggregation':
-      svgPath.attr(
-        'marker-end',
-        'url(' + url + '#' + id + '_' + diagramType + '-aggregationEnd' + ')'
-      );
-      break;
-    case 'extension':
-      svgPath.attr(
-        'marker-end',
-        'url(' + url + '#' + id + '_' + diagramType + '-extensionEnd' + ')'
-      );
-      break;
-    case 'composition':
-      svgPath.attr(
-        'marker-end',
-        'url(' + url + '#' + id + '_' + diagramType + '-compositionEnd' + ')'
-      );
-      break;
-    case 'dependency':
-      svgPath.attr(
-        'marker-end',
-        'url(' + url + '#' + id + '_' + diagramType + '-dependencyEnd' + ')'
-      );
-      break;
-    case 'lollipop':
-      svgPath.attr(
-        'marker-end',
-        'url(' + url + '#' + id + '_' + diagramType + '-lollipopEnd' + ')'
-      );
-      break;
-    default:
-  }
+  addEdgeMarkers(svgPath, edge, url, id, diagramType);
+
   let paths = {};
   if (pointsHasChanged) {
     paths.updatedPath = points;

packages/mermaid/src/diagrams/flowchart/elk/flowRenderer-elk.js

@@ -11,6 +11,7 @@ import common from '../../common/common.js';
 import { interpolateToCurve, getStylesFromArray } from '../../../utils.js';
 import ELK from 'elkjs/lib/elk.bundled.js';
 import { getLineFunctionsWithOffset } from '../../../utils/lineWithOffset.js';
+import { addEdgeMarkers } from '../../../dagre-wrapper/edgeMarker.js';
 
 const elk = new ELK();
 
@@ -586,108 +587,7 @@ const addMarkersToEdge = function (svgPath, edgeData, diagramType, arrowMarkerAb
   }
 
   // look in edge data and decide which marker to use
-  switch (edgeData.arrowTypeStart) {
-    case 'arrow_cross':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-crossStart' + ')'
-      );
-      break;
-    case 'arrow_point':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-pointStart' + ')'
-      );
-      break;
-    case 'arrow_barb':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-barbStart' + ')'
-      );
-      break;
-    case 'arrow_circle':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-circleStart' + ')'
-      );
-      break;
-    case 'aggregation':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-aggregationStart' + ')'
-      );
-      break;
-    case 'extension':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-extensionStart' + ')'
-      );
-      break;
-    case 'composition':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-compositionStart' + ')'
-      );
-      break;
-    case 'dependency':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-dependencyStart' + ')'
-      );
-      break;
-    case 'lollipop':
-      svgPath.attr(
-        'marker-start',
-        'url(' + url + '#' + id + '_' + diagramType + '-lollipopStart' + ')'
-      );
-      break;
-    default:
-  }
-  switch (edgeData.arrowTypeEnd) {
-    case 'arrow_cross':
-      svgPath.attr('marker-end', 'url(' + url + '#' + id + '_' + diagramType + '-crossEnd' + ')');
-      break;
-    case 'arrow_point':
-      svgPath.attr('marker-end', 'url(' + url + '#' + id + '_' + diagramType + '-pointEnd' + ')');
-      break;
-    case 'arrow_barb':
-      svgPath.attr('marker-end', 'url(' + url + '#' + id + '_' + diagramType + '-barbEnd' + ')');
-      break;
-    case 'arrow_circle':
-      svgPath.attr('marker-end', 'url(' + url + '#' + id + '_' + diagramType + '-circleEnd' + ')');
-      break;
-    case 'aggregation':
-      svgPath.attr(
-        'marker-end',
-        'url(' + url + '#' + id + '_' + diagramType + '-aggregationEnd' + ')'
-      );
-      break;
-    case 'extension':
-      svgPath.attr(
-        'marker-end',
-        'url(' + url + '#' + id + '_' + diagramType + '-extensionEnd' + ')'
-      );
-      break;
-    case 'composition':
-      svgPath.attr(
-        'marker-end',
-        'url(' + url + '#' + id + '_' + diagramType + '-compositionEnd' + ')'
-      );
-      break;
-    case 'dependency':
-      svgPath.attr(
-        'marker-end',
-        'url(' + url + '#' + id + '_' + diagramType + '-dependencyEnd' + ')'
-      );
-      break;
-    case 'lollipop':
-      svgPath.attr(
-        'marker-end',
-        'url(' + url + '#' + id + '_' + diagramType + '-lollipopEnd' + ')'
-      );
-      break;
-    default:
-  }
+  addEdgeMarkers(svgPath, edgeData, url, id, diagramType);
 };
 
 /**