encoding.go raw

   1  package negentropy
   2  
   3  import (
   4  	"encoding/binary"
   5  	"encoding/hex"
   6  	"errors"
   7  	"io"
   8  )
   9  
  10  // Errors
  11  var (
  12  	ErrInvalidVarInt  = errors.New("invalid varint encoding")
  13  	ErrUnexpectedEOF  = errors.New("unexpected end of data")
  14  	ErrInvalidBound   = errors.New("invalid bound encoding")
  15  	ErrInvalidMessage = errors.New("invalid negentropy message")
  16  )
  17  
  18  // EncodeVarInt encodes an unsigned integer as a big-endian variable-length integer.
  19  // Matches the C++ negentropy encoding: extract 7-bit groups from LSB,
  20  // reverse to big-endian order, set continuation bits on all but last byte.
  21  func EncodeVarInt(n uint64) []byte {
  22  	if n == 0 {
  23  		return []byte{0}
  24  	}
  25  
  26  	// Extract 7-bit groups from LSB
  27  	var buf [10]byte
  28  	i := 0
  29  	for n > 0 {
  30  		buf[i] = byte(n & 0x7F)
  31  		n >>= 7
  32  		i++
  33  	}
  34  
  35  	// Reverse to big-endian order and set continuation bits
  36  	result := make([]byte, i)
  37  	for j := 0; j < i; j++ {
  38  		result[j] = buf[i-1-j]
  39  	}
  40  	for j := 0; j < len(result)-1; j++ {
  41  		result[j] |= 0x80
  42  	}
  43  
  44  	return result
  45  }
  46  
  47  // DecodeVarInt decodes a big-endian variable-length integer from a reader.
  48  // Matches the C++ negentropy decoding: accumulate with res = (res << 7) | (byte & 0x7F).
  49  func DecodeVarInt(r io.ByteReader) (uint64, error) {
  50  	var n uint64
  51  
  52  	for i := 0; i < 10; i++ {
  53  		b, err := r.ReadByte()
  54  		if err != nil {
  55  			return 0, ErrUnexpectedEOF
  56  		}
  57  
  58  		n = (n << 7) | uint64(b&0x7F)
  59  		if b&0x80 == 0 {
  60  			return n, nil
  61  		}
  62  	}
  63  
  64  	return 0, ErrInvalidVarInt
  65  }
  66  
  67  // WriteVarInt writes a varint to a byte slice, returning bytes written.
  68  func WriteVarInt(buf []byte, n uint64) int {
  69  	encoded := EncodeVarInt(n)
  70  	copy(buf, encoded)
  71  	return len(encoded)
  72  }
  73  
  74  // Encoder handles encoding negentropy protocol messages.
  75  type Encoder struct {
  76  	buf           []byte
  77  	lastTimestamp int64
  78  }
  79  
  80  // NewEncoder creates a new encoder with the given initial capacity.
  81  func NewEncoder(capacity int) *Encoder {
  82  	return &Encoder{
  83  		buf:           make([]byte, 0, capacity),
  84  		lastTimestamp: 0,
  85  	}
  86  }
  87  
  88  // Reset resets the encoder for reuse.
  89  func (e *Encoder) Reset() {
  90  	e.buf = e.buf[:0]
  91  	e.lastTimestamp = 0
  92  }
  93  
  94  // Bytes returns the encoded bytes.
  95  func (e *Encoder) Bytes() []byte {
  96  	return e.buf
  97  }
  98  
  99  // WriteByte writes a single byte.
 100  func (e *Encoder) WriteByte(b byte) {
 101  	e.buf = append(e.buf, b)
 102  }
 103  
 104  // WriteBytes writes a byte slice.
 105  func (e *Encoder) WriteBytes(b []byte) {
 106  	e.buf = append(e.buf, b...)
 107  }
 108  
 109  // WriteVarInt writes a variable-length integer.
 110  func (e *Encoder) WriteVarInt(n uint64) {
 111  	e.buf = append(e.buf, EncodeVarInt(n)...)
 112  }
 113  
 114  // WriteTimestamp writes a timestamp as a delta from the last timestamp.
 115  func (e *Encoder) WriteTimestamp(ts int64) {
 116  	if ts == MaxTimestamp {
 117  		// Special case: max timestamp is encoded as 0
 118  		e.WriteVarInt(0)
 119  	} else {
 120  		delta := ts - e.lastTimestamp
 121  		if delta < 0 {
 122  			delta = 0
 123  		}
 124  		e.WriteVarInt(uint64(delta) + 1)
 125  		e.lastTimestamp = ts
 126  	}
 127  }
 128  
 129  // WriteBound writes a bound (timestamp + optional ID).
 130  func (e *Encoder) WriteBound(b Bound) {
 131  	e.WriteTimestamp(b.Timestamp)
 132  
 133  	if b.Timestamp != MaxTimestamp && b.ID != "" {
 134  		// Write ID length and hex-decoded ID
 135  		idBytes, err := hex.DecodeString(b.ID)
 136  		if err != nil {
 137  			e.WriteVarInt(0)
 138  			return
 139  		}
 140  		// Bound ID prefix can be 0-32 bytes per spec
 141  		if len(idBytes) > FullIDSize {
 142  			idBytes = idBytes[:FullIDSize]
 143  		}
 144  		e.WriteVarInt(uint64(len(idBytes)))
 145  		e.WriteBytes(idBytes)
 146  	} else {
 147  		e.WriteVarInt(0)
 148  	}
 149  }
 150  
 151  // WriteFingerprint writes a fingerprint.
 152  func (e *Encoder) WriteFingerprint(fp Fingerprint) {
 153  	e.WriteBytes(fp[:])
 154  }
 155  
 156  // WriteMode writes a mode byte.
 157  func (e *Encoder) WriteMode(m Mode) {
 158  	e.WriteByte(byte(m))
 159  }
 160  
 161  // Decoder handles decoding negentropy protocol messages.
 162  type Decoder struct {
 163  	data          []byte
 164  	pos           int
 165  	lastTimestamp int64
 166  }
 167  
 168  // NewDecoder creates a new decoder for the given data.
 169  func NewDecoder(data []byte) *Decoder {
 170  	return &Decoder{
 171  		data:          data,
 172  		pos:           0,
 173  		lastTimestamp: 0,
 174  	}
 175  }
 176  
 177  // ReadByte reads a single byte.
 178  func (d *Decoder) ReadByte() (byte, error) {
 179  	if d.pos >= len(d.data) {
 180  		return 0, ErrUnexpectedEOF
 181  	}
 182  	b := d.data[d.pos]
 183  	d.pos++
 184  	return b, nil
 185  }
 186  
 187  // ReadBytes reads n bytes.
 188  func (d *Decoder) ReadBytes(n int) ([]byte, error) {
 189  	if d.pos+n > len(d.data) {
 190  		return nil, ErrUnexpectedEOF
 191  	}
 192  	b := d.data[d.pos : d.pos+n]
 193  	d.pos += n
 194  	return b, nil
 195  }
 196  
 197  // ReadVarInt reads a variable-length integer.
 198  func (d *Decoder) ReadVarInt() (uint64, error) {
 199  	return DecodeVarInt(d)
 200  }
 201  
 202  // ReadTimestamp reads a timestamp (delta encoded).
 203  func (d *Decoder) ReadTimestamp() (int64, error) {
 204  	delta, err := d.ReadVarInt()
 205  	if err != nil {
 206  		return 0, err
 207  	}
 208  
 209  	if delta == 0 {
 210  		return MaxTimestamp, nil
 211  	}
 212  
 213  	ts := d.lastTimestamp + int64(delta-1)
 214  	d.lastTimestamp = ts
 215  	return ts, nil
 216  }
 217  
 218  // ReadBound reads a bound (timestamp + optional ID).
 219  func (d *Decoder) ReadBound() (Bound, error) {
 220  	ts, err := d.ReadTimestamp()
 221  	if err != nil {
 222  		return Bound{}, err
 223  	}
 224  
 225  	idLen, err := d.ReadVarInt()
 226  	if err != nil {
 227  		return Bound{}, err
 228  	}
 229  
 230  	var id string
 231  	if idLen > 0 {
 232  		idBytes, err := d.ReadBytes(int(idLen))
 233  		if err != nil {
 234  			return Bound{}, err
 235  		}
 236  		id = hex.EncodeToString(idBytes)
 237  	}
 238  
 239  	return Bound{Item{Timestamp: ts, ID: id}}, nil
 240  }
 241  
 242  // ReadFingerprint reads a fingerprint.
 243  func (d *Decoder) ReadFingerprint() (Fingerprint, error) {
 244  	data, err := d.ReadBytes(DefaultIDSize)
 245  	if err != nil {
 246  		return Fingerprint{}, err
 247  	}
 248  	var fp Fingerprint
 249  	copy(fp[:], data)
 250  	return fp, nil
 251  }
 252  
 253  // ReadMode reads a mode byte.
 254  func (d *Decoder) ReadMode() (Mode, error) {
 255  	b, err := d.ReadByte()
 256  	if err != nil {
 257  		return 0, err
 258  	}
 259  	return Mode(b), nil
 260  }
 261  
 262  // Remaining returns the number of unread bytes.
 263  func (d *Decoder) Remaining() int {
 264  	return len(d.data) - d.pos
 265  }
 266  
 267  // HasMore returns true if there are more bytes to read.
 268  func (d *Decoder) HasMore() bool {
 269  	return d.pos < len(d.data)
 270  }
 271  
 272  // encodeHex encodes bytes to hex string.
 273  func encodeHex(b []byte) string {
 274  	return hex.EncodeToString(b)
 275  }
 276  
 277  // decodeHex decodes hex string to bytes.
 278  func decodeHex(s string) ([]byte, error) {
 279  	return hex.DecodeString(s)
 280  }
 281  
 282  // EncodeMessage encodes raw bytes to hex for NIP-77 transport.
 283  func EncodeMessage(msg []byte) string {
 284  	return hex.EncodeToString(msg)
 285  }
 286  
 287  // DecodeMessage decodes hex-encoded NIP-77 message to raw bytes.
 288  func DecodeMessage(hexMsg string) ([]byte, error) {
 289  	return hex.DecodeString(hexMsg)
 290  }
 291  
 292  // VarIntSize returns the number of bytes needed to encode n as a varint.
 293  func VarIntSize(n uint64) int {
 294  	if n == 0 {
 295  		return 1
 296  	}
 297  	size := 0
 298  	for n > 0 {
 299  		size++
 300  		n >>= 7
 301  	}
 302  	return size
 303  }
 304  
 305  // BoundSize returns the approximate size of a bound when encoded.
 306  func BoundSize(b Bound) int {
 307  	size := VarIntSize(uint64(b.Timestamp)) + 1 // timestamp + id length byte
 308  	if b.ID != "" {
 309  		size += min(len(b.ID)/2, DefaultIDSize)
 310  	}
 311  	return size
 312  }
 313  
 314  func min(a, b int) int {
 315  	if a < b {
 316  		return a
 317  	}
 318  	return b
 319  }
 320  
 321  // Read implements io.ByteReader for binary.Read compatibility.
 322  func (d *Decoder) Read(p []byte) (int, error) {
 323  	if d.pos >= len(d.data) {
 324  		return 0, io.EOF
 325  	}
 326  	n := copy(p, d.data[d.pos:])
 327  	d.pos += n
 328  	return n, nil
 329  }
 330  
 331  // ReadUint64 reads a little-endian uint64.
 332  func (d *Decoder) ReadUint64() (uint64, error) {
 333  	data, err := d.ReadBytes(8)
 334  	if err != nil {
 335  		return 0, err
 336  	}
 337  	return binary.LittleEndian.Uint64(data), nil
 338  }
 339