Matteo Scandolo | a6a3aee | 2019-11-26 13:30:14 -0700 | [diff] [blame^] | 1 | // Copyright 2014 The Go Authors. All rights reserved. |
| 2 | // Use of this source code is governed by a BSD-style |
| 3 | // license that can be found in the LICENSE file. |
| 4 | |
| 5 | package hpack |
| 6 | |
| 7 | import ( |
| 8 | "bytes" |
| 9 | "errors" |
| 10 | "io" |
| 11 | "sync" |
| 12 | ) |
| 13 | |
| 14 | var bufPool = sync.Pool{ |
| 15 | New: func() interface{} { return new(bytes.Buffer) }, |
| 16 | } |
| 17 | |
| 18 | // HuffmanDecode decodes the string in v and writes the expanded |
| 19 | // result to w, returning the number of bytes written to w and the |
| 20 | // Write call's return value. At most one Write call is made. |
| 21 | func HuffmanDecode(w io.Writer, v []byte) (int, error) { |
| 22 | buf := bufPool.Get().(*bytes.Buffer) |
| 23 | buf.Reset() |
| 24 | defer bufPool.Put(buf) |
| 25 | if err := huffmanDecode(buf, 0, v); err != nil { |
| 26 | return 0, err |
| 27 | } |
| 28 | return w.Write(buf.Bytes()) |
| 29 | } |
| 30 | |
| 31 | // HuffmanDecodeToString decodes the string in v. |
| 32 | func HuffmanDecodeToString(v []byte) (string, error) { |
| 33 | buf := bufPool.Get().(*bytes.Buffer) |
| 34 | buf.Reset() |
| 35 | defer bufPool.Put(buf) |
| 36 | if err := huffmanDecode(buf, 0, v); err != nil { |
| 37 | return "", err |
| 38 | } |
| 39 | return buf.String(), nil |
| 40 | } |
| 41 | |
| 42 | // ErrInvalidHuffman is returned for errors found decoding |
| 43 | // Huffman-encoded strings. |
| 44 | var ErrInvalidHuffman = errors.New("hpack: invalid Huffman-encoded data") |
| 45 | |
| 46 | // huffmanDecode decodes v to buf. |
| 47 | // If maxLen is greater than 0, attempts to write more to buf than |
| 48 | // maxLen bytes will return ErrStringLength. |
| 49 | func huffmanDecode(buf *bytes.Buffer, maxLen int, v []byte) error { |
| 50 | rootHuffmanNode := getRootHuffmanNode() |
| 51 | n := rootHuffmanNode |
| 52 | // cur is the bit buffer that has not been fed into n. |
| 53 | // cbits is the number of low order bits in cur that are valid. |
| 54 | // sbits is the number of bits of the symbol prefix being decoded. |
| 55 | cur, cbits, sbits := uint(0), uint8(0), uint8(0) |
| 56 | for _, b := range v { |
| 57 | cur = cur<<8 | uint(b) |
| 58 | cbits += 8 |
| 59 | sbits += 8 |
| 60 | for cbits >= 8 { |
| 61 | idx := byte(cur >> (cbits - 8)) |
| 62 | n = n.children[idx] |
| 63 | if n == nil { |
| 64 | return ErrInvalidHuffman |
| 65 | } |
| 66 | if n.children == nil { |
| 67 | if maxLen != 0 && buf.Len() == maxLen { |
| 68 | return ErrStringLength |
| 69 | } |
| 70 | buf.WriteByte(n.sym) |
| 71 | cbits -= n.codeLen |
| 72 | n = rootHuffmanNode |
| 73 | sbits = cbits |
| 74 | } else { |
| 75 | cbits -= 8 |
| 76 | } |
| 77 | } |
| 78 | } |
| 79 | for cbits > 0 { |
| 80 | n = n.children[byte(cur<<(8-cbits))] |
| 81 | if n == nil { |
| 82 | return ErrInvalidHuffman |
| 83 | } |
| 84 | if n.children != nil || n.codeLen > cbits { |
| 85 | break |
| 86 | } |
| 87 | if maxLen != 0 && buf.Len() == maxLen { |
| 88 | return ErrStringLength |
| 89 | } |
| 90 | buf.WriteByte(n.sym) |
| 91 | cbits -= n.codeLen |
| 92 | n = rootHuffmanNode |
| 93 | sbits = cbits |
| 94 | } |
| 95 | if sbits > 7 { |
| 96 | // Either there was an incomplete symbol, or overlong padding. |
| 97 | // Both are decoding errors per RFC 7541 section 5.2. |
| 98 | return ErrInvalidHuffman |
| 99 | } |
| 100 | if mask := uint(1<<cbits - 1); cur&mask != mask { |
| 101 | // Trailing bits must be a prefix of EOS per RFC 7541 section 5.2. |
| 102 | return ErrInvalidHuffman |
| 103 | } |
| 104 | |
| 105 | return nil |
| 106 | } |
| 107 | |
| 108 | type node struct { |
| 109 | // children is non-nil for internal nodes |
| 110 | children *[256]*node |
| 111 | |
| 112 | // The following are only valid if children is nil: |
| 113 | codeLen uint8 // number of bits that led to the output of sym |
| 114 | sym byte // output symbol |
| 115 | } |
| 116 | |
| 117 | func newInternalNode() *node { |
| 118 | return &node{children: new([256]*node)} |
| 119 | } |
| 120 | |
| 121 | var ( |
| 122 | buildRootOnce sync.Once |
| 123 | lazyRootHuffmanNode *node |
| 124 | ) |
| 125 | |
| 126 | func getRootHuffmanNode() *node { |
| 127 | buildRootOnce.Do(buildRootHuffmanNode) |
| 128 | return lazyRootHuffmanNode |
| 129 | } |
| 130 | |
| 131 | func buildRootHuffmanNode() { |
| 132 | if len(huffmanCodes) != 256 { |
| 133 | panic("unexpected size") |
| 134 | } |
| 135 | lazyRootHuffmanNode = newInternalNode() |
| 136 | for i, code := range huffmanCodes { |
| 137 | addDecoderNode(byte(i), code, huffmanCodeLen[i]) |
| 138 | } |
| 139 | } |
| 140 | |
| 141 | func addDecoderNode(sym byte, code uint32, codeLen uint8) { |
| 142 | cur := lazyRootHuffmanNode |
| 143 | for codeLen > 8 { |
| 144 | codeLen -= 8 |
| 145 | i := uint8(code >> codeLen) |
| 146 | if cur.children[i] == nil { |
| 147 | cur.children[i] = newInternalNode() |
| 148 | } |
| 149 | cur = cur.children[i] |
| 150 | } |
| 151 | shift := 8 - codeLen |
| 152 | start, end := int(uint8(code<<shift)), int(1<<shift) |
| 153 | for i := start; i < start+end; i++ { |
| 154 | cur.children[i] = &node{sym: sym, codeLen: codeLen} |
| 155 | } |
| 156 | } |
| 157 | |
| 158 | // AppendHuffmanString appends s, as encoded in Huffman codes, to dst |
| 159 | // and returns the extended buffer. |
| 160 | func AppendHuffmanString(dst []byte, s string) []byte { |
| 161 | rembits := uint8(8) |
| 162 | |
| 163 | for i := 0; i < len(s); i++ { |
| 164 | if rembits == 8 { |
| 165 | dst = append(dst, 0) |
| 166 | } |
| 167 | dst, rembits = appendByteToHuffmanCode(dst, rembits, s[i]) |
| 168 | } |
| 169 | |
| 170 | if rembits < 8 { |
| 171 | // special EOS symbol |
| 172 | code := uint32(0x3fffffff) |
| 173 | nbits := uint8(30) |
| 174 | |
| 175 | t := uint8(code >> (nbits - rembits)) |
| 176 | dst[len(dst)-1] |= t |
| 177 | } |
| 178 | |
| 179 | return dst |
| 180 | } |
| 181 | |
| 182 | // HuffmanEncodeLength returns the number of bytes required to encode |
| 183 | // s in Huffman codes. The result is round up to byte boundary. |
| 184 | func HuffmanEncodeLength(s string) uint64 { |
| 185 | n := uint64(0) |
| 186 | for i := 0; i < len(s); i++ { |
| 187 | n += uint64(huffmanCodeLen[s[i]]) |
| 188 | } |
| 189 | return (n + 7) / 8 |
| 190 | } |
| 191 | |
| 192 | // appendByteToHuffmanCode appends Huffman code for c to dst and |
| 193 | // returns the extended buffer and the remaining bits in the last |
| 194 | // element. The appending is not byte aligned and the remaining bits |
| 195 | // in the last element of dst is given in rembits. |
| 196 | func appendByteToHuffmanCode(dst []byte, rembits uint8, c byte) ([]byte, uint8) { |
| 197 | code := huffmanCodes[c] |
| 198 | nbits := huffmanCodeLen[c] |
| 199 | |
| 200 | for { |
| 201 | if rembits > nbits { |
| 202 | t := uint8(code << (rembits - nbits)) |
| 203 | dst[len(dst)-1] |= t |
| 204 | rembits -= nbits |
| 205 | break |
| 206 | } |
| 207 | |
| 208 | t := uint8(code >> (nbits - rembits)) |
| 209 | dst[len(dst)-1] |= t |
| 210 | |
| 211 | nbits -= rembits |
| 212 | rembits = 8 |
| 213 | |
| 214 | if nbits == 0 { |
| 215 | break |
| 216 | } |
| 217 | |
| 218 | dst = append(dst, 0) |
| 219 | } |
| 220 | |
| 221 | return dst, rembits |
| 222 | } |