compact_helpers.go 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283
  1. package krpc
  2. import (
  3. "encoding"
  4. "fmt"
  5. "reflect"
  6. "github.com/anacrolix/missinggo/slices"
  7. "github.com/anacrolix/torrent/bencode"
  8. )
  9. func unmarshalBencodedBinary(u encoding.BinaryUnmarshaler, b []byte) (err error) {
  10. var ub string
  11. err = bencode.Unmarshal(b, &ub)
  12. if err != nil {
  13. return
  14. }
  15. return u.UnmarshalBinary([]byte(ub))
  16. }
  17. type elemSizer interface {
  18. ElemSize() int
  19. }
  20. func unmarshalBinarySlice(slice elemSizer, b []byte) (err error) {
  21. sliceValue := reflect.ValueOf(slice).Elem()
  22. elemType := sliceValue.Type().Elem()
  23. bytesPerElem := slice.ElemSize()
  24. elem := reflect.New(elemType)
  25. for len(b) != 0 {
  26. if len(b) < bytesPerElem {
  27. err = fmt.Errorf("%d trailing bytes < %d required for element", len(b), bytesPerElem)
  28. break
  29. }
  30. if bu, ok := elem.Interface().(encoding.BinaryUnmarshaler); ok {
  31. err = bu.UnmarshalBinary(b[:bytesPerElem])
  32. } else if elem.Elem().Len() == bytesPerElem {
  33. reflect.Copy(elem.Elem(), reflect.ValueOf(b[:bytesPerElem]))
  34. } else {
  35. err = fmt.Errorf("can't unmarshal %v bytes into %v", bytesPerElem, elem.Type())
  36. }
  37. if err != nil {
  38. return
  39. }
  40. sliceValue.Set(reflect.Append(sliceValue, elem.Elem()))
  41. b = b[bytesPerElem:]
  42. }
  43. return
  44. }
  45. func marshalBinarySlice(slice elemSizer) (ret []byte, err error) {
  46. var elems []encoding.BinaryMarshaler
  47. slices.MakeInto(&elems, slice)
  48. for _, e := range elems {
  49. var b []byte
  50. b, err = e.MarshalBinary()
  51. if err != nil {
  52. return
  53. }
  54. if len(b) != slice.ElemSize() {
  55. panic(fmt.Sprintf("marshalled %d bytes, but expected %d", len(b), slice.ElemSize()))
  56. }
  57. ret = append(ret, b...)
  58. }
  59. return
  60. }
  61. func bencodeBytesResult(b []byte, err error) ([]byte, error) {
  62. if err != nil {
  63. return b, err
  64. }
  65. return bencode.Marshal(b)
  66. }
  67. // returns position of x in v, or -1 if not found
  68. func addrIndex(v []NodeAddr, x NodeAddr) int {
  69. for i := 0; i < len(v); i += 1 {
  70. if v[i].Equal(x) {
  71. return i
  72. }
  73. }
  74. return -1
  75. }