github.com/linapex/ethereum-dpos-chinese@v0.0.0-20190316121959-b78b3a4a1ece/swarm/storage/mock/mem/mem.go (about)

     1  
     2  //<developer>
     3  //    <name>linapex 曹一峰</name>
     4  //    <email>linapex@163.com</email>
     5  //    <wx>superexc</wx>
     6  //    <qqgroup>128148617</qqgroup>
     7  //    <url>https://jsq.ink</url>
     8  //    <role>pku engineer</role>
     9  //    <date>2019-03-16 12:09:49</date>
    10  //</624342682084315136>
    11  
    12  //
    13  //
    14  //
    15  //
    16  //
    17  //
    18  //
    19  //
    20  //
    21  //
    22  //
    23  //
    24  //
    25  //
    26  //
    27  
    28  //
    29  //
    30  //
    31  package mem
    32  
    33  import (
    34  	"archive/tar"
    35  	"bytes"
    36  	"encoding/json"
    37  	"io"
    38  	"io/ioutil"
    39  	"sync"
    40  
    41  	"github.com/ethereum/go-ethereum/common"
    42  	"github.com/ethereum/go-ethereum/swarm/storage/mock"
    43  )
    44  
    45  //
    46  //
    47  type GlobalStore struct {
    48  	nodes map[string]map[common.Address]struct{}
    49  	data  map[string][]byte
    50  	mu    sync.Mutex
    51  }
    52  
    53  //
    54  func NewGlobalStore() *GlobalStore {
    55  	return &GlobalStore{
    56  		nodes: make(map[string]map[common.Address]struct{}),
    57  		data:  make(map[string][]byte),
    58  	}
    59  }
    60  
    61  //
    62  //
    63  func (s *GlobalStore) NewNodeStore(addr common.Address) *mock.NodeStore {
    64  	return mock.NewNodeStore(addr, s)
    65  }
    66  
    67  //
    68  //
    69  func (s *GlobalStore) Get(addr common.Address, key []byte) (data []byte, err error) {
    70  	s.mu.Lock()
    71  	defer s.mu.Unlock()
    72  
    73  	if _, ok := s.nodes[string(key)][addr]; !ok {
    74  		return nil, mock.ErrNotFound
    75  	}
    76  
    77  	data, ok := s.data[string(key)]
    78  	if !ok {
    79  		return nil, mock.ErrNotFound
    80  	}
    81  	return data, nil
    82  }
    83  
    84  //
    85  func (s *GlobalStore) Put(addr common.Address, key []byte, data []byte) error {
    86  	s.mu.Lock()
    87  	defer s.mu.Unlock()
    88  
    89  	if _, ok := s.nodes[string(key)]; !ok {
    90  		s.nodes[string(key)] = make(map[common.Address]struct{})
    91  	}
    92  	s.nodes[string(key)][addr] = struct{}{}
    93  	s.data[string(key)] = data
    94  	return nil
    95  }
    96  
    97  //
    98  func (s *GlobalStore) HasKey(addr common.Address, key []byte) bool {
    99  	s.mu.Lock()
   100  	defer s.mu.Unlock()
   101  
   102  	_, ok := s.nodes[string(key)][addr]
   103  	return ok
   104  }
   105  
   106  //
   107  //
   108  func (s *GlobalStore) Import(r io.Reader) (n int, err error) {
   109  	s.mu.Lock()
   110  	defer s.mu.Unlock()
   111  
   112  	tr := tar.NewReader(r)
   113  
   114  	for {
   115  		hdr, err := tr.Next()
   116  		if err != nil {
   117  			if err == io.EOF {
   118  				break
   119  			}
   120  			return n, err
   121  		}
   122  
   123  		data, err := ioutil.ReadAll(tr)
   124  		if err != nil {
   125  			return n, err
   126  		}
   127  
   128  		var c mock.ExportedChunk
   129  		if err = json.Unmarshal(data, &c); err != nil {
   130  			return n, err
   131  		}
   132  
   133  		addrs := make(map[common.Address]struct{})
   134  		for _, a := range c.Addrs {
   135  			addrs[a] = struct{}{}
   136  		}
   137  
   138  		key := string(common.Hex2Bytes(hdr.Name))
   139  		s.nodes[key] = addrs
   140  		s.data[key] = c.Data
   141  		n++
   142  	}
   143  	return n, err
   144  }
   145  
   146  //
   147  //
   148  func (s *GlobalStore) Export(w io.Writer) (n int, err error) {
   149  	s.mu.Lock()
   150  	defer s.mu.Unlock()
   151  
   152  	tw := tar.NewWriter(w)
   153  	defer tw.Close()
   154  
   155  	buf := bytes.NewBuffer(make([]byte, 0, 1024))
   156  	encoder := json.NewEncoder(buf)
   157  	for key, addrs := range s.nodes {
   158  		al := make([]common.Address, 0, len(addrs))
   159  		for a := range addrs {
   160  			al = append(al, a)
   161  		}
   162  
   163  		buf.Reset()
   164  		if err = encoder.Encode(mock.ExportedChunk{
   165  			Addrs: al,
   166  			Data:  s.data[key],
   167  		}); err != nil {
   168  			return n, err
   169  		}
   170  
   171  		data := buf.Bytes()
   172  		hdr := &tar.Header{
   173  			Name: common.Bytes2Hex([]byte(key)),
   174  			Mode: 0644,
   175  			Size: int64(len(data)),
   176  		}
   177  		if err := tw.WriteHeader(hdr); err != nil {
   178  			return n, err
   179  		}
   180  		if _, err := tw.Write(data); err != nil {
   181  			return n, err
   182  		}
   183  		n++
   184  	}
   185  	return n, err
   186  }
   187