stun-go/attribute.go at master · vanstee/stun-go · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
package stun

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
)

const (
	MappedAddressAttribute = 0x0001
)

type AttributeValue interface {
	Serialize() []byte
	Length() uint16
	String() string
}

type Attribute struct {
	Type   uint16
	Length uint16
	Value  AttributeValue
}

func NewAttribute(attributeType uint16, value AttributeValue) *Attribute {
	return &Attribute{
		Type:   attributeType,
		Length: value.Length(),
		Value:  value,
	}
}

func (attribute *Attribute) Serialize() []byte {
	buffer := new(bytes.Buffer)
	binary.Write(buffer, binary.BigEndian, attribute.Type)
	binary.Write(buffer, binary.BigEndian, attribute.Length)
	bytes := buffer.Bytes()

	bytes = append(bytes, attribute.Value.Serialize()...)

	return bytes
}

func ParseAttributes(rawAttributes []byte) ([]*Attribute, error) {
	buffer := bytes.NewBuffer(rawAttributes)
	attributes := []*Attribute{}

	for buffer.Len() > 0 {
		attribute := &Attribute{}
		binary.Read(buffer, binary.BigEndian, &attribute.Type)
		binary.Read(buffer, binary.BigEndian, &attribute.Length)

		rawValue := make([]byte, attribute.Length)
		binary.Read(buffer, binary.BigEndian, &rawValue)
		value, err := ParseAttributeValue(rawValue)
		if err != nil {
			return nil, err
		}

		attribute.Value = value

		attributes = append(attributes, attribute)
	}

	return attributes, nil
}

func ParseAttributeValue(rawValue []byte) (AttributeValue, error) {
	buffer := bytes.NewBuffer(rawValue)
	var attributeType uint16
	binary.Read(buffer, binary.BigEndian, &attributeType)

	switch attributeType {
	case MappedAddressAttribute:
		return ParseMappedAddress(rawValue)
	}

	return nil, errors.New("Attribute type is invalid")
}

func (attribute *Attribute) String() string {
	var buffer bytes.Buffer
	buffer.WriteString(fmt.Sprintf("Type: %d\n", attribute.Type))
	buffer.WriteString(fmt.Sprintf("Length: %d\n", attribute.Length))
	buffer.WriteString("Value:\n")
	buffer.WriteString(attribute.Value.String())
	return buffer.String()
}