Support marshalling http tracker response peers

[btrtrc.git] / tracker / http / protocol.go

diff --git a/tracker/http/protocol.go b/tracker/http/protocol.go
index 5ce9fa90cf2d32beea36ee1a18afca5ea1b75f29..11a90b2a315a150cddcb3340be48ee99d2883916 100644 (file)
--- a/tracker/http/protocol.go
+++ b/tracker/http/protocol.go
@@ -19,7 +19,30 @@ type HttpResponse struct {
        Peers6 krpc.CompactIPv6NodeAddrs `bencode:"peers6"`
 }
 
-type Peers []Peer
+type Peers struct {
+       List    []Peer
+       Compact bool
+}
+
+func (me Peers) MarshalBencode() ([]byte, error) {
+       if me.Compact {
+               cnas := make([]krpc.NodeAddr, 0, len(me.List))
+               for _, peer := range me.List {
+                       cnas = append(cnas, krpc.NodeAddr{
+                               IP:   peer.IP,
+                               Port: peer.Port,
+                       })
+               }
+               return krpc.CompactIPv4NodeAddrs(cnas).MarshalBencode()
+       } else {
+               return bencode.Marshal(me.List)
+       }
+}
+
+var (
+       _ bencode.Unmarshaler = (*Peers)(nil)
+       _ bencode.Marshaler   = Peers{}
+)
 
 func (me *Peers) UnmarshalBencode(b []byte) (err error) {
        var _v interface{}
@@ -35,19 +58,21 @@ func (me *Peers) UnmarshalBencode(b []byte) (err error) {
                if err != nil {
                        return
                }
+               me.Compact = true
                for _, cp := range cnas {
-                       *me = append(*me, Peer{
+                       me.List = append(me.List, Peer{
                                IP:   cp.IP[:],
-                               Port: int(cp.Port),
+                               Port: cp.Port,
                        })
                }
                return
        case []interface{}:
                vars.Add("http responses with list peers", 1)
+               me.Compact = false
                for _, i := range v {
                        var p Peer
                        p.FromDictInterface(i.(map[string]interface{}))
-                       *me = append(*me, p)
+                       me.List = append(me.List, p)
                }
                return
        default: