]>
Commit | Line | Data |
---|---|---|
15c0b25d AP |
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 { | |
107c1cdb | 50 | rootHuffmanNode := getRootHuffmanNode() |
15c0b25d AP |
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 | |
107c1cdb | 110 | children *[256]*node |
15c0b25d AP |
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 { | |
107c1cdb | 118 | return &node{children: new([256]*node)} |
15c0b25d AP |
119 | } |
120 | ||
107c1cdb ND |
121 | var ( |
122 | buildRootOnce sync.Once | |
123 | lazyRootHuffmanNode *node | |
124 | ) | |
125 | ||
126 | func getRootHuffmanNode() *node { | |
127 | buildRootOnce.Do(buildRootHuffmanNode) | |
128 | return lazyRootHuffmanNode | |
129 | } | |
15c0b25d | 130 | |
107c1cdb | 131 | func buildRootHuffmanNode() { |
15c0b25d AP |
132 | if len(huffmanCodes) != 256 { |
133 | panic("unexpected size") | |
134 | } | |
107c1cdb | 135 | lazyRootHuffmanNode = newInternalNode() |
15c0b25d AP |
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) { | |
107c1cdb | 142 | cur := lazyRootHuffmanNode |
15c0b25d AP |
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 | } |