socks5/datagram.go

package socks5

import (
	"bytes"
	"io"
)

// Datagram is the datagram packet
type Datagram struct {
	ATYP    ATYP
	DstAddr []byte
	DstPort []byte
	Data    []byte
}

// NewDatagram creates new datagram packet
func NewDatagram(atyp ATYP, host []byte, port []byte, data []byte) *Datagram {
	return &Datagram{
		ATYP:    atyp,
		DstAddr: host,
		DstPort: port,
		Data:    data,
	}
}

// NewDatagramFrom reads datagram packet from reader
func NewDatagramFrom(r io.Reader) (*Datagram, error) {
	buf := make([]byte, 4)
	if _, err := io.ReadFull(r, buf); err != nil {
		return nil, err
	}

	if !bytes.Equal([]byte{Reserved, Reserved}, buf[:2]) {
		return nil, ErrBadDatagram
	}

	if buf[2] != DatagramStandalone {
		return nil, ErrFragmentedDatagram
	}

	addr, err := ReadAddress(r, ATYP(buf[3]))
	if err != nil {
		return nil, err
	}

	port := make([]byte, 2)
	if _, err := io.ReadFull(r, port); err != nil {
		return nil, err
	}

	data, err := io.ReadAll(r)
	if err != nil {
		return nil, err
	}

	datagram := &Datagram{
		ATYP:    ATYP(buf[3]),
		DstAddr: addr,
		DstPort: port,
		Data:    data,
	}

	return datagram, nil
}

// WriteTo writes datagram packet
func (s *Datagram) WriteTo(w io.Writer) (int64, error) {
	n, err := w.Write(s.Bytes())
	if err != nil {
		return int64(n), err
	}

	return int64(n), nil
}

func (s *Datagram) Bytes() []byte {
	buf := bytes.NewBuffer([]byte{Reserved, Reserved, DatagramStandalone, byte(s.ATYP)})
	buf.Write(s.DstAddr)
	buf.Write(s.DstPort)
	buf.Write(s.Data)

	return buf.Bytes()
}