84 lines
1.6 KiB
Go
84 lines
1.6 KiB
Go
package dht
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
|
|
"github.com/jackpal/bencode-go"
|
|
)
|
|
|
|
// Msg represents a KRPC message.
|
|
type Msg struct {
|
|
T string `bencode:"t"`
|
|
Y string `bencode:"y"`
|
|
Q string `bencode:"q,omitempty"`
|
|
A map[string]interface{} `bencode:"a,omitempty"`
|
|
R map[string]interface{} `bencode:"r,omitempty"`
|
|
E []interface{} `bencode:"e,omitempty"`
|
|
}
|
|
|
|
// EncodeMsg marshals a KRPC message to bencode.
|
|
func EncodeMsg(msg Msg) ([]byte, error) {
|
|
m := make(map[string]interface{})
|
|
m["t"] = msg.T
|
|
m["y"] = msg.Y
|
|
if msg.Q != "" {
|
|
m["q"] = msg.Q
|
|
}
|
|
if msg.A != nil {
|
|
m["a"] = msg.A
|
|
}
|
|
if msg.R != nil {
|
|
m["r"] = msg.R
|
|
}
|
|
if msg.E != nil {
|
|
m["e"] = msg.E
|
|
}
|
|
|
|
var buf bytes.Buffer
|
|
err := bencode.Marshal(&buf, m)
|
|
return buf.Bytes(), err
|
|
}
|
|
|
|
// DecodeMsg unmarshals a KRPC message from bencode.
|
|
func DecodeMsg(data []byte) (Msg, error) {
|
|
val, err := bencode.Decode(bytes.NewReader(data))
|
|
if err != nil {
|
|
return Msg{}, err
|
|
}
|
|
m, ok := val.(map[string]interface{})
|
|
if !ok {
|
|
return Msg{}, fmt.Errorf("expected map[string]interface{}, got %T", val)
|
|
}
|
|
|
|
var msg Msg
|
|
if t, ok := m["t"].(string); ok {
|
|
msg.T = t
|
|
}
|
|
if y, ok := m["y"].(string); ok {
|
|
msg.Y = y
|
|
}
|
|
if q, ok := m["q"].(string); ok {
|
|
msg.Q = q
|
|
}
|
|
if a, ok := m["a"].(map[string]interface{}); ok {
|
|
msg.A = a
|
|
}
|
|
if r, ok := m["r"].(map[string]interface{}); ok {
|
|
msg.R = r
|
|
}
|
|
if e, ok := m["e"].([]interface{}); ok {
|
|
msg.E = e
|
|
}
|
|
return msg, nil
|
|
}
|
|
|
|
// NewQuery creates a KRPC query message.
|
|
func NewQuery(t string, q string, a map[string]interface{}) Msg {
|
|
return Msg{
|
|
T: t,
|
|
Y: "q",
|
|
Q: q,
|
|
A: a,
|
|
}
|
|
}
|