mirror of
https://github.com/ginuerzh/gosocks5.git
synced 2024-08-11 17:54:31 +00:00
add Selector interface
This commit is contained in:
@@ -8,19 +8,18 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
type Config struct {
|
type Selector interface {
|
||||||
Methods []uint8
|
// return supported methods
|
||||||
SelectMethod func(methods ...uint8) uint8
|
Methods() []uint8
|
||||||
MethodSelected func(method uint8, conn net.Conn) (net.Conn, error)
|
// select method
|
||||||
}
|
Select(methods ...uint8) (method uint8)
|
||||||
|
// on method selected
|
||||||
func defaultConfig() *Config {
|
OnSelected(method uint8, conn net.Conn) (net.Conn, error)
|
||||||
return &Config{}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
type Conn struct {
|
type Conn struct {
|
||||||
c net.Conn
|
c net.Conn
|
||||||
config *Config
|
selector Selector
|
||||||
method uint8
|
method uint8
|
||||||
isClient bool
|
isClient bool
|
||||||
handshaked bool
|
handshaked bool
|
||||||
@@ -28,18 +27,18 @@ type Conn struct {
|
|||||||
handshakeErr error
|
handshakeErr error
|
||||||
}
|
}
|
||||||
|
|
||||||
func ClientConn(conn net.Conn, config *Config) *Conn {
|
func ClientConn(conn net.Conn, selector Selector) *Conn {
|
||||||
return &Conn{
|
return &Conn{
|
||||||
c: conn,
|
c: conn,
|
||||||
config: config,
|
selector: selector,
|
||||||
isClient: true,
|
isClient: true,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func ServerConn(conn net.Conn, config *Config) *Conn {
|
func ServerConn(conn net.Conn, selector Selector) *Conn {
|
||||||
return &Conn{
|
return &Conn{
|
||||||
c: conn,
|
c: conn,
|
||||||
config: config,
|
selector: selector,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -64,11 +63,13 @@ func (conn *Conn) Handleshake() error {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (conn *Conn) clientHandshake() error {
|
func (conn *Conn) clientHandshake() error {
|
||||||
if conn.config == nil {
|
var methods []uint8
|
||||||
conn.config = defaultConfig()
|
var nm int
|
||||||
}
|
|
||||||
|
|
||||||
nm := len(conn.config.Methods)
|
if conn.selector != nil {
|
||||||
|
methods = conn.selector.Methods()
|
||||||
|
}
|
||||||
|
nm = len(methods)
|
||||||
if nm == 0 {
|
if nm == 0 {
|
||||||
nm = 1
|
nm = 1
|
||||||
}
|
}
|
||||||
@@ -76,7 +77,7 @@ func (conn *Conn) clientHandshake() error {
|
|||||||
b := make([]byte, 2+nm)
|
b := make([]byte, 2+nm)
|
||||||
b[0] = Ver5
|
b[0] = Ver5
|
||||||
b[1] = uint8(nm)
|
b[1] = uint8(nm)
|
||||||
copy(b[2:], conn.config.Methods)
|
copy(b[2:], methods)
|
||||||
|
|
||||||
if _, err := conn.c.Write(b); err != nil {
|
if _, err := conn.c.Write(b); err != nil {
|
||||||
return err
|
return err
|
||||||
@@ -90,8 +91,8 @@ func (conn *Conn) clientHandshake() error {
|
|||||||
return ErrBadVersion
|
return ErrBadVersion
|
||||||
}
|
}
|
||||||
|
|
||||||
if conn.config.MethodSelected != nil {
|
if conn.selector != nil {
|
||||||
c, err := conn.config.MethodSelected(b[1], conn.c)
|
c, err := conn.selector.OnSelected(b[1], conn.c)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
@@ -104,26 +105,22 @@ func (conn *Conn) clientHandshake() error {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (conn *Conn) serverHandshake() error {
|
func (conn *Conn) serverHandshake() error {
|
||||||
if conn.config == nil {
|
|
||||||
conn.config = defaultConfig()
|
|
||||||
}
|
|
||||||
|
|
||||||
methods, err := ReadMethods(conn.c)
|
methods, err := ReadMethods(conn.c)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
method := MethodNoAuth
|
method := MethodNoAuth
|
||||||
if conn.config.SelectMethod != nil {
|
if conn.selector != nil {
|
||||||
method = conn.config.SelectMethod(methods...)
|
method = conn.selector.Select(methods...)
|
||||||
}
|
}
|
||||||
|
|
||||||
if _, err := conn.c.Write([]byte{Ver5, method}); err != nil {
|
if _, err := conn.c.Write([]byte{Ver5, method}); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
if conn.config.MethodSelected != nil {
|
if conn.selector != nil {
|
||||||
c, err := conn.config.MethodSelected(method, conn.c)
|
c, err := conn.selector.OnSelected(method, conn.c)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user