Bug in Text Visualization Code

[peiyangL]

Overview

The affected code snippet can be found here: https://github.com/pytorch/captum/blob/master/captum/attr/_utils/visualization.py#L829-L836

Within the format_word_importances function, the word variable is not being properly escaped before being incorporated into the HTML template. This oversight could potentially result in incorrect visualization outputs or even present a security vulnerability if untrusted input is used. Below, I provide two examples to demonstrate these issues:

To Reproduce

Case 1: Commented-out User Input

from captum.attr import visualization

words = ['hello', 'world', 'entrar<!--#exec', 'cmd', '"ls', '-->']
scores = [0 for _ in words]

_ = visualization.visualize_text([visualization.VisualizationDataRecord(
    word_attributions=scores,
    raw_input_ids=words,
    attr_score=sum(scores),
    pred_class=0,
    pred_prob=0,
    true_class=0,
    attr_class=0,
    convergence_score=0
)])

The last four words are not displayed correctly.

Case 2: Cross-Site Scripting (XSS) in a Jupyter Environment

from captum.attr import visualization

words = ['hello', 'world', "x<script>alert('xss');</script>x"]
scores = [-1, 0, 1]

_ = visualization.visualize_text([visualization.VisualizationDataRecord(
    word_attributions=scores,
    raw_input_ids=words,
    attr_score=sum(scores),
    pred_class=0,
    pred_prob=0,
    true_class=0,
    attr_class=0,
    convergence_score=0
)])

Solution

the fix for this issue is straightforward—by implementing proper escaping for the word variable. Just modify the line 829.

word = format_special_tokens(word)

->

import html
word = html.escape(word)
# word = html.escape(format_special_tokens(word)) # if you want to keep the format_special_tokens function