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 }