huffman.go 5.0 KB

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