armor/decoder.go

   1 package armor
   2
   3 import (
   4         "bufio"
   5         "bytes"
   6         "encoding/base64"
   7         "fmt"
   8         "io"
   9
  10         "golang.org/x/net/html"
  11 )
  12
  13 // ErrUnknownVersion is the error returned when the first character inside the
  14 // element encoding (but outside the base64 encoding) is not '0'.
  15 type ErrUnknownVersion byte
  16
  17 func (err ErrUnknownVersion) Error() string {
  18         return fmt.Sprintf("unknown armor version indicator %+q", byte(err))
  19 }
  20
  21 func isASCIIWhitespace(b byte) bool {
  22         switch b {
  23         // https://infra.spec.whatwg.org/#ascii-whitespace
  24         case '\x09', '\x0a', '\x0c', '\x0d', '\x20':
  25                 return true
  26         default:
  27                 return false
  28         }
  29 }
  30
  31 func splitASCIIWhitespace(data []byte, atEOF bool) (advance int, token []byte, err error) {
  32         var i, j int
  33         // Skip initial whitespace.
  34         for i = 0; i < len(data); i++ {
  35                 if !isASCIIWhitespace(data[i]) {
  36                         break
  37                 }
  38         }
  39         // Look for next whitespace.
  40         for j = i; j < len(data); j++ {
  41                 if isASCIIWhitespace(data[j]) {
  42                         return j + 1, data[i:j], nil
  43                 }
  44         }
  45         // We reached the end of data without finding more whitespace. Only
  46         // consider it a token if we are at EOF.
  47         if atEOF && i < j {
  48                 return j, data[i:j], nil
  49         }
  50         // Otherwise, request more data.
  51         return i, nil, nil
  52 }
  53
  54 func decodeToWriter(w io.Writer, r io.Reader) (int64, error) {
  55         tokenizer := html.NewTokenizer(r)
  56         // Set a memory limit on token sizes, otherwise the tokenizer will
  57         // buffer text indefinitely if it is not broken up by other token types.
  58         tokenizer.SetMaxBuf(elementSizeLimit)
  59         active := false
  60         total := int64(0)
  61         for {
  62                 tt := tokenizer.Next()
  63                 switch tt {
  64                 case html.ErrorToken:
  65                         err := tokenizer.Err()
  66                         if err == io.EOF {
  67                                 err = nil
  68                         }
  69                         if err == nil && active {
  70                                 return total, fmt.Errorf("missing </pre> tag")
  71                         }
  72                         return total, err
  73                 case html.TextToken:
  74                         if active {
  75                                 // Re-join the separate chunks of text and
  76                                 // feed them to the decoder.
  77                                 scanner := bufio.NewScanner(bytes.NewReader(tokenizer.Text()))
  78                                 scanner.Split(splitASCIIWhitespace)
  79                                 for scanner.Scan() {
  80                                         n, err := w.Write(scanner.Bytes())
  81                                         total += int64(n)
  82                                         if err != nil {
  83                                                 return total, err
  84                                         }
  85                                 }
  86                                 if err := scanner.Err(); err != nil {
  87                                         return total, err
  88                                 }
  89                         }
  90                 case html.StartTagToken:
  91                         tn, _ := tokenizer.TagName()
  92                         if string(tn) == "pre" {
  93                                 if active {
  94                                         // nesting not allowed
  95                                         return total, fmt.Errorf("unexpected %s", tokenizer.Token())
  96                                 }
  97                                 active = true
  98                         }
  99                 case html.EndTagToken:
 100                         tn, _ := tokenizer.TagName()
 101                         if string(tn) == "pre" {
 102                                 if !active {
 103                                         // stray end tag
 104                                         return total, fmt.Errorf("unexpected %s", tokenizer.Token())
 105                                 }
 106                                 active = false
 107                         }
 108                 }
 109         }
 110 }
 111
 112 // NewDecoder returns a new AMP armor decoder.
 113 func NewDecoder(r io.Reader) (io.Reader, error) {
 114         pr, pw := io.Pipe()
 115         go func() {
 116                 _, err := decodeToWriter(pw, r)
 117                 pw.CloseWithError(err)
 118         }()
 119
 120         // The first byte inside the element encoding is a server–client
 121         // protocol version indicator.
 122         var version [1]byte
 123         _, err := pr.Read(version[:])
 124         if err != nil {
 125                 pr.CloseWithError(err)
 126                 return nil, err
 127         }
 128         switch version[0] {
 129         case '0':
 130                 return base64.NewDecoder(base64.StdEncoding, pr), nil
 131         default:
 132                 err := ErrUnknownVersion(version[0])
 133                 pr.CloseWithError(err)
 134                 return nil, err
 135         }
 136 }