add Selector interface

This commit is contained in:
rui.zheng
2015-09-25 17:56:52 +08:00
parent 2d00801412
commit a77d04567a
+26 -29
View File
@@ -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
} }