github.com/keysonzzz/kmg@v0.0.0-20151121023212-05317bfd7d39/kmgNet/kmgTfo/TfoDial_linux.go (about)

     1  package kmgTfo
     2  
     3  import (
     4  	"net"
     5  	"os"
     6  	"sync"
     7  
     8  	"github.com/bronze1man/kmg/kmgNet"
     9  	"github.com/bronze1man/kmg/kmgNet/kmgUnix"
    10  	"golang.org/x/sys/unix"
    11  )
    12  
    13  // dial tcp with tcp fastopen
    14  // you should use
    15  //     echo 3 > /proc/sys/net/ipv4/tcp_fastopen
    16  //   to enable it
    17  // you should use linux kernel version >3.7
    18  // you should write something before read to use this function.
    19  // network is useless, it always use tcp4
    20  func TfoLazyDial(network string, nextAddr string) (conn net.Conn, err error) {
    21  	return &tfoLazyConn{nextAddr: nextAddr}, nil
    22  }
    23  
    24  type tfoLazyConn struct {
    25  	net.Conn
    26  	nextAddr string
    27  	lock     sync.Mutex
    28  	isClosed bool
    29  }
    30  
    31  func (c *tfoLazyConn) Read(b []byte) (n int, err error) {
    32  	c.lock.Lock()
    33  	if c.Conn != nil && !c.isClosed {
    34  		c.lock.Unlock()
    35  		return c.Conn.Read(b)
    36  	}
    37  	if c.isClosed {
    38  		c.lock.Unlock()
    39  		return 0, kmgNet.ErrClosing
    40  	}
    41  	c.Conn, err = net.Dial("tcp", c.nextAddr)
    42  	if err != nil {
    43  		c.lock.Unlock()
    44  		return
    45  	}
    46  	c.lock.Unlock()
    47  	return c.Conn.Read(b)
    48  }
    49  
    50  func (c *tfoLazyConn) Write(b []byte) (n int, err error) {
    51  	c.lock.Lock()
    52  	if c.Conn != nil && !c.isClosed {
    53  		c.lock.Unlock()
    54  		return c.Conn.Write(b)
    55  	}
    56  	if c.isClosed {
    57  		c.lock.Unlock()
    58  		return 0, kmgNet.ErrClosing
    59  	}
    60  	if len(b) <= tfoFirstSize {
    61  		c.Conn, err = TfoDial(c.nextAddr, b)
    62  		if err != nil {
    63  			c.lock.Unlock()
    64  			return 0, err
    65  		}
    66  		c.lock.Unlock()
    67  		return len(b), err
    68  	} else {
    69  		c.Conn, err = TfoDial(c.nextAddr, b[:tfoFirstSize])
    70  		if err != nil {
    71  			c.lock.Unlock()
    72  			return 0, err
    73  		}
    74  		c.lock.Unlock()
    75  		_, err = c.Conn.Write(b[tfoFirstSize:])
    76  		if err != nil {
    77  			return 0, err
    78  		}
    79  		return len(b), err
    80  	}
    81  }
    82  
    83  func (c *tfoLazyConn) Close() error {
    84  	c.lock.Lock()
    85  	if c.isClosed {
    86  		c.lock.Unlock()
    87  		return kmgNet.ErrClosing
    88  	}
    89  	c.isClosed = true
    90  	if c.Conn != nil {
    91  		c.lock.Unlock()
    92  		return c.Conn.Close()
    93  	}
    94  	c.lock.Unlock()
    95  	return nil
    96  }
    97  
    98  const tfoFirstSize = 1000
    99  
   100  // dial tcp with tcp fastopen
   101  // 第一个包体积不要太大,需要小于一定数量,否则会被吃掉(正确性问题),
   102  // 如果过大,此处会在连接时发送前一部分,连接后又发送一部分
   103  func TfoDial(nextAddr string, data []byte) (conn net.Conn, err error) {
   104  	s, err := unix.Socket(unix.AF_INET, unix.SOCK_STREAM, 0)
   105  	if err != nil {
   106  		return nil, err
   107  	}
   108  	defer unix.Close(s)
   109  	sa, err := kmgUnix.IPv4TcpAddrToUnixSocksAddr(nextAddr)
   110  	if err != nil {
   111  		return nil, err
   112  	}
   113  	if len(data) <= tfoFirstSize {
   114  		err = unix.Sendto(s, data, unix.MSG_FASTOPEN, sa)
   115  		if err != nil {
   116  			return
   117  		}
   118  	} else {
   119  		err = unix.Sendto(s, data[:tfoFirstSize], unix.MSG_FASTOPEN, sa)
   120  		if err != nil {
   121  			return
   122  		}
   123  	}
   124  	f := os.NewFile(uintptr(s), "TFODial")
   125  	defer f.Close()
   126  	conn, err = net.FileConn(f)
   127  	if err != nil {
   128  		return
   129  	}
   130  	if len(data) > tfoFirstSize {
   131  		_, err = conn.Write(data[tfoFirstSize:])
   132  		if err != nil {
   133  			return nil, err
   134  		}
   135  	}
   136  	return conn, nil
   137  }