github.com/zignig/go-ipfs@v0.0.0-20141111235910-c9e5fdf55a52/exchange/bitswap/message/message.go (about)

     1  package message
     2  
     3  import (
     4  	proto "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/goprotobuf/proto"
     5  	blocks "github.com/jbenet/go-ipfs/blocks"
     6  	pb "github.com/jbenet/go-ipfs/exchange/bitswap/message/internal/pb"
     7  	netmsg "github.com/jbenet/go-ipfs/net/message"
     8  	nm "github.com/jbenet/go-ipfs/net/message"
     9  	peer "github.com/jbenet/go-ipfs/peer"
    10  	u "github.com/jbenet/go-ipfs/util"
    11  )
    12  
    13  // TODO move message.go into the bitswap package
    14  // TODO move bs/msg/internal/pb to bs/internal/pb and rename pb package to bitswap_pb
    15  
    16  type BitSwapMessage interface {
    17  	// Wantlist returns a slice of unique keys that represent data wanted by
    18  	// the sender.
    19  	Wantlist() []u.Key
    20  
    21  	// Blocks returns a slice of unique blocks
    22  	Blocks() []blocks.Block
    23  
    24  	// AddWanted adds the key to the Wantlist.
    25  	//
    26  	// Insertion order determines priority. That is, earlier insertions are
    27  	// deemed higher priority than keys inserted later.
    28  	//
    29  	// t = 0, msg.AddWanted(A)
    30  	// t = 1, msg.AddWanted(B)
    31  	//
    32  	// implies Priority(A) > Priority(B)
    33  	AddWanted(u.Key)
    34  
    35  	AddBlock(blocks.Block)
    36  	Exportable
    37  }
    38  
    39  type Exportable interface {
    40  	ToProto() *pb.Message
    41  	ToNet(p peer.Peer) (nm.NetMessage, error)
    42  }
    43  
    44  type impl struct {
    45  	existsInWantlist map[u.Key]struct{}     // map to detect duplicates
    46  	wantlist         []u.Key                // slice to preserve ordering
    47  	blocks           map[u.Key]blocks.Block // map to detect duplicates
    48  }
    49  
    50  func New() BitSwapMessage {
    51  	return &impl{
    52  		blocks:           make(map[u.Key]blocks.Block),
    53  		existsInWantlist: make(map[u.Key]struct{}),
    54  		wantlist:         make([]u.Key, 0),
    55  	}
    56  }
    57  
    58  func newMessageFromProto(pbm pb.Message) BitSwapMessage {
    59  	m := New()
    60  	for _, s := range pbm.GetWantlist() {
    61  		m.AddWanted(u.Key(s))
    62  	}
    63  	for _, d := range pbm.GetBlocks() {
    64  		b := blocks.NewBlock(d)
    65  		m.AddBlock(*b)
    66  	}
    67  	return m
    68  }
    69  
    70  func (m *impl) Wantlist() []u.Key {
    71  	return m.wantlist
    72  }
    73  
    74  func (m *impl) Blocks() []blocks.Block {
    75  	bs := make([]blocks.Block, 0)
    76  	for _, block := range m.blocks {
    77  		bs = append(bs, block)
    78  	}
    79  	return bs
    80  }
    81  
    82  func (m *impl) AddWanted(k u.Key) {
    83  	_, exists := m.existsInWantlist[k]
    84  	if exists {
    85  		return
    86  	}
    87  	m.existsInWantlist[k] = struct{}{}
    88  	m.wantlist = append(m.wantlist, k)
    89  }
    90  
    91  func (m *impl) AddBlock(b blocks.Block) {
    92  	m.blocks[b.Key()] = b
    93  }
    94  
    95  func FromNet(nmsg netmsg.NetMessage) (BitSwapMessage, error) {
    96  	pb := new(pb.Message)
    97  	if err := proto.Unmarshal(nmsg.Data(), pb); err != nil {
    98  		return nil, err
    99  	}
   100  	m := newMessageFromProto(*pb)
   101  	return m, nil
   102  }
   103  
   104  func (m *impl) ToProto() *pb.Message {
   105  	pb := new(pb.Message)
   106  	for _, k := range m.Wantlist() {
   107  		pb.Wantlist = append(pb.Wantlist, string(k))
   108  	}
   109  	for _, b := range m.Blocks() {
   110  		pb.Blocks = append(pb.Blocks, b.Data)
   111  	}
   112  	return pb
   113  }
   114  
   115  func (m *impl) ToNet(p peer.Peer) (nm.NetMessage, error) {
   116  	return nm.FromObject(p, m.ToProto())
   117  }