socks5/datagram.go
2021-07-10 11:50:31 +03:00

84 lines
1.5 KiB
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()
}