github.com/mit-dci/lit@v0.0.0-20221102210550-8c3d3b49f2ce/db/lnbolt/peerdb.go (about)

     1  package lnbolt
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"github.com/boltdb/bolt"
     7  	"github.com/mit-dci/lit/lncore"
     8  	"github.com/mit-dci/lit/lnutil"
     9  	"github.com/mit-dci/lit/logging"
    10  )
    11  
    12  var (
    13  	peersLabel    = []byte(`peers`)
    14  	peerMetaLabel = []byte(`peersmeta`)
    15  	pdbbuckets    = [][]byte{
    16  		peersLabel,
    17  		peerMetaLabel,
    18  	}
    19  
    20  	peerIdxLast = []byte(`lastpeeridx`)
    21  )
    22  
    23  type peerboltdb struct {
    24  	db *bolt.DB
    25  }
    26  
    27  func (pdb *peerboltdb) init() error {
    28  	err := pdb.db.Update(func(tx *bolt.Tx) error {
    29  		for _, n := range pdbbuckets {
    30  			_, err := tx.CreateBucketIfNotExists(n)
    31  			if err != nil {
    32  				return err
    33  			}
    34  		}
    35  		return nil
    36  	})
    37  	if err != nil {
    38  		return err
    39  	}
    40  	return nil
    41  }
    42  
    43  func (pdb *peerboltdb) GetPeerAddrs() ([]lncore.LnAddr, error) {
    44  
    45  	addrs := make([]lncore.LnAddr, 0)
    46  
    47  	err := pdb.db.View(func(tx *bolt.Tx) error {
    48  		b := tx.Bucket(peersLabel)
    49  
    50  		// Iterate over all of the members of the bucket.
    51  		cur := b.Cursor()
    52  		atmp := make([]lncore.LnAddr, 0)
    53  		for {
    54  			k, _ := cur.Next()
    55  			if k == nil {
    56  				break
    57  			}
    58  			lnaddr, err := lncore.ParseLnAddr(string(k))
    59  			if err != nil {
    60  				logging.Warnf("lnbolt/peerdb: found invalid key in DB as lnaddr: %s (error: %s)", string(k), err.Error())
    61  				continue
    62  			}
    63  
    64  			atmp = append(atmp, lnaddr)
    65  		}
    66  
    67  		// Now that we have the final array return it.
    68  		addrs = atmp
    69  		return nil
    70  	})
    71  	if err != nil {
    72  		return nil, err
    73  	}
    74  
    75  	return addrs, nil
    76  }
    77  
    78  func (pdb *peerboltdb) GetPeerInfo(addr lncore.LnAddr) (*lncore.PeerInfo, error) {
    79  
    80  	var raw []byte
    81  	var err error
    82  
    83  	if pdb.db == nil {
    84  		logging.Warnf("PDB.db is nil!")
    85  		return nil, fmt.Errorf("PDB.db is nil")
    86  	}
    87  
    88  	// Just get the raw data from the DB.
    89  	err = pdb.db.View(func(tx *bolt.Tx) error {
    90  		b := tx.Bucket(peersLabel)
    91  		raw = b.Get([]byte(string(addr)))
    92  		return nil
    93  	})
    94  	if err != nil {
    95  		return nil, err
    96  	}
    97  
    98  	if raw == nil {
    99  		return nil, nil
   100  	}
   101  
   102  	var pi lncore.PeerInfo
   103  	err = json.Unmarshal(raw, &pi)
   104  	if err != nil {
   105  		return nil, err
   106  	}
   107  
   108  	return &pi, nil
   109  
   110  }
   111  
   112  func (pdb *peerboltdb) GetPeerInfos() (map[lncore.LnAddr]lncore.PeerInfo, error) {
   113  
   114  	var out map[lncore.LnAddr]lncore.PeerInfo
   115  	var err error
   116  
   117  	err = pdb.db.View(func(tx *bolt.Tx) error {
   118  		b := tx.Bucket(peersLabel)
   119  
   120  		// Iterate over everything.
   121  		cur := b.Cursor()
   122  		mtmp := map[lncore.LnAddr]lncore.PeerInfo{}
   123  		for {
   124  			k, v := cur.Next()
   125  			if k == nil {
   126  				break
   127  			}
   128  
   129  			var pi lncore.PeerInfo
   130  			err2 := json.Unmarshal(v, &pi) // TODO Move outside tx block.
   131  			if err2 != nil {
   132  				return err2
   133  			}
   134  
   135  			ka, err2 := lncore.ParseLnAddr(string(k))
   136  			if err2 != nil {
   137  				logging.Warnf("lnbolt/peerdb: found invalid key in DB as lnaddr: %s (error: %s)", string(k), err.Error())
   138  				continue
   139  			}
   140  			mtmp[ka] = pi
   141  
   142  		}
   143  
   144  		out = mtmp
   145  		return nil
   146  	})
   147  	if err != nil {
   148  		return nil, err
   149  	}
   150  
   151  	return out, nil
   152  
   153  }
   154  
   155  func (pdb *peerboltdb) AddPeer(addr lncore.LnAddr, pi lncore.PeerInfo) error {
   156  	return pdb.UpdatePeer(addr, &pi)
   157  }
   158  
   159  func (pdb *peerboltdb) UpdatePeer(addr lncore.LnAddr, pi *lncore.PeerInfo) error {
   160  
   161  	var err error
   162  
   163  	araw := []byte(addr)
   164  	piraw, err := json.Marshal(pi)
   165  	if err != nil {
   166  		return err
   167  	}
   168  
   169  	err = pdb.db.Update(func(tx *bolt.Tx) error {
   170  		b := tx.Bucket(peersLabel)
   171  
   172  		err2 := b.Put(araw, piraw)
   173  		if err2 != nil {
   174  			return err2
   175  		}
   176  
   177  		return nil
   178  	})
   179  	if err != nil {
   180  		return err
   181  	}
   182  
   183  	return nil
   184  
   185  }
   186  
   187  func (pdb *peerboltdb) DeletePeer(addr lncore.LnAddr) error {
   188  
   189  	var err error
   190  
   191  	araw := []byte(addr)
   192  	err = pdb.db.Update(func(tx *bolt.Tx) error {
   193  		b := tx.Bucket(peersLabel)
   194  
   195  		err2 := b.Delete(araw)
   196  		if err2 != nil {
   197  			return err2
   198  		}
   199  
   200  		return nil
   201  	})
   202  	if err != nil {
   203  		return err
   204  	}
   205  
   206  	return nil
   207  
   208  }
   209  
   210  func (pdb *peerboltdb) GetUniquePeerIdx() (uint32, error) {
   211  
   212  	var err error
   213  	var pidx uint32
   214  
   215  	err = pdb.db.Update(func(tx *bolt.Tx) error {
   216  		b := tx.Bucket(peerMetaLabel)
   217  
   218  		// Get the last unique peer idx, or create it.
   219  		v := b.Get(peerIdxLast)
   220  		if v == nil {
   221  			x := lnutil.U32tB(1)
   222  			err2 := b.Put(peerIdxLast, x)
   223  			if err2 != nil {
   224  				return err2
   225  			}
   226  			v = x
   227  		}
   228  		pidx = lnutil.BtU32(v)
   229  
   230  		// Increment it.
   231  		err2 := b.Put(peerIdxLast, lnutil.U32tB(pidx+1))
   232  		if err2 != nil {
   233  			return err2
   234  		}
   235  
   236  		return nil
   237  	})
   238  
   239  	if err != nil {
   240  		return 0, err
   241  	}
   242  	return pidx, nil
   243  }