func getRequest(conn *ss.Conn) (host string, extra []byte, err error) { const ( idType = 0 // address type index idIP0 = 1 // ip addres start index idDmLen = 1 // domain address length index idDm0 = 2 // domain address start index typeIPv4 = 1 // type is ipv4 address typeDm = 3 // type is domain address typeIPv6 = 4 // type is ipv6 address lenIPv4 = 1 + net.IPv4len + 2 // 1addrType + ipv4 + 2port lenIPv6 = 1 + net.IPv6len + 2 // 1addrType + ipv6 + 2port lenDmBase = 1 + 1 + 2 // 1addrType + 1addrLen + 2port, plus addrLen ) // buf size should at least have the same size with the largest possible // request size (when addrType is 3, domain name has at most 256 bytes) // 1(addrType) + 1(lenByte) + 256(max length address) + 2(port) buf := make([]byte, 260) var n int // read till we get possible domain length field ss.SetReadTimeout(conn) if n, err = io.ReadAtLeast(conn, buf, idDmLen+1); err != nil { return } //fmt.Println(buf) reqLen := -1 switch buf[idType] { case typeIPv4: reqLen = lenIPv4 case typeIPv6: reqLen = lenIPv6 case typeDm: reqLen = int(buf[idDmLen]) + lenDmBase default: err = fmt.Errorf("addr type %d not supported", buf[idType]) return } if n < reqLen { // rare case ss.SetReadTimeout(conn) if _, err = io.ReadFull(conn, buf[n:reqLen]); err != nil { return } } else if n > reqLen { // it's possible to read more than just the request head extra = buf[reqLen:n] } tp := strings.Split(conn.LocalAddr().String(), ":") if len(tp) == 2 { p := tp[1] now := time.Now().Format("2006-01-02 15:04:05") if len(flowData.Usage[p]) == 0 { flowData.Usage[p] = []string{"0", now} } used, _ := strconv.ParseUint(flowData.Usage[p][0], 10, 64) used += uint64(n - reqLen) flowData.Usage[p][0] = strconv.FormatUint(used, 10) flowData.Usage[p][1] = now } // Return string for typeIP is not most efficient, but browsers (Chrome, // Safari, Firefox) all seems using typeDm exclusively. So this is not a // big problem. switch buf[idType] { case typeIPv4: host = net.IP(buf[idIP0 : idIP0+net.IPv4len]).String() case typeIPv6: host = net.IP(buf[idIP0 : idIP0+net.IPv6len]).String() case typeDm: host = string(buf[idDm0 : idDm0+buf[idDmLen]]) } // parse port port := binary.BigEndian.Uint16(buf[reqLen-2 : reqLen]) host = net.JoinHostPort(host, strconv.Itoa(int(port))) return }
func handleConnection(conn *ss.Conn, port string) { var host string //fmt.Println(strings.Split(conn.LocalAddr().String(), ":")[1]) connCnt++ // this maybe not accurate, but should be enough if connCnt-nextLogConnCnt >= 0 { // XXX There's no xadd in the atomic package, so it's difficult to log // the message only once with low cost. Also note nextLogConnCnt maybe // added twice for current peak connection number level. log.Printf("Number of client connections reaches %d\n", nextLogConnCnt) nextLogConnCnt += logCntDelta } // function arguments are always evaluated, so surround debug statement // with if statement if debug { debug.Printf("new client %s->%s\n", conn.RemoteAddr().String(), conn.LocalAddr()) } closed := false defer func() { if debug { debug.Printf("closed pipe %s<->%s\n", conn.RemoteAddr(), host) } connCnt-- if !closed { conn.Close() } }() host, extra, err := getRequest(conn) if err != nil { log.Println("error getting request", conn.RemoteAddr(), conn.LocalAddr(), err) return } debug.Println("connecting", host) remote, err := net.Dial("tcp", host) if err != nil { if ne, ok := err.(*net.OpError); ok && (ne.Err == syscall.EMFILE || ne.Err == syscall.ENFILE) { // log too many open file error // EMFILE is process reaches open file limits, ENFILE is system limit log.Println("dial error:", err) } else { log.Println("error connecting to:", host, err) } return } defer func() { if !closed { remote.Close() } }() // write extra bytes read from if extra != nil { // debug.Println("getRequest read extra data, writing to remote, len", len(extra)) if _, err = remote.Write(extra); err != nil { debug.Println("write request extra error:", err) return } } if debug { debug.Printf("piping %s<->%s", conn.RemoteAddr(), host) } go ss.PipeThenClose(conn, remote, ss.SET_TIMEOUT, flowData, port) ss.PipeThenClose(remote, conn, ss.NO_TIMEOUT, flowData, port) closed = true return }