torrent/bencode/encode.go

package bencode

import (
	"io"
	"math/big"
	"reflect"
	"runtime"
	"sort"
	"strconv"
	"sync"

	"github.com/anacrolix/missinggo"
)

func isEmptyValue(v reflect.Value) bool {
	return missinggo.IsEmptyValue(v)
}

type Encoder struct {
	w       io.Writer
	scratch [64]byte
}

func (e *Encoder) Encode(v interface{}) (err error) {
	if v == nil {
		return
	}
	defer func() {
		if e := recover(); e != nil {
			if _, ok := e.(runtime.Error); ok {
				panic(e)
			}
			var ok bool
			err, ok = e.(error)
			if !ok {
				panic(e)
			}
		}
	}()
	e.reflectValue(reflect.ValueOf(v))
	return nil
}

type stringValues []reflect.Value

func (sv stringValues) Len() int           { return len(sv) }
func (sv stringValues) Swap(i, j int)      { sv[i], sv[j] = sv[j], sv[i] }
func (sv stringValues) Less(i, j int) bool { return sv.get(i) < sv.get(j) }
func (sv stringValues) get(i int) string   { return sv[i].String() }

func (e *Encoder) write(s []byte) {
	_, err := e.w.Write(s)
	if err != nil {
		panic(err)
	}
}

func (e *Encoder) writeString(s string) {
	for s != "" {
		n := copy(e.scratch[:], s)
		s = s[n:]
		e.write(e.scratch[:n])
	}
}

func (e *Encoder) reflectString(s string) {
	e.writeStringPrefix(int64(len(s)))
	e.writeString(s)
}

func (e *Encoder) writeStringPrefix(l int64) {
	b := strconv.AppendInt(e.scratch[:0], l, 10)
	e.write(b)
	e.writeString(":")
}

func (e *Encoder) reflectByteSlice(s []byte) {
	e.writeStringPrefix(int64(len(s)))
	e.write(s)
}

// Returns true if the value implements Marshaler interface and marshaling was
// done successfully.
func (e *Encoder) reflectMarshaler(v reflect.Value) bool {
	if !v.Type().Implements(marshalerType) {
		if v.Kind() != reflect.Ptr && v.CanAddr() && v.Addr().Type().Implements(marshalerType) {
			v = v.Addr()
		} else {
			return false
		}
	}
	m := v.Interface().(Marshaler)
	data, err := m.MarshalBencode()
	if err != nil {
		panic(&MarshalerError{v.Type(), err})
	}
	e.write(data)
	return true
}

var bigIntType = reflect.TypeOf((*big.Int)(nil)).Elem()

func (e *Encoder) reflectValue(v reflect.Value) {
	if e.reflectMarshaler(v) {
		return
	}

	if v.Type() == bigIntType {
		e.writeString("i")
		bi := v.Interface().(big.Int)
		e.writeString(bi.String())
		e.writeString("e")
		return
	}

	switch v.Kind() {
	case reflect.Bool:
		if v.Bool() {
			e.writeString("i1e")
		} else {
			e.writeString("i0e")
		}
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		e.writeString("i")
		b := strconv.AppendInt(e.scratch[:0], v.Int(), 10)
		e.write(b)
		e.writeString("e")
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		e.writeString("i")
		b := strconv.AppendUint(e.scratch[:0], v.Uint(), 10)
		e.write(b)
		e.writeString("e")
	case reflect.String:
		e.reflectString(v.String())
	case reflect.Struct:
		e.writeString("d")
		for _, ef := range getEncodeFields(v.Type()) {
			fieldValue := ef.i(v)
			if !fieldValue.IsValid() {
				continue
			}
			if ef.omitEmpty && isEmptyValue(fieldValue) {
				continue
			}
			e.reflectString(ef.tag)
			e.reflectValue(fieldValue)
		}
		e.writeString("e")
	case reflect.Map:
		if v.Type().Key().Kind() != reflect.String {
			panic(&MarshalTypeError{v.Type()})
		}
		if v.IsNil() {
			e.writeString("de")
			break
		}
		e.writeString("d")
		sv := stringValues(v.MapKeys())
		sort.Sort(sv)
		for _, key := range sv {
			e.reflectString(key.String())
			e.reflectValue(v.MapIndex(key))
		}
		e.writeString("e")
	case reflect.Slice, reflect.Array:
		e.reflectSequence(v)
	case reflect.Interface:
		e.reflectValue(v.Elem())
	case reflect.Ptr:
		if v.IsNil() {
			v = reflect.Zero(v.Type().Elem())
		} else {
			v = v.Elem()
		}
		e.reflectValue(v)
	default:
		panic(&MarshalTypeError{v.Type()})
	}
}

func (e *Encoder) reflectSequence(v reflect.Value) {
	// Use bencode string-type
	if v.Type().Elem().Kind() == reflect.Uint8 {
		if v.Kind() != reflect.Slice {
			// Can't use []byte optimization
			if !v.CanAddr() {
				e.writeStringPrefix(int64(v.Len()))
				for i := 0; i < v.Len(); i++ {
					var b [1]byte
					b[0] = byte(v.Index(i).Uint())
					e.write(b[:])
				}
				return
			}
			v = v.Slice(0, v.Len())
		}
		s := v.Bytes()
		e.reflectByteSlice(s)
		return
	}
	if v.IsNil() {
		e.writeString("le")
		return
	}
	e.writeString("l")
	for i, n := 0, v.Len(); i < n; i++ {
		e.reflectValue(v.Index(i))
	}
	e.writeString("e")
}

type encodeField struct {
	i         func(v reflect.Value) reflect.Value
	tag       string
	omitEmpty bool
}

type encodeFieldsSortType []encodeField

func (ef encodeFieldsSortType) Len() int           { return len(ef) }
func (ef encodeFieldsSortType) Swap(i, j int)      { ef[i], ef[j] = ef[j], ef[i] }
func (ef encodeFieldsSortType) Less(i, j int) bool { return ef[i].tag < ef[j].tag }

var (
	typeCacheLock     sync.RWMutex
	encodeFieldsCache = make(map[reflect.Type][]encodeField)
)

func getEncodeFields(t reflect.Type) []encodeField {
	typeCacheLock.RLock()
	fs, ok := encodeFieldsCache[t]
	typeCacheLock.RUnlock()
	if ok {
		return fs
	}
	fs = makeEncodeFields(t)
	typeCacheLock.Lock()
	defer typeCacheLock.Unlock()
	encodeFieldsCache[t] = fs
	return fs
}

func makeEncodeFields(t reflect.Type) (fs []encodeField) {
	for _i, n := 0, t.NumField(); _i < n; _i++ {
		i := _i
		f := t.Field(i)
		if f.PkgPath != "" {
			continue
		}
		if f.Anonymous {
			t := f.Type
			if t.Kind() == reflect.Ptr {
				t = t.Elem()
			}
			anonEFs := makeEncodeFields(t)
			for aefi := range anonEFs {
				anonEF := anonEFs[aefi]
				bottomField := anonEF
				bottomField.i = func(v reflect.Value) reflect.Value {
					v = v.Field(i)
					if v.Kind() == reflect.Ptr {
						if v.IsNil() {
							// This will skip serializing this value.
							return reflect.Value{}
						}
						v = v.Elem()
					}
					return anonEF.i(v)
				}
				fs = append(fs, bottomField)
			}
			continue
		}
		var ef encodeField
		ef.i = func(v reflect.Value) reflect.Value {
			return v.Field(i)
		}
		ef.tag = f.Name

		tv := getTag(f.Tag)
		if tv.Ignore() {
			continue
		}
		if tv.Key() != "" {
			ef.tag = tv.Key()
		}
		ef.omitEmpty = tv.OmitEmpty()
		fs = append(fs, ef)
	}
	fss := encodeFieldsSortType(fs)
	sort.Sort(fss)
	return fs
}