pubkey-serial.go raw

   1  //go:build !(js && wasm)
   2  
   3  package database
   4  
   5  import (
   6  	"bytes"
   7  	"errors"
   8  
   9  	"github.com/dgraph-io/badger/v4"
  10  	"next.orly.dev/pkg/lol/chk"
  11  	"next.orly.dev/pkg/database/indexes"
  12  	"next.orly.dev/pkg/database/indexes/types"
  13  	"next.orly.dev/pkg/nostr/encoders/hex"
  14  )
  15  
  16  // GetOrCreatePubkeySerial returns the serial for a pubkey, creating one if it doesn't exist.
  17  // The pubkey parameter should be 32 bytes (schnorr public key).
  18  // This function is thread-safe and uses transactions to ensure atomicity.
  19  func (d *D) GetOrCreatePubkeySerial(pubkey []byte) (ser *types.Uint40, err error) {
  20  	if len(pubkey) != 32 {
  21  		err = errors.New("pubkey must be 32 bytes")
  22  		return
  23  	}
  24  
  25  	// Create pubkey hash
  26  	pubHash := new(types.PubHash)
  27  	if err = pubHash.FromPubkey(pubkey); chk.E(err) {
  28  		return
  29  	}
  30  
  31  	// First, try to get existing serial (separate transaction for read)
  32  	var existingSer *types.Uint40
  33  	existingSer, err = d.GetPubkeySerial(pubkey)
  34  	if err == nil && existingSer != nil {
  35  		// Serial already exists
  36  		ser = existingSer
  37  		return ser, nil
  38  	}
  39  
  40  	// Serial doesn't exist, create a new one
  41  	var serial uint64
  42  	if serial, err = d.pubkeySeq.Next(); chk.E(err) {
  43  		return
  44  	}
  45  
  46  	ser = new(types.Uint40)
  47  	if err = ser.Set(serial); chk.E(err) {
  48  		return
  49  	}
  50  
  51  	// Store both mappings in a transaction
  52  	err = d.Update(func(txn *badger.Txn) error {
  53  		// Double-check that the serial wasn't created by another goroutine
  54  		// while we were getting the sequence number
  55  		prefixBuf := new(bytes.Buffer)
  56  		prefixBuf.Write([]byte(indexes.PubkeySerialPrefix))
  57  		if terr := pubHash.MarshalWrite(prefixBuf); chk.E(terr) {
  58  			return terr
  59  		}
  60  		searchPrefix := prefixBuf.Bytes()
  61  
  62  		opts := badger.DefaultIteratorOptions
  63  		opts.PrefetchValues = false
  64  		opts.Prefix = searchPrefix
  65  		it := txn.NewIterator(opts)
  66  		it.Seek(searchPrefix)
  67  		if it.Valid() {
  68  			// Another goroutine created it, extract and return that serial
  69  			key := it.Item().KeyCopy(nil)
  70  			it.Close()
  71  			if len(key) == 16 {
  72  				serialBytes := key[11:16]
  73  				serialBuf := bytes.NewReader(serialBytes)
  74  				existSer := new(types.Uint40)
  75  				if terr := existSer.UnmarshalRead(serialBuf); terr == nil {
  76  					ser = existSer
  77  					return nil // Don't write, just return the existing serial
  78  				}
  79  			}
  80  		}
  81  		it.Close()
  82  
  83  		// Store pubkey hash -> serial mapping
  84  		keyBuf := new(bytes.Buffer)
  85  		if terr := indexes.PubkeySerialEnc(pubHash, ser).MarshalWrite(keyBuf); chk.E(terr) {
  86  			return terr
  87  		}
  88  		fullKey := make([]byte, len(keyBuf.Bytes()))
  89  		copy(fullKey, keyBuf.Bytes())
  90  		// DEBUG: log the key being written
  91  		if len(fullKey) > 0 {
  92  			// log.T.F("Writing PubkeySerial: key=%s (len=%d), prefix=%s", hex.Enc(fullKey), len(fullKey), string(fullKey[:3]))
  93  		}
  94  		if terr := txn.Set(fullKey, nil); chk.E(terr) {
  95  			return terr
  96  		}
  97  
  98  		// Store serial -> full pubkey mapping (pubkey stored as value)
  99  		keyBuf.Reset()
 100  		if terr := indexes.SerialPubkeyEnc(ser).MarshalWrite(keyBuf); chk.E(terr) {
 101  			return terr
 102  		}
 103  		if terr := txn.Set(keyBuf.Bytes(), pubkey); chk.E(terr) {
 104  			return terr
 105  		}
 106  
 107  		return nil
 108  	})
 109  
 110  	return
 111  }
 112  
 113  // GetPubkeySerial returns the serial for a pubkey if it exists.
 114  // Returns an error if the pubkey doesn't have a serial yet.
 115  func (d *D) GetPubkeySerial(pubkey []byte) (ser *types.Uint40, err error) {
 116  	if len(pubkey) != 32 {
 117  		err = errors.New("pubkey must be 32 bytes")
 118  		return
 119  	}
 120  
 121  	// Create pubkey hash
 122  	pubHash := new(types.PubHash)
 123  	if err = pubHash.FromPubkey(pubkey); chk.E(err) {
 124  		return
 125  	}
 126  
 127  	// Build search key with just prefix + pubkey hash (no serial)
 128  	prefixBuf := new(bytes.Buffer)
 129  	prefixBuf.Write([]byte(indexes.PubkeySerialPrefix)) // 3 bytes
 130  	if err = pubHash.MarshalWrite(prefixBuf); chk.E(err) {
 131  		return
 132  	}
 133  	searchPrefix := prefixBuf.Bytes() // Should be 11 bytes: 3 (prefix) + 8 (pubkey hash)
 134  
 135  	ser = new(types.Uint40)
 136  	err = d.View(func(txn *badger.Txn) error {
 137  		opts := badger.DefaultIteratorOptions
 138  		opts.PrefetchValues = false // We only need the key
 139  		it := txn.NewIterator(opts)
 140  		defer it.Close()
 141  
 142  		// Seek to the prefix and check if we found a matching key
 143  		it.Seek(searchPrefix)
 144  		if !it.ValidForPrefix(searchPrefix) {
 145  			return errors.New("pubkey serial not found")
 146  		}
 147  
 148  		// Extract serial from key (last 5 bytes)
 149  		// Key format: prefix(3) + pubkey_hash(8) + serial(5) = 16 bytes
 150  		key := it.Item().KeyCopy(nil)
 151  		if len(key) != 16 {
 152  			return errors.New("invalid key length for pubkey serial")
 153  		}
 154  
 155  		// Verify the prefix matches
 156  		if !bytes.HasPrefix(key, searchPrefix) {
 157  			return errors.New("key prefix mismatch")
 158  		}
 159  
 160  		serialBytes := key[11:16] // Extract last 5 bytes (the serial)
 161  
 162  		// Decode serial
 163  		serialBuf := bytes.NewReader(serialBytes)
 164  		if err := ser.UnmarshalRead(serialBuf); chk.E(err) {
 165  			return err
 166  		}
 167  
 168  		return nil
 169  	})
 170  
 171  	return
 172  }
 173  
 174  // GetPubkeyBySerial returns the full 32-byte pubkey for a given serial.
 175  func (d *D) GetPubkeyBySerial(ser *types.Uint40) (pubkey []byte, err error) {
 176  	keyBuf := new(bytes.Buffer)
 177  	if err = indexes.SerialPubkeyEnc(ser).MarshalWrite(keyBuf); chk.E(err) {
 178  		return
 179  	}
 180  
 181  	err = d.View(func(txn *badger.Txn) error {
 182  		item, gerr := txn.Get(keyBuf.Bytes())
 183  		if chk.E(gerr) {
 184  			return gerr
 185  		}
 186  
 187  		return item.Value(func(val []byte) error {
 188  			pubkey = make([]byte, len(val))
 189  			copy(pubkey, val)
 190  			return nil
 191  		})
 192  	})
 193  
 194  	if err != nil {
 195  		err = errors.New("pubkey not found for serial: " + hex.Enc([]byte{byte(ser.Get())}))
 196  	}
 197  
 198  	return
 199  }
 200