Extract getRequest in local.go

cmd/shadowsocks-local/local.go

@@ -1,83 +1,184 @@
 package main
 
 import (
+	"bytes"
+	"encoding/binary"
+	"errors"
 	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
+	"io"
 	"log"
 	"net"
 	"strconv"
 )
 
 var debug ss.DebugLog
 
-func handleConnection(conn net.Conn, server string, encTbl *ss.EncryptTable) {
-	debug.Printf("socks connect from %s\n", conn.RemoteAddr().String())
-	b := make([]byte, 262)
-	var err error = nil
-	var hasError = false
+var (
+	errAddr   = errors.New("socks addr type not supported")
+	errVer    = errors.New("socks version not supported")
+	errMethod = errors.New("socks only support 1 method now")
+	errAuth   = errors.New("socks authentication not required")
+	errCmd    = errors.New("socks command not supported")
+)
+
+func handShake(conn net.Conn) (err error) {
+	// version identification and method selection message in theory can have
+	// at most 256 methods, plus version and nmethod field in total 258 bytes
+	// the current rfc defines only 3 authentication methods (plus 2 reserved)
+
+	// currently, we only support version 5 and 1 method
+	buf := make([]byte, 3)
+	// make sure we get 3 bytes
+	if _, err = io.ReadFull(conn, buf); err != nil {
+		return
+	}
+	if buf[0] != 5 {
+		err = errVer
+	}
+	if buf[1] != 1 {
+		err = errMethod
+	}
+	if buf[2] != 0 {
+		err = errAuth
+	}
+	if err != nil {
+		// send "no acceptable methods" to client
+		conn.Write([]byte{5, 0xff})
+		return
+	}
+	_, err = conn.Write([]byte{5, 0})
+	return
+}
+
+func getRequest(conn net.Conn) (rawaddr []byte, extra []byte, host string, err error) {
+	const (
+		idVer   = 0
+		idCmd   = 1
+		idType  = 3 // address type index
+		idIP0   = 4 // ip addres start index
+		idDmLen = 4 // domain address length index
+		idDm0   = 5 // domain address start index
+
+		typeIP = 1 // type is ip address
+		typeDm = 3 // type is domain address
+
+		lenIP     = 3 + 1 + 4 + 2 // 3(ver+cmd+rsv) + 1addrType + 4ip + 2port
+		lenDmBase = 3 + 1 + 1 + 2 // 3 + 1addrType + 1addrLen + 2port, plus addrLen
+	)
+	// refer to getRequest in server.go for why set buffer size to 263
+	buf := make([]byte, 263, 263)
+	cur := 0 // current location in buf
+	reqLen := 0
+
 	for {
-		var _ int
-		buf := make([]byte, 4096)
-		_, err = conn.Read(b)
-		if err != nil {
-			hasError = true
-			break
+		ss.SetReadTimeout(conn)
+		var n int
+		// usually need to read only once
+		if n, err = conn.Read(buf[cur:]); err != nil {
+			// debug.Println("read request error:", err)
+			return
 		}
-		conn.Write([]byte{0x05, 0x00})
-		_, err = conn.Read(buf)
-		mode := buf[1]
-		if mode != 1 {
-			hasError = true
-			log.Println("mode != 1")
-			break
+		cur += n
+		if cur < idType+1 { // read till we get addr type
+			continue
 		}
-		var addr string
-		addrType := buf[3]
-		var addrToSend []byte
-		if addrType == 1 {
-			var addrIp net.IP = make(net.IP, 4)
-			copy(addrIp, buf[4:8])
-			addr = addrIp.String()
-			addrToSend = buf[3:10]
-		} else if addrType == 3 {
-			addrLen := buf[4]
-			addr = string(buf[5 : 5+addrLen])
-			addrToSend = buf[3 : 5+addrLen+2]
-		} else {
-			hasError = true
-			log.Println("unsurpported addr type")
-			break
+		// check version and cmd
+		if buf[idVer] != 5 {
+			err = errVer
+			return
 		}
-		debug.Println("connecting", addr)
-		conn.Write([]byte{0x05, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x08, 0x43})
-
-		remote, err := ss.DialWithRawAddr(addrToSend, server, encTbl)
-		if err != nil {
-			hasError = true
-			break
+		if buf[idCmd] != 1 {
+			err = errCmd
+			return
 		}
-		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")
-		err = conn.Close()
-		err1 := remote.Close()
-		if err == nil {
-			err = err1
+		// TODO following code is copied from server.go, fix code duplication?
+		if buf[idType] == typeIP {
+			if cur >= lenIP {
+				// debug.Println("ip request complete, cur:", cur)
+				reqLen = lenIP
+				break
+			}
+		} else if buf[idType] == typeDm {
+			if cur < idDmLen+1 { // read until we get address length byte
+				continue
+			}
+			if cur >= lenDmBase+int(buf[idDmLen]) {
+				// debug.Println("domain request complete, cur:", cur)
+				reqLen = lenDmBase + int(buf[idDmLen])
+				break
+			}
+		} else {
+			err = errAddr
+			return
 		}
-		break
+		// debug.Println("request not complete, cur:", cur)
 	}
-	if err != nil || hasError {
-		if err != nil {
-			log.Println("error:", err)
-		}
-		err = conn.Close()
-		if err != nil {
-			log.Println("close:", err)
+
+	rawaddr = buf[idType:reqLen]
+	if cur > reqLen {
+		extra = buf[reqLen:cur]
+		// debug.Println("extra:", string(extra))
+	}
+
+	if debug {
+		if buf[idType] == typeIP {
+			addrIp := make(net.IP, 4)
+			copy(addrIp, buf[idIP0:idIP0+4])
+			host = addrIp.String()
+		} else if buf[idType] == typeDm {
+			host = string(buf[idDm0 : idDm0+buf[idDmLen]])
 		}
+		var port int16
+		sb := bytes.NewBuffer(buf[reqLen-2 : reqLen])
+		binary.Read(sb, binary.BigEndian, &port)
+		host += ":" + strconv.Itoa(int(port))
+	}
+	return
+}
+
+func handleConnection(conn net.Conn, server string, encTbl *ss.EncryptTable) {
+	if debug {
+		debug.Printf("socks connect from %s\n", conn.RemoteAddr().String())
+	}
+	defer conn.Close()
+
+	var err error = nil
+	if err = handShake(conn); err != nil {
+		log.Println("socks handshack:", err)
 		return
 	}
+	rawaddr, extra, addr, err := getRequest(conn)
+	if err != nil {
+		log.Println("error getting request:", err)
+		return
+	}
+	// TODO should send error code to client if connect to server failed
+	_, err = conn.Write([]byte{0x05, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x08, 0x43})
+	if err != nil {
+		debug.Println("send connection confirmation:", err)
+		return
+	}
+
+	debug.Println("connecting to", addr)
+	remote, err := ss.DialWithRawAddr(rawaddr, server, encTbl)
+	if err != nil {
+		log.Println("error connect to shadowsocks server:", err)
+		return
+	}
+	defer remote.Close()
+	if extra != nil {
+		debug.Println("writing extra content to remote, len", len(extra))
+		if _, err = remote.Write(extra); err != nil {
+			debug.Println("write request extra error:", err)
+			return
+		}
+	}
 
+	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")
 }
 
 func run(port, password, server string) {

cmd/shadowsocks-server/server.go

@@ -99,7 +99,7 @@ func handleConnection(conn *ss.Conn) {
 
 	host, extra, err := getRequest(conn)
 	if err != nil {
-		debug.Println("error getting request:", err)
+		log.Println("error getting request:", err)
 		return
 	}
 	debug.Println("connecting", host)