From 231f434e73210df2e1ddef101c625d8c731a7759 Mon Sep 17 00:00:00 2001
From: Dingjun <fangdingjun@gmail.com>
Date: Sat, 14 Jan 2017 10:54:25 +0800
Subject: [PATCH] ntohs: add byte order determine

---
 client.go         |  3 +--
 redir_iptables.go | 47 +++++++++++++++++++++++++++++++++++++++--------
 2 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/client.go b/client.go
index 2eedaf6..f419fa8 100644
--- a/client.go
+++ b/client.go
@@ -265,8 +265,7 @@ func (cc *Client) handleDynamicForward(conn net.Conn) {
 		}
 	} else {
 		// SO_ORIGNAL_DST failed
-		// ipv6 not support this syscall
-		// so ignore it
+		// just ignore it
 		Log(DEBUG, "get original destination on %s failed: %s, ignore",
 			conn.LocalAddr(), err)
 	}
diff --git a/redir_iptables.go b/redir_iptables.go
index 922a87f..14b15e4 100644
--- a/redir_iptables.go
+++ b/redir_iptables.go
@@ -31,6 +31,15 @@ func getOriginDst(c net.Conn) (net.Addr, error) {
 
 	defer f.Close()
 
+	// ipv4
+	level := syscall.SOL_IP
+
+	remoteIP := c.RemoteAddr().(*net.TCPAddr).IP
+	if remoteIP.To4() == nil {
+		// ipv6
+		level = syscall.SOL_IPV6
+	}
+
 	// get original ip destination, in C like this
 	//
 	// struct sockaddr addr;
@@ -39,7 +48,7 @@ func getOriginDst(c net.Conn) (net.Addr, error) {
 	// getsocketopt(fd, SOL_IP, SO_ORIGINAL_DST, &addr, &len);
 	//
 	_, _, errno := syscall.Syscall6(sysGetSockOpt, f.Fd(),
-		uintptr(syscall.SOL_IP), uintptr(soOriginalDst),
+		uintptr(level), uintptr(soOriginalDst),
 		uintptr(unsafe.Pointer(&sockaddr)),
 		uintptr(unsafe.Pointer(&len)), 0)
 
@@ -54,11 +63,11 @@ func getOriginDst(c net.Conn) (net.Addr, error) {
 	case syscall.AF_INET:
 		a := (*syscall.RawSockaddrInet4)(unsafe.Pointer(&sockaddr))
 		ip = net.IP(a.Addr[0:])
-		port = ntohl(a.Port)
+		port = ntohs(a.Port)
 	case syscall.AF_INET6:
 		a := (*syscall.RawSockaddrInet6)(unsafe.Pointer(&sockaddr))
 		ip = net.IP(a.Addr[0:])
-		port = ntohl(a.Port)
+		port = ntohs(a.Port)
 	default:
 		return nil, fmt.Errorf("unknown socket family: %d",
 			sockaddr.Addr.Family)
@@ -68,9 +77,31 @@ func getOriginDst(c net.Conn) (net.Addr, error) {
 	return addr, nil
 }
 
-func ntohl(a uint16) uint16 {
-	b := make([]byte, 2)
-	binary.BigEndian.PutUint16(b, a)
-	c := binary.LittleEndian.Uint16(b)
-	return c
+func ntohs(a uint16) uint16 {
+	if isLittleEndian {
+		b := make([]byte, 2)
+		binary.BigEndian.PutUint16(b, a)
+		c := binary.LittleEndian.Uint16(b)
+		return c
+	}
+	return a
+}
+
+var isLittleEndian = isHostLittleEndian()
+
+func isHostLittleEndian() bool {
+	// determine the byte order
+
+	var num uint16 = 0x1234
+
+	buf := make([]byte, 2)
+
+	binary.BigEndian.PutUint16(buf, num)
+	p := (*[2]byte)(unsafe.Pointer(&num))
+	if p[0] != buf[0] {
+		// little endian
+		return true
+	}
+	// big endian
+	return false
 }