encode.go 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293
  1. package bencode
  2. import (
  3. "io"
  4. "math/big"
  5. "reflect"
  6. "runtime"
  7. "sort"
  8. "strconv"
  9. "sync"
  10. "github.com/anacrolix/missinggo"
  11. )
  12. func isEmptyValue(v reflect.Value) bool {
  13. return missinggo.IsEmptyValue(v)
  14. }
  15. type Encoder struct {
  16. w io.Writer
  17. scratch [64]byte
  18. }
  19. func (e *Encoder) Encode(v interface{}) (err error) {
  20. if v == nil {
  21. return
  22. }
  23. defer func() {
  24. if e := recover(); e != nil {
  25. if _, ok := e.(runtime.Error); ok {
  26. panic(e)
  27. }
  28. var ok bool
  29. err, ok = e.(error)
  30. if !ok {
  31. panic(e)
  32. }
  33. }
  34. }()
  35. e.reflectValue(reflect.ValueOf(v))
  36. return nil
  37. }
  38. type stringValues []reflect.Value
  39. func (sv stringValues) Len() int { return len(sv) }
  40. func (sv stringValues) Swap(i, j int) { sv[i], sv[j] = sv[j], sv[i] }
  41. func (sv stringValues) Less(i, j int) bool { return sv.get(i) < sv.get(j) }
  42. func (sv stringValues) get(i int) string { return sv[i].String() }
  43. func (e *Encoder) write(s []byte) {
  44. _, err := e.w.Write(s)
  45. if err != nil {
  46. panic(err)
  47. }
  48. }
  49. func (e *Encoder) writeString(s string) {
  50. for s != "" {
  51. n := copy(e.scratch[:], s)
  52. s = s[n:]
  53. e.write(e.scratch[:n])
  54. }
  55. }
  56. func (e *Encoder) reflectString(s string) {
  57. e.writeStringPrefix(int64(len(s)))
  58. e.writeString(s)
  59. }
  60. func (e *Encoder) writeStringPrefix(l int64) {
  61. b := strconv.AppendInt(e.scratch[:0], l, 10)
  62. e.write(b)
  63. e.writeString(":")
  64. }
  65. func (e *Encoder) reflectByteSlice(s []byte) {
  66. e.writeStringPrefix(int64(len(s)))
  67. e.write(s)
  68. }
  69. // Returns true if the value implements Marshaler interface and marshaling was
  70. // done successfully.
  71. func (e *Encoder) reflectMarshaler(v reflect.Value) bool {
  72. if !v.Type().Implements(marshalerType) {
  73. if v.Kind() != reflect.Ptr && v.CanAddr() && v.Addr().Type().Implements(marshalerType) {
  74. v = v.Addr()
  75. } else {
  76. return false
  77. }
  78. }
  79. m := v.Interface().(Marshaler)
  80. data, err := m.MarshalBencode()
  81. if err != nil {
  82. panic(&MarshalerError{v.Type(), err})
  83. }
  84. e.write(data)
  85. return true
  86. }
  87. var bigIntType = reflect.TypeOf((*big.Int)(nil)).Elem()
  88. func (e *Encoder) reflectValue(v reflect.Value) {
  89. if e.reflectMarshaler(v) {
  90. return
  91. }
  92. if v.Type() == bigIntType {
  93. e.writeString("i")
  94. bi := v.Interface().(big.Int)
  95. e.writeString(bi.String())
  96. e.writeString("e")
  97. return
  98. }
  99. switch v.Kind() {
  100. case reflect.Bool:
  101. if v.Bool() {
  102. e.writeString("i1e")
  103. } else {
  104. e.writeString("i0e")
  105. }
  106. case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
  107. e.writeString("i")
  108. b := strconv.AppendInt(e.scratch[:0], v.Int(), 10)
  109. e.write(b)
  110. e.writeString("e")
  111. case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
  112. e.writeString("i")
  113. b := strconv.AppendUint(e.scratch[:0], v.Uint(), 10)
  114. e.write(b)
  115. e.writeString("e")
  116. case reflect.String:
  117. e.reflectString(v.String())
  118. case reflect.Struct:
  119. e.writeString("d")
  120. for _, ef := range getEncodeFields(v.Type()) {
  121. fieldValue := ef.i(v)
  122. if !fieldValue.IsValid() {
  123. continue
  124. }
  125. if ef.omitEmpty && isEmptyValue(fieldValue) {
  126. continue
  127. }
  128. e.reflectString(ef.tag)
  129. e.reflectValue(fieldValue)
  130. }
  131. e.writeString("e")
  132. case reflect.Map:
  133. if v.Type().Key().Kind() != reflect.String {
  134. panic(&MarshalTypeError{v.Type()})
  135. }
  136. if v.IsNil() {
  137. e.writeString("de")
  138. break
  139. }
  140. e.writeString("d")
  141. sv := stringValues(v.MapKeys())
  142. sort.Sort(sv)
  143. for _, key := range sv {
  144. e.reflectString(key.String())
  145. e.reflectValue(v.MapIndex(key))
  146. }
  147. e.writeString("e")
  148. case reflect.Slice, reflect.Array:
  149. e.reflectSequence(v)
  150. case reflect.Interface:
  151. e.reflectValue(v.Elem())
  152. case reflect.Ptr:
  153. if v.IsNil() {
  154. v = reflect.Zero(v.Type().Elem())
  155. } else {
  156. v = v.Elem()
  157. }
  158. e.reflectValue(v)
  159. default:
  160. panic(&MarshalTypeError{v.Type()})
  161. }
  162. }
  163. func (e *Encoder) reflectSequence(v reflect.Value) {
  164. // Use bencode string-type
  165. if v.Type().Elem().Kind() == reflect.Uint8 {
  166. if v.Kind() != reflect.Slice {
  167. // Can't use []byte optimization
  168. if !v.CanAddr() {
  169. e.writeStringPrefix(int64(v.Len()))
  170. for i := 0; i < v.Len(); i++ {
  171. var b [1]byte
  172. b[0] = byte(v.Index(i).Uint())
  173. e.write(b[:])
  174. }
  175. return
  176. }
  177. v = v.Slice(0, v.Len())
  178. }
  179. s := v.Bytes()
  180. e.reflectByteSlice(s)
  181. return
  182. }
  183. if v.IsNil() {
  184. e.writeString("le")
  185. return
  186. }
  187. e.writeString("l")
  188. for i, n := 0, v.Len(); i < n; i++ {
  189. e.reflectValue(v.Index(i))
  190. }
  191. e.writeString("e")
  192. }
  193. type encodeField struct {
  194. i func(v reflect.Value) reflect.Value
  195. tag string
  196. omitEmpty bool
  197. }
  198. type encodeFieldsSortType []encodeField
  199. func (ef encodeFieldsSortType) Len() int { return len(ef) }
  200. func (ef encodeFieldsSortType) Swap(i, j int) { ef[i], ef[j] = ef[j], ef[i] }
  201. func (ef encodeFieldsSortType) Less(i, j int) bool { return ef[i].tag < ef[j].tag }
  202. var (
  203. typeCacheLock sync.RWMutex
  204. encodeFieldsCache = make(map[reflect.Type][]encodeField)
  205. )
  206. func getEncodeFields(t reflect.Type) []encodeField {
  207. typeCacheLock.RLock()
  208. fs, ok := encodeFieldsCache[t]
  209. typeCacheLock.RUnlock()
  210. if ok {
  211. return fs
  212. }
  213. fs = makeEncodeFields(t)
  214. typeCacheLock.Lock()
  215. defer typeCacheLock.Unlock()
  216. encodeFieldsCache[t] = fs
  217. return fs
  218. }
  219. func makeEncodeFields(t reflect.Type) (fs []encodeField) {
  220. for _i, n := 0, t.NumField(); _i < n; _i++ {
  221. i := _i
  222. f := t.Field(i)
  223. if f.PkgPath != "" {
  224. continue
  225. }
  226. if f.Anonymous {
  227. t := f.Type
  228. if t.Kind() == reflect.Ptr {
  229. t = t.Elem()
  230. }
  231. anonEFs := makeEncodeFields(t)
  232. for aefi := range anonEFs {
  233. anonEF := anonEFs[aefi]
  234. bottomField := anonEF
  235. bottomField.i = func(v reflect.Value) reflect.Value {
  236. v = v.Field(i)
  237. if v.Kind() == reflect.Ptr {
  238. if v.IsNil() {
  239. // This will skip serializing this value.
  240. return reflect.Value{}
  241. }
  242. v = v.Elem()
  243. }
  244. return anonEF.i(v)
  245. }
  246. fs = append(fs, bottomField)
  247. }
  248. continue
  249. }
  250. var ef encodeField
  251. ef.i = func(v reflect.Value) reflect.Value {
  252. return v.Field(i)
  253. }
  254. ef.tag = f.Name
  255. tv := getTag(f.Tag)
  256. if tv.Ignore() {
  257. continue
  258. }
  259. if tv.Key() != "" {
  260. ef.tag = tv.Key()
  261. }
  262. ef.omitEmpty = tv.OmitEmpty()
  263. fs = append(fs, ef)
  264. }
  265. fss := encodeFieldsSortType(fs)
  266. sort.Sort(fss)
  267. return fs
  268. }