From d091e8c10be90c24dff8772e5229da3e4e25d413 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Tue, 12 May 2015 03:29:43 +0000
Subject: [PATCH] add socks5 conn

---
 conn.go   | 172 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 server.go |  70 ----------------------
 2 files changed, 172 insertions(+), 70 deletions(-)
 create mode 100644 conn.go
 delete mode 100644 server.go

diff --git a/conn.go b/conn.go
new file mode 100644
index 0000000..0260592
--- /dev/null
+++ b/conn.go
@@ -0,0 +1,172 @@
+package gosocks5
+
+import (
+	//"log"
+	"net"
+	"io"
+	"sync"
+	"time"
+)
+
+type Config struct {
+    Methods []uint8
+    SelectMethod   func(methods ...uint8) uint8
+	MethodSelected func(method uint8, conn net.Conn) (net.Conn, error)
+}
+
+func defaultConfig() *Config {
+    return &Config{}
+}
+
+type Conn struct {
+    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,
+    }
+}
+
+func ServerConn(conn net.Conn, config *Config) *Conn {
+    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
+}
+
+func (conn *Conn) clientHandshake() error {
+    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)
+	copy(b[2:], conn.config.Methods)
+
+	if _, err := conn.c.Write(b); err != nil {
+		return err
+	}
+
+	if _, err := io.ReadFull(conn.c, b[:2]); err != nil {
+		return err
+	}
+
+	if b[0] != Ver5 {
+		return ErrBadVersion
+	}
+	
+	if conn.config.MethodSelected != 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
+	return nil
+}
+
+func (conn *Conn) serverHandshake() error {
+    if conn.config == nil {
+        conn.config = defaultConfig()
+    }
+    
+	methods, err := ReadMethods(conn.c)
+	if err != nil {
+		return err
+	}
+
+	method := MethodNoAuth
+	if conn.config.SelectMethod != nil {
+		method = conn.config.SelectMethod(methods...)
+	}
+
+	if _, err := conn.c.Write([]byte{Ver5, method}); err != nil {
+		return err
+	}
+
+	if conn.config.MethodSelected != nil {
+		c, err := conn.config.MethodSelected(method, conn.c)
+		if err != nil {
+			return err
+		}
+		conn.c = c
+	}
+	conn.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)
+}
+
+func (conn *Conn) Write(b []byte) (n int, err error) {
+    if err = conn.Handleshake(); err != nil {
+        return
+    }
+    return conn.c.Write(b)
+}
+
+func (conn *Conn) Close() error {
+    return conn.c.Close()
+}
+
+func (conn *Conn) LocalAddr() net.Addr {
+    return conn.c.LocalAddr()
+}
+
+func (conn *Conn) RemoteAddr() net.Addr {
+    return conn.c.RemoteAddr()
+}
+
+func (conn *Conn) SetDeadline(t time.Time) error {
+    return conn.c.SetDeadline(t)
+}
+
+func (conn *Conn) SetReadDeadline(t time.Time) error {
+    return conn.c.SetReadDeadline(t)
+}
+
+func (conn *Conn) SetWriteDeadline(t time.Time) error {
+    return conn.c.SetWriteDeadline(t)
+}
\ No newline at end of file
diff --git a/server.go b/server.go
deleted file mode 100644
index 834093b..0000000
--- a/server.go
+++ /dev/null
@@ -1,70 +0,0 @@
-package gosocks5
-
-import (
-	"log"
-	"net"
-)
-
-type Server struct {
-	Addr string // TCP address to listen on
-
-	SelectMethod   func(methods ...uint8) uint8
-	MethodSelected func(method uint8, conn net.Conn) (net.Conn, error)
-	Handle         func(conn net.Conn)
-}
-
-func (s *Server) ListenAndServe() error {
-	addr, err := net.ResolveTCPAddr("tcp", s.Addr)
-	if err != nil {
-		return err
-	}
-
-	ln, err := net.ListenTCP("tcp", addr)
-	if err != nil {
-		return err
-	}
-	defer ln.Close()
-
-	for {
-		conn, err := ln.AcceptTCP()
-		if err != nil {
-			log.Println("accept:", err)
-			continue
-		}
-		//log.Println("accept", conn.RemoteAddr())
-		go s.handle(conn)
-	}
-}
-
-func (s *Server) handle(conn net.Conn) {
-	defer conn.Close()
-
-	methods, err := ReadMethods(conn)
-	if err != nil {
-		log.Println(err)
-		return
-	}
-
-	method := MethodNoAuth
-	if s.SelectMethod != nil {
-		method = s.SelectMethod(methods...)
-	}
-
-	if _, err := conn.Write([]byte{Ver5, method}); err != nil {
-		log.Println(err)
-		return
-	}
-
-	if s.MethodSelected != nil {
-		c, err := s.MethodSelected(method, conn)
-		if err != nil {
-			log.Println(err)
-			return
-		}
-		conn = c
-	}
-
-	if s.Handle != nil {
-		s.Handle(conn)
-	}
-}