torrent-gateway/internal/tracker/bencode.go

package tracker

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"strconv"
)

// BencodeEncoder provides additional bencode utilities beyond the anacrolix library
type BencodeEncoder struct{}

// NewBencodeEncoder creates a new bencode encoder
func NewBencodeEncoder() *BencodeEncoder {
	return &BencodeEncoder{}
}

// EncodeResponse encodes a tracker response with proper bencode formatting
func (e *BencodeEncoder) EncodeResponse(resp *AnnounceResponse) ([]byte, error) {
	var buf bytes.Buffer

	// Start dictionary
	buf.WriteString("d")

	// Add fields in alphabetical order (bencode requirement)
	if resp.Complete >= 0 {
		buf.WriteString("8:complete")
		buf.WriteString(e.encodeInt(resp.Complete))
	}

	if resp.FailureReason != "" {
		buf.WriteString("14:failure reason")
		buf.WriteString(e.encodeString(resp.FailureReason))
	}

	if resp.Incomplete >= 0 {
		buf.WriteString("10:incomplete")
		buf.WriteString(e.encodeInt(resp.Incomplete))
	}

	if resp.Interval > 0 {
		buf.WriteString("8:interval")
		buf.WriteString(e.encodeInt(resp.Interval))
	}

	if resp.MinInterval > 0 {
		buf.WriteString("12:min interval")
		buf.WriteString(e.encodeInt(resp.MinInterval))
	}

	// Encode peers
	if resp.Peers != nil {
		buf.WriteString("5:peers")
		if peerBytes, ok := resp.Peers.([]byte); ok {
			// Compact format
			buf.WriteString(e.encodeBytes(peerBytes))
		} else if dictPeers, ok := resp.Peers.([]DictPeer); ok {
			// Dictionary format
			buf.WriteString("l") // Start list
			for _, peer := range dictPeers {
				buf.WriteString("d") // Start peer dict
				buf.WriteString("2:ip")
				buf.WriteString(e.encodeString(peer.IP))
				buf.WriteString("7:peer id")
				buf.WriteString(e.encodeString(peer.PeerID))
				buf.WriteString("4:port")
				buf.WriteString(e.encodeInt(peer.Port))
				buf.WriteString("e") // End peer dict
			}
			buf.WriteString("e") // End list
		}
	}

	if resp.TrackerID != "" {
		buf.WriteString("10:tracker id")
		buf.WriteString(e.encodeString(resp.TrackerID))
	}

	if resp.WarningMessage != "" {
		buf.WriteString("15:warning message")
		buf.WriteString(e.encodeString(resp.WarningMessage))
	}

	// End dictionary
	buf.WriteString("e")

	return buf.Bytes(), nil
}

// encodeString encodes a string in bencode format
func (e *BencodeEncoder) encodeString(s string) string {
	return fmt.Sprintf("%d:%s", len(s), s)
}

// encodeBytes encodes bytes in bencode format
func (e *BencodeEncoder) encodeBytes(b []byte) string {
	return fmt.Sprintf("%d:", len(b)) + string(b)
}

// encodeInt encodes an integer in bencode format
func (e *BencodeEncoder) encodeInt(i int) string {
	return fmt.Sprintf("i%de", i)
}

// ParseAnnounceQuery parses URL-encoded announce parameters with proper bencode handling
func ParseAnnounceQuery(query map[string][]string) (map[string]interface{}, error) {
	result := make(map[string]interface{})

	for key, values := range query {
		if len(values) == 0 {
			continue
		}

		value := values[0]
		switch key {
		case "info_hash", "peer_id":
			// These are binary data that may be URL-encoded
			result[key] = value
		case "port", "uploaded", "downloaded", "left", "numwant":
			if i, err := strconv.ParseInt(value, 10, 64); err == nil {
				result[key] = i
			}
		case "compact":
			result[key] = value == "1"
		default:
			result[key] = value
		}
	}

	return result, nil
}

// BencodeDecoder provides bencode decoding utilities
type BencodeDecoder struct {
	reader *bufio.Reader
}

// NewBencodeDecoder creates a new bencode decoder
func NewBencodeDecoder(r io.Reader) *BencodeDecoder {
	return &BencodeDecoder{
		reader: bufio.NewReader(r),
	}
}

// DecodeDictionary decodes a bencode dictionary
func (d *BencodeDecoder) DecodeDictionary() (map[string]interface{}, error) {
	// Read 'd' marker
	b, err := d.reader.ReadByte()
	if err != nil {
		return nil, err
	}
	if b != 'd' {
		return nil, fmt.Errorf("expected dictionary marker 'd', got %c", b)
	}

	dict := make(map[string]interface{})

	for {
		// Check for end marker
		b, err := d.reader.ReadByte()
		if err != nil {
			return nil, err
		}
		if b == 'e' {
			break
		}
		
		// Put byte back
		d.reader.UnreadByte()

		// Read key (always a string)
		key, err := d.decodeString()
		if err != nil {
			return nil, fmt.Errorf("error reading dictionary key: %w", err)
		}

		// Read value
		value, err := d.decodeValue()
		if err != nil {
			return nil, fmt.Errorf("error reading dictionary value for key %s: %w", key, err)
		}

		dict[key] = value
	}

	return dict, nil
}

// decodeValue decodes any bencode value
func (d *BencodeDecoder) decodeValue() (interface{}, error) {
	b, err := d.reader.ReadByte()
	if err != nil {
		return nil, err
	}

	switch {
	case b >= '0' && b <= '9':
		// String - put byte back and decode
		d.reader.UnreadByte()
		return d.decodeString()
	case b == 'i':
		// Integer
		return d.decodeInteger()
	case b == 'l':
		// List
		return d.decodeList()
	case b == 'd':
		// Dictionary - put byte back and decode
		d.reader.UnreadByte()
		return d.DecodeDictionary()
	default:
		return nil, fmt.Errorf("unexpected bencode marker: %c", b)
	}
}

// decodeString decodes a bencode string
func (d *BencodeDecoder) decodeString() (string, error) {
	// Read length
	var lengthBytes []byte
	for {
		b, err := d.reader.ReadByte()
		if err != nil {
			return "", err
		}
		if b == ':' {
			break
		}
		lengthBytes = append(lengthBytes, b)
	}

	length, err := strconv.Atoi(string(lengthBytes))
	if err != nil {
		return "", fmt.Errorf("invalid string length: %s", string(lengthBytes))
	}

	// Read string data
	data := make([]byte, length)
	_, err = io.ReadFull(d.reader, data)
	if err != nil {
		return "", err
	}

	return string(data), nil
}

// decodeInteger decodes a bencode integer
func (d *BencodeDecoder) decodeInteger() (int64, error) {
	var intBytes []byte
	for {
		b, err := d.reader.ReadByte()
		if err != nil {
			return 0, err
		}
		if b == 'e' {
			break
		}
		intBytes = append(intBytes, b)
	}

	return strconv.ParseInt(string(intBytes), 10, 64)
}

// decodeList decodes a bencode list
func (d *BencodeDecoder) decodeList() ([]interface{}, error) {
	var list []interface{}

	for {
		// Check for end marker
		b, err := d.reader.ReadByte()
		if err != nil {
			return nil, err
		}
		if b == 'e' {
			break
		}
		
		// Put byte back
		d.reader.UnreadByte()

		// Read value
		value, err := d.decodeValue()
		if err != nil {
			return nil, err
		}

		list = append(list, value)
	}

	return list, nil
}