From 5686c46ac1519672e18496ade67b7039363d66d8 Mon Sep 17 00:00:00 2001
From: Chen Yufei <cyfdecyf@gmail.com>
Date: Thu, 13 Dec 2012 20:54:09 +0800
Subject: [PATCH] Refactor handleConnection in server.go

---
 cmd/shadowsocks-local/local.go   |   4 +-
 cmd/shadowsocks-server/server.go | 134 +++++++++++++------------------
 shadowsocks/pipe.go              |   2 +-
 3 files changed, 61 insertions(+), 79 deletions(-)

diff --git a/cmd/shadowsocks-local/local.go b/cmd/shadowsocks-local/local.go
index a461d0c..77cce0d 100644
--- a/cmd/shadowsocks-local/local.go
+++ b/cmd/shadowsocks-local/local.go
@@ -47,7 +47,7 @@ func handleConnection(conn net.Conn, server string, encTbl *ss.EncryptTable) {
 			log.Println("unsurpported addr type")
 			break
 		}
-		debug.Println("connecting ", addr)
+		debug.Println("connecting", addr)
 		conn.Write([]byte{0x05, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x08, 0x43})
 
 		remote, err := ss.DialWithRawAddr(addrToSend, server, encTbl)
@@ -55,7 +55,7 @@ func handleConnection(conn net.Conn, server string, encTbl *ss.EncryptTable) {
 			hasError = true
 			break
 		}
-		c := make(chan int, 2)
+		c := make(chan byte, 2)
 		go ss.Pipe(conn, remote, c)
 		go ss.Pipe(remote, conn, c)
 		<-c // close the other connection whenever one connection is closed
diff --git a/cmd/shadowsocks-server/server.go b/cmd/shadowsocks-server/server.go
index 1c93d38..77a66c6 100644
--- a/cmd/shadowsocks-server/server.go
+++ b/cmd/shadowsocks-server/server.go
@@ -3,6 +3,7 @@ package main
 import (
 	"bytes"
 	"encoding/binary"
+	"errors"
 	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
 	"log"
 	"net"
@@ -13,89 +14,70 @@ import (
 
 var debug ss.DebugLog
 
+var errAddr = errors.New("addr type not supported")
+
 func handleConnection(conn *ss.Conn) {
-	debug.Printf("socks connect from %s\n", conn.RemoteAddr().String())
-	var err error = nil
-	var hasError = false
-	for {
-		var _ int
-		buf := make([]byte, 1)
-		_, err = conn.Read(buf)
-		if err != nil {
-			hasError = true
-			break
-		}
-		addrType := buf[0]
-		var addr string
-		var port int16
-		if addrType == 1 {
-			buf = make([]byte, 6)
-			_, err = conn.Read(buf)
-			if err != nil {
-				hasError = true
-				break
-			}
-			var addrIp net.IP = make(net.IP, 4)
-			copy(addrIp, buf[0:4])
-			addr = addrIp.String()
-			sb := bytes.NewBuffer(buf[4:6])
-			binary.Read(sb, binary.BigEndian, &port)
-		} else if addrType == 3 {
-			_, err = conn.Read(buf)
-			if err != nil {
-				hasError = true
-				break
-			}
-			addrLen := buf[0]
-			buf = make([]byte, addrLen+2)
-			_, err = conn.Read(buf)
-			if err != nil {
-				hasError = true
-				break
-			}
-			sb := bytes.NewBuffer(buf[0:addrLen])
-			addr = sb.String()
-			sb = bytes.NewBuffer(buf[addrLen : addrLen+2])
-			binary.Read(sb, binary.BigEndian, &port)
-		} else {
-			hasError = true
-			log.Println("unsurpported addr type")
-			break
-		}
-		debug.Println("connecting ", addr)
-		var remote net.Conn
-		remote, err = net.Dial("tcp", addr+":"+strconv.Itoa(int(port)))
-		if err != nil {
-			hasError = true
-			break
+	if debug {
+		// function arguments are always evaluated, so surround debug
+		// statement with if statement
+		debug.Printf("socks connect from %s\n", conn.RemoteAddr().String())
+	}
+	defer conn.Close()
+
+	var addr string
+	var port int16
+	var addrType byte
+	var remote net.Conn
+	var c chan byte
+	var err error
+
+	buf := make([]byte, 1)
+	if _, err = conn.Read(buf); err != nil {
+		goto onError
+	}
+	addrType = buf[0]
+	if addrType == 1 {
+		buf = make([]byte, 6)
+		if _, err = conn.Read(buf); err != nil {
+			goto onError
 		}
-		if err != nil {
-			hasError = true
-			break
+		sb := bytes.NewBuffer(buf[4:6])
+		binary.Read(sb, binary.BigEndian, &port)
+		addrIp := make(net.IP, 4)
+		copy(addrIp, buf[0:4])
+		addr = addrIp.String()
+	} else if addrType == 3 {
+		if _, err = conn.Read(buf); err != nil {
+			goto onError
 		}
-		c := make(chan int, 2)
-		go ss.Pipe(conn, remote, c)
-		go ss.Pipe(remote, conn, c)
-		<-c // close the other connection whenever one connection is closed
-		debug.Println("closing")
-		err = conn.Close()
-		err1 := remote.Close()
-		if err == nil {
-			err = err1
+		addrLen := buf[0]
+		buf = make([]byte, addrLen+2)
+		if _, err = conn.Read(buf); err != nil {
+			goto onError
 		}
-		break
+		sb := bytes.NewBuffer(buf[addrLen : addrLen+2])
+		binary.Read(sb, binary.BigEndian, &port)
+		addr = string(buf[0:addrLen])
+	} else {
+		log.Println("unsurpported addr type")
+		err = errAddr
+		goto onError
 	}
-	if err != nil || hasError {
-		if err != nil {
-			debug.Println("error:", err)
-		}
-		err = conn.Close()
-		if err != nil {
-			debug.Println("close:", err)
-		}
-		return
+	debug.Println("connecting", addr)
+	if remote, err = net.Dial("tcp", addr+":"+strconv.Itoa(int(port))); err != nil {
+		goto onError
 	}
+	defer remote.Close()
+	debug.Println("piping", addr)
+	c = make(chan byte, 2)
+	go ss.Pipe(conn, remote, c)
+	go ss.Pipe(remote, conn, c)
+	<-c // close the other connection whenever one connection is closed
+	debug.Println("closing", addr)
+	return
 
+onError:
+	debug.Println("error", addr, err)
 }
 
 // Add a encrypt table cache to save memory and startup time in case of many
diff --git a/shadowsocks/pipe.go b/shadowsocks/pipe.go
index c3c40c4..befdd38 100644
--- a/shadowsocks/pipe.go
+++ b/shadowsocks/pipe.go
@@ -4,7 +4,7 @@ import (
 	"net"
 )
 
-func Pipe(src, dst net.Conn, end chan int) {
+func Pipe(src, dst net.Conn, end chan byte) {
 	// Should not use io.Copy here.
 	// io.Copy will try to use the ReadFrom interface of TCPConn, but the src
 	// here is not a regular file, so sendfile is not applicable.