This commit is contained in:
rui.zheng
2015-05-12 16:12:35 +08:00
parent d091e8c10b
commit df98baade5
+69 -67
View File
@@ -1,78 +1,78 @@
package gosocks5
import (
//"log"
"net"
"io"
"log"
"net"
"sync"
"time"
)
type Config struct {
Methods []uint8
SelectMethod func(methods ...uint8) uint8
Methods []uint8
SelectMethod func(methods ...uint8) uint8
MethodSelected func(method uint8, conn net.Conn) (net.Conn, error)
}
func defaultConfig() *Config {
return &Config{}
return &Config{}
}
type Conn struct {
c net.Conn
config *Config
method uint8
isClient bool
handshaked bool
handshakeMutex sync.Mutex
handshakeErr error
c net.Conn
config *Config
method uint8
isClient bool
handshaked bool
handshakeMutex sync.Mutex
handshakeErr error
}
func ClientConn(conn net.Conn, config *Config) *Conn {
return &Conn{
c: conn,
config: config,
isClient: true,
}
return &Conn{
c: conn,
config: config,
isClient: true,
}
}
func ServerConn(conn net.Conn, config *Config) *Conn {
return &Conn{
c: conn,
config: config,
}
return &Conn{
c: conn,
config: config,
}
}
func (conn *Conn) Handleshake() error {
conn.handshakeMutex.Lock()
defer conn.handshakeMutex.Unlock()
if err := conn.handshakeErr; err != nil {
return err
}
if conn.handshaked {
return nil
}
if conn.isClient {
conn.handshakeErr = conn.clientHandshake()
} else {
conn.handshakeErr = conn.serverHandshake()
}
return conn.handshakeErr
conn.handshakeMutex.Lock()
defer conn.handshakeMutex.Unlock()
if err := conn.handshakeErr; err != nil {
return err
}
if conn.handshaked {
return nil
}
if conn.isClient {
conn.handshakeErr = conn.clientHandshake()
} else {
conn.handshakeErr = conn.serverHandshake()
}
return conn.handshakeErr
}
func (conn *Conn) clientHandshake() error {
if conn.config == nil {
conn.config = defaultConfig()
}
nm := len(conn.config.Methods)
if conn.config == nil {
conn.config = defaultConfig()
}
nm := len(conn.config.Methods)
if nm == 0 {
nm = 1
}
b := make([]byte, 2+nm)
b[0] = Ver5
b[1] = uint8(nm)
@@ -89,24 +89,25 @@ func (conn *Conn) clientHandshake() error {
if b[0] != Ver5 {
return ErrBadVersion
}
if conn.config.MethodSelected != nil {
c, err := conn.config.MethodSelected(b[1], conn.c)
if err != nil {
c, err := conn.config.MethodSelected(b[1], conn.c)
if err != nil {
return err
}
conn.c = c
}
conn.method = b[1]
conn.handshaked = true
log.Println("method:", conn.method)
conn.handshaked = true
return nil
}
func (conn *Conn) serverHandshake() error {
if conn.config == nil {
conn.config = defaultConfig()
}
if conn.config == nil {
conn.config = defaultConfig()
}
methods, err := ReadMethods(conn.c)
if err != nil {
return err
@@ -129,44 +130,45 @@ func (conn *Conn) serverHandshake() error {
conn.c = c
}
conn.method = method
log.Println("method:", method)
conn.handshaked = true
return nil
}
func (conn *Conn) Read(b []byte) (n int, err error) {
if err = conn.Handleshake(); err != nil {
return
}
return conn.c.Read(b)
if err = conn.Handleshake(); err != nil {
return
}
return conn.c.Read(b)
}
func (conn *Conn) Write(b []byte) (n int, err error) {
if err = conn.Handleshake(); err != nil {
return
}
return conn.c.Write(b)
if err = conn.Handleshake(); err != nil {
return
}
return conn.c.Write(b)
}
func (conn *Conn) Close() error {
return conn.c.Close()
return conn.c.Close()
}
func (conn *Conn) LocalAddr() net.Addr {
return conn.c.LocalAddr()
return conn.c.LocalAddr()
}
func (conn *Conn) RemoteAddr() net.Addr {
return conn.c.RemoteAddr()
return conn.c.RemoteAddr()
}
func (conn *Conn) SetDeadline(t time.Time) error {
return conn.c.SetDeadline(t)
return conn.c.SetDeadline(t)
}
func (conn *Conn) SetReadDeadline(t time.Time) error {
return conn.c.SetReadDeadline(t)
return conn.c.SetReadDeadline(t)
}
func (conn *Conn) SetWriteDeadline(t time.Time) error {
return conn.c.SetWriteDeadline(t)
}
return conn.c.SetWriteDeadline(t)
}