bt/krpc/message.go

// Copyright 2020~2023 xgfone
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package krpc

import (
	"bytes"
	"fmt"

	"github.com/xgfone/bt/bencode"
	"github.com/xgfone/bt/metainfo"
)

// Predefine some error code.
const (
	// BEP 5
	ErrorCodeGeneric       = 201
	ErrorCodeServer        = 202
	ErrorCodeProtocol      = 203
	ErrorCodeMethodUnknown = 204

	// BEP 44
	ErrorCodeMessageValueFieldTooBig       = 205
	ErrorCodeInvalidSignature              = 206
	ErrorCodeSaltFieldTooBig               = 207
	ErrorCodeCasHashMismatched             = 301
	ErrorCodeSequenceNumberLessThanCurrent = 302
)

// Message represents messages that nodes in the network send to each other
// as specified by the KRPC protocol, which are also referred to as the KRPC
// messages.
//
// The message is a dictonary that is "bencoded" (serialization & compression
// format adopted by the BitTorrent) and sent via the UDP connection to peers.
//
// There are three types of messages: QUERY, RESPONSE, ERROR.
//
// A KRPC message is a single dictionary with two keys common to every message
// and additional keys depending on the type of message. Every message has a key
// "t" with a string value representing a transaction ID. This transaction ID
// is generated by the querying node and is echoed in the response, so responses
// may be correlated with multiple queries to the same node. The transaction ID
// should be encoded as a short string of binary numbers, typically 2 characters
// are enough as they cover 2^16 outstanding queries. The other key contained
// in every KRPC message is "y" with a single character value describing
// the type of message. The value of the "y" key is one of "q" for query,
// "r" for response, or "e" for error.
type Message struct {
	// Transaction ID
	//
	// Required
	T string `bencode:"t"` // BEP 5

	// Message Type: "q" for QUERY, "r" for RESPONSE, "e" for ERROR
	//
	// Required
	Y string `bencode:"y"` // BEP 5

	// Query Method: one of "ping", "find_node", "get_peers", "announce_peer"
	//
	// Required only if "y" is equal to "q".
	Q string `bencode:"q,omitempty"` // BEP 5

	A QueryArg       `bencode:"a,omitempty"` // BEP 5: Only for the QUERY message
	R ResponseResult `bencode:"r,omitempty"` // BEP 5: Only for the RESPONSE message
	E Error          `bencode:"e,omitempty"` // BEP 5: Only for the ERROR message

	RO bool `bencode:"ro,omitempty"` // BEP 43: ReadOnly
}

// NewQueryMsg return a QUERY message.
func NewQueryMsg(tid, method string, arg QueryArg) Message {
	return Message{T: tid, Y: "q", Q: method, A: arg}
}

// NewResponseMsg return a RESPONSE message.
func NewResponseMsg(tid string, data ResponseResult) Message {
	return Message{T: tid, Y: "r", R: data}
}

// NewErrorMsg return a ERROR message.
func NewErrorMsg(tid string, code int, reason string) Message {
	return Message{T: tid, Y: "e", E: Error{Code: code, Reason: reason}}
}

// IsQuery reports whether the message is a QUERY message.
func (m Message) IsQuery() bool { return m.Y == "q" }

// IsResponse reports whether the message is a RESPONSE message.
func (m Message) IsResponse() bool { return m.Y == "r" }

// IsError reports whether the message is an ERROR message.
func (m Message) IsError() bool { return m.Y == "e" }

// RID returns the value named "id" in "r".
//
// Return the ZERO value instead if no "id".
func (m Message) RID() metainfo.Hash { return m.R.ID }

// QID returns the value named "id" in "a", that's, the query arguments.
//
// Return the ZERO value instead if no "id".
func (m Message) QID() metainfo.Hash { return m.A.ID }

// ID returns the QID or RID.
//
// Notice: it will panic if the the message is not a query or response.
func (m Message) ID() metainfo.Hash {
	switch m.Y {
	case "q":
		return m.QID()
	case "r":
		return m.RID()
	default:
		panic(fmt.Errorf("unknown message type '%s'", m.Y))
	}
}

// Error represents a response error.
type Error struct {
	Code   int    // BEP 5
	Reason string // BEP 5
}

// NewError returns a new Error.
func NewError(code int, reason string) Error {
	return Error{Code: code, Reason: reason}
}

func (e *Error) decode(vs []interface{}) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("panic: %v", r)
		}
	}()

	e.Code = int(vs[0].(int64))
	e.Reason = vs[1].(string)
	return
}

// UnmarshalBencode implements the interface bencode.Unmarshaler.
func (e *Error) UnmarshalBencode(b []byte) (err error) {
	var iface interface{}
	if err = bencode.NewDecoder(bytes.NewBuffer(b)).Decode(&iface); err != nil {
		return
	}

	switch v := iface.(type) {
	case []interface{}:
		err = e.decode(v)
	case string:
		e.Reason = v
	default:
		err = fmt.Errorf(`KRPC error bencode value has unexpected type: %T`, v)
	}

	return
}

// MarshalBencode implements the interface bencode.Marshaler.
func (e Error) MarshalBencode() (ret []byte, err error) {
	if e.Code == 0 && e.Reason == "" {
		return nil, nil
	}

	buf := bytes.NewBuffer(nil)
	buf.Grow(32)
	err = bencode.NewEncoder(buf).Encode([]interface{}{e.Code, e.Reason})
	if err == nil {
		ret = buf.Bytes()
	}
	return
}

func (e Error) Error() string {
	return fmt.Sprintf("KRPC error %d: %s", e.Code, e.Reason)
}

// Want represents the type of nodes that the request wants.
//
// BEP 32
type Want string

// Predefine some wants.
//
// BEP 32
const (
	WantNodes  Want = "n4"
	WantNodes6 Want = "n6"
)

// QueryArg represents the arguments used by the QUERY message.
type QueryArg struct {
	// ID is used to identify a querying node.
	ID metainfo.Hash `bencode:"id"` // BEP 5

	// Target is used to identify the node sought by the queryer.
	//
	// find_node
	Target metainfo.Hash `bencode:"target,omitempty"` // BEP 5

	// InfoHash is the infohash of the torrent file.
	//
	// get_peers, announce_peer
	InfoHash metainfo.Hash `bencode:"info_hash,omitempty"` // BEP 5

	// Port is the port on where sender is listening to allow others
	// to download the torrent.
	//
	// announce_peer
	Port uint16 `bencode:"port,omitempty"` // BEP 5

	// Token is the received one from an earlier "get_peers" query.
	//
	// announce_peer
	Token string `bencode:"token,omitempty"` // BEP 5

	// ImpliedPort is used by the senders to apparent DHT port
	// to improve NAT support.
	//
	// announce_peer
	ImpliedPort bool `bencode:"implied_port,omitempty"` // BEP 5

	// Wants is only used to represent to expect "nodes" for "n4"
	// or "nodes6" for "n6".
	//
	// Notice: It only governs the presence of the "nodes" and "nodes6"
	// parameters, not the interpretation of "values".
	//
	// find_node, get_peers
	Wants []Want `bencode:"want,omitempty"` // BEP 32
}

// ContainsWant reports whether the request contains the given Want.
func (a QueryArg) ContainsWant(w Want) bool {
	for _, want := range a.Wants {
		if want == w {
			return true
		}
	}
	return false
}

// GetPort returns the real port of the peer.
func (a QueryArg) GetPort(port int) uint16 {
	if a.ImpliedPort {
		return uint16(port)
	}
	return a.Port
}

// ResponseResult represents the results used by the RESPONSE message.
type ResponseResult struct {
	// ID is used to indentify the queried node, that's, the response node.
	ID metainfo.Hash `bencode:"id"` // BEP 5

	// Nodes is a string containing the compact node information for the list
	// of the ipv4 target node, or the K(8) closest good nodes in routing table
	// of the requested ipv4 target.
	//
	// find_node
	Nodes CompactIPv4Nodes `bencode:"nodes,omitempty"` // BEP 5

	// Nodes6 is a string containing the compact node information for the list
	// of the ipv6 target node, or the K(8) closest good nodes in routing table
	// of the requested ipv6 target.
	//
	// find_node
	Nodes6 CompactIPv6Nodes `bencode:"nodes6,omitempty"` // BEP 32

	// Token is used for future "announce_peer".
	//
	// get_peers
	Token string `bencode:"token,omitempty"` // BEP 5

	// Values is a list of the torrent peers.
	//
	// For each element, in general, it is a compact IP-address/port info and
	// may be decoded to Addr, for example,
	//
	//   addrs := make([]Addr, len(values))
	//   for i, v := range values {
	//       addrs[i].UnmarshalBinary([]byte(v))
	//   }
	//
	// get_peers
	Values []string `bencode:"values,omitempty"` // BEP 5
}