socks5/negotiation.go

package socks5

import (
	"io"
)

// NegotiationRequest is the negotiation request packet
type NegotiationRequest struct {
	NMethods int
	Methods  []AuthMethod
}

// NewNegotiationRequest returns new negotiation request packet
func NewNegotiationRequest(methods []AuthMethod) *NegotiationRequest {
	return &NegotiationRequest{
		NMethods: len(methods),
		Methods:  methods,
	}
}

// NewNegotiationRequestFrom reads negotiation request packet from reader
func NewNegotiationRequestFrom(r io.Reader) (*NegotiationRequest, error) {
	buf := make([]byte, 2)
	if _, err := io.ReadFull(r, buf); err != nil {
		return nil, err
	}

	if buf[0] != Version {
		return nil, ErrVersion
	}

	methodsLen := int(buf[1])
	if methodsLen == 0 {
		return nil, ErrBadRequest
	}

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

	methods := make([]AuthMethod, methodsLen)
	for i := range methodsBytes {
		methods[i] = AuthMethod(methodsBytes[i])
	}

	req := &NegotiationRequest{
		NMethods: methodsLen,
		Methods:  methods,
	}

	return req, nil
}

// WriteTo writes negotiation request packet
func (s *NegotiationRequest) WriteTo(w io.Writer) (int64, error) {
	buf := []byte{Version, byte(s.NMethods)}
	for i := range s.Methods {
		buf = append(buf, byte(s.Methods[i]))
	}

	n, err := w.Write(buf)
	if err != nil {
		return int64(n), err
	}

	return int64(n), nil
}

// NegotiationReply is the negotiation reply packet
type NegotiationReply struct {
	Method AuthMethod
}

// NewNegotiationReply returns new negotiation reply packet
func NewNegotiationReply(method AuthMethod) *NegotiationReply {
	return &NegotiationReply{
		Method: method,
	}
}

// NewNegotiationReplyFrom reads negotiation reply packet from reader
func NewNegotiationReplyFrom(r io.Reader) (*NegotiationReply, error) {
	buf := make([]byte, 2)
	if _, err := io.ReadFull(r, buf); err != nil {
		return nil, err
	}

	if buf[0] != Version {
		return nil, ErrVersion
	}

	rep := &NegotiationReply{
		Method: AuthMethod(buf[1]),
	}

	return rep, nil
}

// WriteTo writes negotiation reply packet
func (s *NegotiationReply) WriteTo(w io.Writer) (int64, error) {
	n, err := w.Write([]byte{Version, byte(s.Method)})
	if err != nil {
		return int64(n), err
	}

	return int64(n), nil
}