Merge pull request #14 from bouk/add-html-escaping

Add support for HTML escaped string outputting

Author: benbjohnson

README.md

@@ -26,7 +26,9 @@ An ego template is made up of several types of blocks:
 
 * **Code Block** - These blocks execute raw Go code: `<% var foo = "bar" %>`
 
-* **Print Block** - These blocks print a Go expression: `<%= myVar %>`
+* **Print Block** - These blocks print a Go expression. They use `html.EscapeString` to escape it before outputting: `<%= myVar %>`
+
+* **Raw Print Block** - These blocks print a Go expression raw into the HTML: `<%== "<script>" %>`
 
 * **Header Block** - These blocks allow you to import packages: `<%% import "encoding/json" %%>`
 

scanner.go

@@ -63,7 +63,18 @@ func (s *Scanner) scanCodeBlock() (Block, error) {
 	case '%':
 		return s.scanHeaderBlock()
 	case '=':
-		return s.scanPrintBlock()
+		ch, err := s.read()
+		if err == io.EOF {
+			return nil, io.ErrUnexpectedEOF
+		} else if err != nil {
+			return nil, err
+		}
+		if ch == '=' {
+			return s.scanRawPrintBlock()
+		} else {
+			s.unread()
+			return s.scanPrintBlock()
+		}
 	}
 
 	// Otherwise read the contents of the code block.
@@ -98,6 +109,16 @@ func (s *Scanner) scanHeaderBlock() (Block, error) {
 	return b, nil
 }
 
+func (s *Scanner) scanRawPrintBlock() (Block, error) {
+	b := &RawPrintBlock{Pos: s.pos}
+	content, err := s.scanContent()
+	if err != nil {
+		return nil, err
+	}
+	b.Content = content
+	return b, nil
+}
+
 func (s *Scanner) scanPrintBlock() (Block, error) {
 	b := &PrintBlock{Pos: s.pos}
 	content, err := s.scanContent()

scanner_test.go

@@ -79,6 +79,13 @@ func TestScannerCodeBlockUnexpectedEOF_4(t *testing.T) {
 	assert.Equal(t, err, io.ErrUnexpectedEOF)
 }
 
+// Ensure that a print code block that ends unexpectedly returns an error.
+func TestScannerCodeBlockUnexpectedEOF_5(t *testing.T) {
+	s := NewScanner(bytes.NewBufferString(`<%=`), "tmpl.ego")
+	_, err := s.Scan()
+	assert.Equal(t, err, io.ErrUnexpectedEOF)
+}
+
 // Ensure that a header block can be scanned.
 func TestScannerHeaderBlock(t *testing.T) {
 	s := NewScanner(bytes.NewBufferString(`<%% import "foo" %%>`), "tmpl.ego")
@@ -127,22 +134,33 @@ func TestScannerHeaderBlockUnexpectedEOF_5(t *testing.T) {
 
 // Ensure that a print block can be scanned.
 func TestScannerPrintBlock(t *testing.T) {
-	s := NewScanner(bytes.NewBufferString(`<%= myNum %>`), "tmpl.ego")
+	s := NewScanner(bytes.NewBufferString(`<%== myNum %>`), "tmpl.ego")
 	b, err := s.Scan()
 	assert.NoError(t, err)
-	if b, ok := b.(*PrintBlock); assert.True(t, ok) {
+	if b, ok := b.(*RawPrintBlock); assert.True(t, ok) {
 		assert.Equal(t, b.Content, ` myNum `)
 		assert.Equal(t, b.Pos, Pos{Path: "tmpl.ego", LineNo: 1})
 	}
 }
 
 // Ensure that a print block that ends unexpectedly returns an error.
 func TestScannerPrintBlockUnexpectedEOF(t *testing.T) {
-	s := NewScanner(bytes.NewBufferString(`<%= `), "tmpl.ego")
+	s := NewScanner(bytes.NewBufferString(`<%== `), "tmpl.ego")
 	_, err := s.Scan()
 	assert.Equal(t, err, io.ErrUnexpectedEOF)
 }
 
+// Ensure that an escaped print block can be scanned.
+func TestScannerEscapedPrintBlock(t *testing.T) {
+	s := NewScanner(bytes.NewBufferString(`<%= myNum %>`), "tmpl.ego")
+	b, err := s.Scan()
+	assert.NoError(t, err)
+	if b, ok := b.(*PrintBlock); assert.True(t, ok) {
+		assert.Equal(t, b.Content, ` myNum `)
+		assert.Equal(t, b.Pos, Pos{Path: "tmpl.ego", LineNo: 1})
+	}
+}
+
 // Ensure that a declaration block can be scanned.
 func TestScannerDeclarationBlock(t *testing.T) {
 	s := NewScanner(bytes.NewBufferString(`<%! MyTemplate(w io.Writer) error %>`), "tmpl.ego")
@@ -163,11 +181,11 @@ func TestScannerDeclarationBlockUnexpectedEOF(t *testing.T) {
 
 // Ensure that line numbers are tracked correctly.
 func TestScannerMultiline(t *testing.T) {
-	s := NewScanner(bytes.NewBufferString("hello\nworld<%= x \n\n %>goodbye"), "tmpl.ego")
+	s := NewScanner(bytes.NewBufferString("hello\nworld<%== x \n\n %>goodbye"), "tmpl.ego")
 	b, _ := s.Scan()
 	assert.Equal(t, b.(*TextBlock).Pos, Pos{Path: "tmpl.ego", LineNo: 1})
 	b, _ = s.Scan()
-	assert.Equal(t, b.(*PrintBlock).Pos, Pos{Path: "tmpl.ego", LineNo: 2})
+	assert.Equal(t, b.(*RawPrintBlock).Pos, Pos{Path: "tmpl.ego", LineNo: 2})
 	b, _ = s.Scan()
 	assert.Equal(t, b.(*TextBlock).Pos, Pos{Path: "tmpl.ego", LineNo: 4})
 }

template.go

@@ -14,7 +14,7 @@ import (
 
 // Template represents an entire Ego template.
 // A template consists of a declaration block followed by zero or more blocks.
-// Blocks can be either a TextBlock, a PrintBlock, or a CodeBlock.
+// Blocks can be either a TextBlock, a PrintBlock, a RawPrintBlock, or a CodeBlock.
 type Template struct {
 	Path   string
 	Blocks []Block
@@ -79,6 +79,15 @@ func (t *Template) nonHeaderBlocks() []Block {
 	return blocks
 }
 
+func (t *Template) hasEscapedPrintBlock() bool {
+	for _, b := range t.Blocks {
+		if _, ok := b.(*PrintBlock); ok {
+			return true
+		}
+	}
+	return false
+}
+
 // normalize joins together adjacent text blocks.
 func (t *Template) normalize() {
 	var a []Block
@@ -102,6 +111,7 @@ func (b *DeclarationBlock) block() {}
 func (b *TextBlock) block()        {}
 func (b *CodeBlock) block()        {}
 func (b *HeaderBlock) block()      {}
+func (b *RawPrintBlock) block()    {}
 func (b *PrintBlock) block()       {}
 
 // DeclarationBlock represents a block that declaration the function signature.
@@ -158,15 +168,27 @@ func (b *HeaderBlock) write(buf *bytes.Buffer) error {
 	return nil
 }
 
-// PrintBlock represents a block of the template that is printed out to the writer.
+// RawPrintBlock represents a block of the template that is printed out to the writer.
+type RawPrintBlock struct {
+	Pos     Pos
+	Content string
+}
+
+func (b *RawPrintBlock) write(buf *bytes.Buffer) error {
+	b.Pos.write(buf)
+	fmt.Fprintf(buf, `_, _ = fmt.Fprintf(w, "%%v", %s)`+"\n", b.Content)
+	return nil
+}
+
+// PrintBlock represents a block that will HTML escape the contents before outputting
 type PrintBlock struct {
 	Pos     Pos
 	Content string
 }
 
 func (b *PrintBlock) write(buf *bytes.Buffer) error {
 	b.Pos.write(buf)
-	fmt.Fprintf(buf, `_, _ = fmt.Fprintf(w, "%%v", %s)`+"\n", b.Content)
+	fmt.Fprintf(buf, `_, _ = fmt.Fprint(w, html.EscapeString(fmt.Sprintf("%%v", %s)))`+"\n", b.Content)
 	return nil
 }
 
@@ -237,7 +259,15 @@ func (p *Package) writeHeader(w io.Writer) error {
 	var decls = map[string]bool{`:"fmt"`: true, `:"io"`: true}
 	fmt.Fprint(&buf, "import (\n")
 	fmt.Fprintln(&buf, `"fmt"`)
+	for _, t := range p.Templates {
+		if t.hasEscapedPrintBlock() {
+			fmt.Fprintln(&buf, `"html"`)
+			decls["html"] = true
+			break
+		}
+	}
 	fmt.Fprintln(&buf, `"io"`)
+
 	for _, d := range f.Decls {
 		d, ok := d.(*ast.GenDecl)
 		if !ok || d.Tok != token.IMPORT {

template_test.go

@@ -20,7 +20,7 @@ func TestTemplate_Write(t *testing.T) {
 			&DeclarationBlock{Content: " func MyTemplate(w io.Writer, nums []int) error "},
 			&CodeBlock{Content: "  for _, num := range nums {"},
 			&TextBlock{Content: "    <p>"},
-			&PrintBlock{Content: "num + 1"},
+			&RawPrintBlock{Content: "num + 1"},
 			&TextBlock{Content: "    </p>"},
 			&CodeBlock{Content: "  }"},
 			&TextBlock{Content: "</html>"},