minor changes

add some comments
fix some wrong format
change the method of forward finish check

.gitignore

@@ -1,3 +1,4 @@
 *.exe
 *~
 *bak
+*.swp

README.md

@@ -1,20 +1,20 @@
 socks-go
 =======
 
-A socks server implement by golang, support socks4/4a, socks5.
+A socks server implemented by golang, support socks4/4a, socks5.
 
 
 usage
 ====
 Usage example:
 
-    import "github.com/fangdingjun/socks"
+    import socks "github.com/fangdingjun/socks-go"
 
-    fucn main(){
+    func main(){
         l, _ := net.Listen("tcp", ":1080")
         for {
             conn, _ := l.Accept()
-            s := socks.SocksConn{conn}
-            go s.Serve()
+            s := socks.SocksConn{ClientConn: conn, Dial: nil} // Dial is a function which dial to the upstream server
+            go s.Serve() // serve the socks request
         }
     }

examples/server/.gitignore

@@ -0,0 +1 @@
+server

socks.go

@@ -45,7 +45,7 @@ func (s *SocksConn) Serve() {
 		s5 := socks5Conn{clientConn: s.ClientConn, dial: dial}
 		s5.Serve()
 	default:
-		log.Printf("error version %s", buf[0])
+		log.Printf("error version %d", buf[0])
 		s.ClientConn.Close()
 	}
 }

socks4.go

@@ -16,6 +16,7 @@ type socks4Conn struct {
 
 func (s4 *socks4Conn) Serve() {
 	defer s4.Close()
+
 	if err := s4.processRequest(); err != nil {
 		log.Println(err)
 		return
@@ -26,17 +27,27 @@ func (s4 *socks4Conn) Close() {
 	if s4.clientConn != nil {
 		s4.clientConn.Close()
 	}
+
 	if s4.serverConn != nil {
 		s4.serverConn.Close()
 	}
 }
 
 func (s4 *socks4Conn) forward() {
+
+	c := make(chan int, 2)
+
 	go func() {
 		io.Copy(s4.clientConn, s4.serverConn)
+		c <- 1
 	}()
 
+	go func() {
 		io.Copy(s4.serverConn, s4.clientConn)
+		c <- 1
+	}()
+
+	<-c
 }
 
 func (s4 *socks4Conn) processRequest() error {
@@ -44,18 +55,22 @@ func (s4 *socks4Conn) processRequest() error {
 	// process command and target here
 
 	buf := make([]byte, 128)
+
+	// read header
 	n, err := io.ReadAtLeast(s4.clientConn, buf, 8)
 	if err != nil {
 		return err
 	}
 
-	// only support connect
+	// command only support connect
 	if buf[0] != cmdConnect {
 		return fmt.Errorf("error command %s", buf[0])
 	}
 
+	// get port
 	port := binary.BigEndian.Uint16(buf[1:3])
 
+	// get ip
 	ip := net.IP(buf[3:7])
 
 	// NULL-terminated user string
@@ -90,7 +105,7 @@ func (s4 *socks4Conn) processRequest() error {
 	// if dial to target failed, user will receive connection reset
 	s4.clientConn.Write([]byte{0x00, 0x5a, 0x01, 0x02, 0x00, 0x00, 0x00, 0x00})
 
-	log.Printf("connecting to %s", target)
+	//log.Printf("connecting to %s\r\n", target)
 
 	// connect to the target
 	s4.serverConn, err = s4.dial("tcp", target)

socks5.go

@@ -18,6 +18,7 @@ type socks5Conn struct {
 
 func (s5 *socks5Conn) Serve() {
 	defer s5.Close()
+
 	if err := s5.handshake(); err != nil {
 		log.Println(err)
 		return
@@ -34,6 +35,8 @@ func (s5 *socks5Conn) handshake() error {
 	// only process auth methods here
 
 	buf := make([]byte, 258)
+
+	// read auth methods
 	n, err := io.ReadAtLeast(s5.clientConn, buf, 1)
 	if err != nil {
 		return err
@@ -41,6 +44,7 @@ func (s5 *socks5Conn) handshake() error {
 
 	l := int(buf[0]) + 1
 	if n < l {
+		// read remains data
 		_, err := io.ReadFull(s5.clientConn, buf[n:l])
 		if err != nil {
 			return err
@@ -55,22 +59,26 @@ func (s5 *socks5Conn) handshake() error {
 
 func (s5 *socks5Conn) processRequest() error {
 	buf := make([]byte, 258)
+
+	// read header
 	n, err := io.ReadAtLeast(s5.clientConn, buf, 10)
 	if err != nil {
 		return err
 	}
+
 	if buf[0] != socks5Version {
-		return fmt.Errorf("error version %s", buf[0])
+		return fmt.Errorf("error version %d", buf[0])
 	}
 
-	// only support connect
+	// command only support connect
 	if buf[1] != cmdConnect {
 		return fmt.Errorf("unsupported command %s", buf[1])
 	}
 
-	hlen := 0
-	host := ""
-	msglen := 0
+	hlen := 0   // target address length
+	host := ""  // target address
+	msglen := 0 // header length
+
 	switch buf[3] {
 	case addrTypeIPv4:
 		hlen = 4
@@ -83,6 +91,7 @@ func (s5 *socks5Conn) processRequest() error {
 	msglen = 6 + hlen
 
 	if n < msglen {
+		// read remains header
 		_, err := io.ReadFull(s5.clientConn, buf[n:msglen])
 		if err != nil {
 			return err
@@ -97,15 +106,17 @@ func (s5 *socks5Conn) processRequest() error {
 		host = net.IP(addr).String()
 	}
 
+	// get target port
 	port := binary.BigEndian.Uint16(buf[msglen-2 : msglen])
 
+	// target address
 	target := net.JoinHostPort(host, fmt.Sprintf("%d", port))
 
 	// reply user with connect success
 	// if dial to target failed, user will receive connection reset
 	s5.clientConn.Write([]byte{0x05, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01})
 
-	log.Printf("connecing to %s", target)
+	//log.Printf("connecing to %s\r\n", target)
 
 	// connect to the target
 	s5.serverConn, err = s5.dial("tcp", target)
@@ -120,8 +131,20 @@ func (s5 *socks5Conn) processRequest() error {
 }
 
 func (s5 *socks5Conn) forward() {
-	go io.Copy(s5.clientConn, s5.serverConn)
+
+	c := make(chan int, 2)
+
+	go func() {
+		io.Copy(s5.clientConn, s5.serverConn)
+		c <- 1
+	}()
+
+	go func() {
 		io.Copy(s5.serverConn, s5.clientConn)
+		c < 1
+	}()
+
+	<-c
 }
 
 func (s5 *socks5Conn) Close() {