github.com/keltia/go-ipfs@v0.3.8-0.20150909044612-210793031c63/exchange/bitswap/message/message.go (about)

     1  package message
     2  
     3  import (
     4  	"io"
     5  
     6  	blocks "github.com/ipfs/go-ipfs/blocks"
     7  	key "github.com/ipfs/go-ipfs/blocks/key"
     8  	pb "github.com/ipfs/go-ipfs/exchange/bitswap/message/pb"
     9  	wantlist "github.com/ipfs/go-ipfs/exchange/bitswap/wantlist"
    10  	inet "github.com/ipfs/go-ipfs/p2p/net"
    11  
    12  	ggio "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/gogo/protobuf/io"
    13  	proto "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/gogo/protobuf/proto"
    14  )
    15  
    16  // TODO move message.go into the bitswap package
    17  // TODO move bs/msg/internal/pb to bs/internal/pb and rename pb package to bitswap_pb
    18  
    19  type BitSwapMessage interface {
    20  	// Wantlist returns a slice of unique keys that represent data wanted by
    21  	// the sender.
    22  	Wantlist() []Entry
    23  
    24  	// Blocks returns a slice of unique blocks
    25  	Blocks() []*blocks.Block
    26  
    27  	// AddEntry adds an entry to the Wantlist.
    28  	AddEntry(key key.Key, priority int)
    29  
    30  	Cancel(key key.Key)
    31  
    32  	Empty() bool
    33  
    34  	// A full wantlist is an authoritative copy, a 'non-full' wantlist is a patch-set
    35  	Full() bool
    36  
    37  	AddBlock(*blocks.Block)
    38  	Exportable
    39  
    40  	Loggable() map[string]interface{}
    41  }
    42  
    43  type Exportable interface {
    44  	ToProto() *pb.Message
    45  	ToNet(w io.Writer) error
    46  }
    47  
    48  type impl struct {
    49  	full     bool
    50  	wantlist map[key.Key]Entry
    51  	blocks   map[key.Key]*blocks.Block
    52  }
    53  
    54  func New(full bool) BitSwapMessage {
    55  	return newMsg(full)
    56  }
    57  
    58  func newMsg(full bool) *impl {
    59  	return &impl{
    60  		blocks:   make(map[key.Key]*blocks.Block),
    61  		wantlist: make(map[key.Key]Entry),
    62  		full:     full,
    63  	}
    64  }
    65  
    66  type Entry struct {
    67  	wantlist.Entry
    68  	Cancel bool
    69  }
    70  
    71  func newMessageFromProto(pbm pb.Message) BitSwapMessage {
    72  	m := newMsg(pbm.GetWantlist().GetFull())
    73  	for _, e := range pbm.GetWantlist().GetEntries() {
    74  		m.addEntry(key.Key(e.GetBlock()), int(e.GetPriority()), e.GetCancel())
    75  	}
    76  	for _, d := range pbm.GetBlocks() {
    77  		b := blocks.NewBlock(d)
    78  		m.AddBlock(b)
    79  	}
    80  	return m
    81  }
    82  
    83  func (m *impl) Full() bool {
    84  	return m.full
    85  }
    86  
    87  func (m *impl) Empty() bool {
    88  	return len(m.blocks) == 0 && len(m.wantlist) == 0
    89  }
    90  
    91  func (m *impl) Wantlist() []Entry {
    92  	var out []Entry
    93  	for _, e := range m.wantlist {
    94  		out = append(out, e)
    95  	}
    96  	return out
    97  }
    98  
    99  func (m *impl) Blocks() []*blocks.Block {
   100  	bs := make([]*blocks.Block, 0, len(m.blocks))
   101  	for _, block := range m.blocks {
   102  		bs = append(bs, block)
   103  	}
   104  	return bs
   105  }
   106  
   107  func (m *impl) Cancel(k key.Key) {
   108  	delete(m.wantlist, k)
   109  	m.addEntry(k, 0, true)
   110  }
   111  
   112  func (m *impl) AddEntry(k key.Key, priority int) {
   113  	m.addEntry(k, priority, false)
   114  }
   115  
   116  func (m *impl) addEntry(k key.Key, priority int, cancel bool) {
   117  	e, exists := m.wantlist[k]
   118  	if exists {
   119  		e.Priority = priority
   120  		e.Cancel = cancel
   121  	} else {
   122  		m.wantlist[k] = Entry{
   123  			Entry: wantlist.Entry{
   124  				Key:      k,
   125  				Priority: priority,
   126  			},
   127  			Cancel: cancel,
   128  		}
   129  	}
   130  }
   131  
   132  func (m *impl) AddBlock(b *blocks.Block) {
   133  	m.blocks[b.Key()] = b
   134  }
   135  
   136  func FromNet(r io.Reader) (BitSwapMessage, error) {
   137  	pbr := ggio.NewDelimitedReader(r, inet.MessageSizeMax)
   138  
   139  	pb := new(pb.Message)
   140  	if err := pbr.ReadMsg(pb); err != nil {
   141  		return nil, err
   142  	}
   143  
   144  	m := newMessageFromProto(*pb)
   145  	return m, nil
   146  }
   147  
   148  func (m *impl) ToProto() *pb.Message {
   149  	pbm := new(pb.Message)
   150  	pbm.Wantlist = new(pb.Message_Wantlist)
   151  	for _, e := range m.wantlist {
   152  		pbm.Wantlist.Entries = append(pbm.Wantlist.Entries, &pb.Message_Wantlist_Entry{
   153  			Block:    proto.String(string(e.Key)),
   154  			Priority: proto.Int32(int32(e.Priority)),
   155  			Cancel:   proto.Bool(e.Cancel),
   156  		})
   157  	}
   158  	for _, b := range m.Blocks() {
   159  		pbm.Blocks = append(pbm.Blocks, b.Data)
   160  	}
   161  	return pbm
   162  }
   163  
   164  func (m *impl) ToNet(w io.Writer) error {
   165  	pbw := ggio.NewDelimitedWriter(w)
   166  
   167  	if err := pbw.WriteMsg(m.ToProto()); err != nil {
   168  		return err
   169  	}
   170  	return nil
   171  }
   172  
   173  func (m *impl) Loggable() map[string]interface{} {
   174  	var blocks []string
   175  	for _, v := range m.blocks {
   176  		blocks = append(blocks, v.Key().Pretty())
   177  	}
   178  	return map[string]interface{}{
   179  		"blocks": blocks,
   180  		"wants":  m.Wantlist(),
   181  	}
   182  }