add new file

lib/enet/conn.go

@@ -0,0 +1,119 @@
+package enet
+
+import (
+	"ehang.io/nps/lib/pool"
+	"errors"
+	"net"
+	"sync"
+	"syscall"
+)
+
+type Conn interface {
+	net.Conn
+	Reset(int) error
+	Clear()
+	Readable() bool
+	AllBytes() ([]byte, error)
+	SyscallConn() (syscall.RawConn, error)
+}
+
+var _ Conn = (*ReaderConn)(nil)
+
+var bp = pool.NewBufferPool(MaxReadSize)
+
+const MaxReadSize = 32 * 1024
+
+// ReaderConn is an implement of reusable data connection
+type ReaderConn struct {
+	buf      []byte
+	nowIndex int
+	hasRead  int
+	hasClear bool
+	net.Conn
+	sync.RWMutex
+}
+
+// NewReaderConn returns a new ReaderConn
+func NewReaderConn(conn net.Conn) *ReaderConn {
+	return &ReaderConn{Conn: conn, buf: bp.Get()}
+}
+
+// SyscallConn returns a raw network connection
+func (rc *ReaderConn) SyscallConn() (syscall.RawConn, error) {
+	return rc.Conn.(syscall.Conn).SyscallConn()
+}
+
+// Read is an implement of Net.Conn Read function
+func (rc *ReaderConn) Read(b []byte) (n int, err error) {
+	rc.Lock()
+	defer rc.Unlock()
+	if rc.hasClear || (rc.nowIndex == rc.hasRead && rc.hasRead == MaxReadSize) {
+		if !rc.hasClear {
+			rc.Clear()
+		}
+		return rc.Conn.Read(b)
+	}
+	if rc.hasRead > rc.nowIndex {
+		n = copy(b, rc.buf[rc.nowIndex:rc.hasRead])
+		rc.nowIndex += n
+		return
+	}
+	if rc.hasRead == MaxReadSize {
+		n = copy(b, rc.buf[rc.nowIndex:rc.hasRead])
+		rc.nowIndex += n
+		return
+	}
+	err = rc.readOnce()
+	if err != nil {
+		return
+	}
+	n = copy(b, rc.buf[rc.nowIndex:rc.hasRead])
+	rc.nowIndex += n
+	return
+}
+
+// readOnce
+func (rc *ReaderConn) readOnce() error {
+	// int(math.Min(float64(MaxReadSize-rc.hasRead), float64(len(b)-(rc.hasRead-rc.nowIndex))))
+	// read as much as possible to judge whether there is still readable
+	n, err := rc.Conn.Read(rc.buf[rc.nowIndex : rc.hasRead+MaxReadSize-rc.hasRead])
+	rc.hasRead += n
+	return err
+}
+
+// Readable return whether there is data in the buffer
+func (rc *ReaderConn) Readable() bool {
+	return (rc.hasRead - rc.nowIndex) > 0
+}
+
+// AllBytes return all data in the buffer
+func (rc *ReaderConn) AllBytes() ([]byte, error) {
+	rc.Lock()
+	defer rc.Unlock()
+	if rc.hasRead == 0 {
+		if err := rc.readOnce(); err != nil {
+			return nil, err
+		}
+	}
+	if !rc.Readable() {
+		return nil, errors.New("can not read '")
+	}
+	b := rc.buf[rc.nowIndex:rc.hasRead]
+	rc.nowIndex = rc.hasRead
+	return b, nil
+}
+
+// Reset will reset data index
+func (rc *ReaderConn) Reset(n int) error {
+	if !rc.hasClear {
+		rc.nowIndex = n
+		return nil
+	}
+	return errors.New("the enet can not reset anymore")
+}
+
+// Clear will put buf to pool and can not reuse anymore
+func (rc *ReaderConn) Clear() {
+	rc.hasClear = true
+	bp.Put(rc.buf)
+}

lib/enet/conn_test.go

@@ -0,0 +1,81 @@
+package enet
+
+import (
+	"math/rand"
+	"net"
+	"testing"
+	"time"
+)
+
+func TestReaderConn_Read(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:61254")
+	if err != nil {
+		t.Fatal(err)
+	}
+	b := make([]byte, 33*1024)
+	go func() {
+		conn, err := net.Dial("tcp", "127.0.0.1:61254")
+		if err != nil {
+			t.Fatal(err)
+		}
+		rand.Seed(time.Now().UnixNano())
+		for i := 0; i < 33*1024; i++ {
+			b[i] = byte(rand.Intn(128))
+		}
+		conn.Write(b)
+	}()
+	conn, err := ln.Accept()
+	if err != nil {
+		t.Fatal(err)
+	}
+	rConn := NewReaderConn(conn)
+	buf := make([]byte, 1024)
+	nn := 0
+	times := 0
+	for {
+		n, err := rConn.Read(buf)
+		if err != nil {
+			t.Fatal(err)
+		}
+		for i := 0; i < 1024; i++ {
+			if b[times*1024+i] != buf[i] {
+				t.Fatal("data error")
+			}
+		}
+		times++
+		nn += n
+		if nn > 30*1024 {
+			break
+		}
+		if times > 100 {
+			t.Fatal("read error")
+		}
+	}
+
+	rConn.Reset(0)
+	nn = 0
+	times = 0
+	for {
+		n, err := rConn.Read(buf)
+		if err != nil {
+			t.Fatal(err)
+		}
+		for i := 0; i < 1024; i++ {
+			if b[times*1024+i] != buf[i] {
+				t.Fatal("data error")
+			}
+		}
+		nn += n
+		times++
+		if nn > 32*1024 {
+			break
+		}
+		if times > 100 {
+			t.Fatal("read error")
+		}
+	}
+	if !rConn.hasClear || rConn.hasRead != rConn.nowIndex || rConn.nowIndex != MaxReadSize {
+		t.Fatal("read error")
+	}
+
+}

lib/enet/listener.go

@@ -0,0 +1,62 @@
+package enet
+
+import (
+	"errors"
+	"net"
+	"sync/atomic"
+)
+
+var _ net.Listener = (*Listener)(nil)
+
+// Listener is an implementation of net.Listener
+type Listener struct {
+	ch      chan net.Conn
+	closeCh chan struct{}
+	closed  int32
+	nowNum  int32
+	addr    net.Addr
+}
+
+// NewListener returns an initialized Listener
+func NewListener() *Listener {
+	return &Listener{ch: make(chan net.Conn, 10), closeCh: make(chan struct{})}
+}
+
+// SendConn is used to add connection to the listener
+func (bl *Listener) SendConn(c net.Conn) error {
+	if atomic.LoadInt32(&bl.closed) == 1 {
+		return errors.New("the listener is already closed")
+	}
+	atomic.AddInt32(&bl.nowNum, 1)
+	select {
+	case bl.ch <- c:
+		return nil
+	case <-bl.closeCh:
+	}
+	if atomic.AddInt32(&bl.nowNum, -1) == 0 && atomic.LoadInt32(&bl.closed) == 1 {
+		close(bl.ch)
+	}
+	return errors.New("the listener is already closed")
+}
+
+// Accept is used to get connection from the listener
+func (bl *Listener) Accept() (net.Conn, error) {
+	c := <-bl.ch
+	if c == nil {
+		return nil, errors.New("the listener is already closed")
+	}
+	return c, nil
+}
+
+// Close is used to close the listener, it will discard all  existing connections
+func (bl *Listener) Close() error {
+	if atomic.CompareAndSwapInt32(&bl.closed, 0, 1) {
+		close(bl.closeCh)
+	}
+	return nil
+}
+
+// Addr returns the listener's address'
+func (bl *Listener) Addr() net.Addr {
+	return bl.addr
+}

lib/enet/paket.go

@@ -0,0 +1,161 @@
+package enet
+
+import (
+	"bytes"
+	"ehang.io/nps/lib/common"
+	"ehang.io/nps/lib/pool"
+	"github.com/pkg/errors"
+	"net"
+	"sync/atomic"
+	"time"
+)
+
+var (
+	_ net.PacketConn = (*TcpPacketConn)(nil)
+	_ PacketConn     = (*ReaderPacketConn)(nil)
+)
+
+type PacketConn interface {
+	net.PacketConn
+	SendPacket([]byte, net.Addr) error
+	FirstPacket() ([]byte, net.Addr, error)
+}
+
+var udpBp = pool.NewBufferPool(1500)
+
+// TcpPacketConn is an implement of net.PacketConn by net.Conn
+type TcpPacketConn struct {
+	udpBp []byte
+	net.Conn
+}
+
+// NewTcpPacketConn return a *TcpPacketConn
+func NewTcpPacketConn(conn net.Conn) *TcpPacketConn {
+	return &TcpPacketConn{Conn: conn}
+}
+
+// ReadFrom is a implement of net.PacketConn ReadFrom
+func (tp *TcpPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
+	b := udpBp.Get()
+	defer udpBp.Put(b)
+	n, err = common.ReadLenBytes(tp.Conn, b)
+	if err != nil {
+		return
+	}
+	rAddr, err := common.ReadAddr(bytes.NewReader(b[:n]))
+	if err != nil {
+		return
+	}
+	n = copy(p, b[len(rAddr):n])
+	addr, err = net.ResolveUDPAddr("udp", rAddr.String())
+	return
+}
+
+// WriteTo is a implement of net.PacketConn WriteTo
+func (tp *TcpPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
+	var pAddr common.Addr
+	pAddr, err = common.ParseAddr(addr.String())
+	if err != nil {
+		return
+	}
+	return common.WriteLenBytes(tp.Conn, append(pAddr, p...))
+}
+
+// ReaderPacketConn is an implementation of net.PacketConn
+type ReaderPacketConn struct {
+	ch              chan *packet
+	closeCh         chan struct{}
+	closed          int32
+	nowNum          int32
+	addr            net.Addr
+	writePacketConn net.PacketConn
+	readTimer       *time.Timer
+	firstPacket     []byte
+}
+
+type packet struct {
+	b    []byte
+	addr net.Addr
+}
+
+// NewReaderPacketConn returns an initialized PacketConn
+func NewReaderPacketConn(writePacketConn net.PacketConn, firstPacket []byte, addr net.Addr) *ReaderPacketConn {
+	return &ReaderPacketConn{
+		ch:              make(chan *packet, 10),
+		closeCh:         make(chan struct{}),
+		addr:            addr,
+		writePacketConn: writePacketConn,
+		readTimer:       time.NewTimer(time.Hour * 24 * 3650),
+		firstPacket:     firstPacket,
+	}
+}
+
+func (pc *ReaderPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
+	var pt *packet
+	select {
+	case pt = <-pc.ch:
+	case <-pc.readTimer.C:
+	}
+	if pt == nil {
+		return 0, nil, errors.New("the PacketConn is already closed")
+	}
+	copy(p, pt.b)
+	return len(pt.b), pt.addr, nil
+}
+
+func (pc *ReaderPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
+	return pc.writePacketConn.WriteTo(p, addr)
+}
+
+// LocalAddr returns the listener's address
+func (pc *ReaderPacketConn) LocalAddr() net.Addr {
+	return pc.addr
+}
+
+func (pc *ReaderPacketConn) SetDeadline(t time.Time) error {
+	pc.readTimer.Reset(t.Sub(time.Now()))
+	return pc.writePacketConn.SetWriteDeadline(t)
+}
+
+func (pc *ReaderPacketConn) SetReadDeadline(t time.Time) error {
+	pc.readTimer.Reset(t.Sub(time.Now()))
+	return nil
+}
+
+func (pc *ReaderPacketConn) SetWriteDeadline(t time.Time) error {
+	return pc.writePacketConn.SetWriteDeadline(t)
+}
+
+func (pc *ReaderPacketConn) FirstPacket() ([]byte, net.Addr, error) {
+	if pc.firstPacket == nil || pc.addr == nil {
+		return nil, nil, errors.New("not found first packet")
+	}
+	return pc.firstPacket, pc.addr, nil
+}
+
+// SendPacket is used to add connection to the listener
+func (pc *ReaderPacketConn) SendPacket(b []byte, addr net.Addr) error {
+	if atomic.LoadInt32(&pc.closed) == 1 {
+		return errors.New("the listener is already closed")
+	}
+	atomic.AddInt32(&pc.nowNum, 1)
+	select {
+	case pc.ch <- &packet{b: b, addr: addr}:
+		return nil
+	case <-pc.closeCh:
+	case <-pc.readTimer.C:
+		_ = pc.Close()
+	}
+	if atomic.AddInt32(&pc.nowNum, -1) == 0 && atomic.LoadInt32(&pc.closed) == 1 {
+		close(pc.ch)
+	}
+	return errors.New("the packetConn is already closed")
+}
+
+// Close is used to close the listener, it will discard all  existing connections
+func (pc *ReaderPacketConn) Close() error {
+	if atomic.CompareAndSwapInt32(&pc.closed, 0, 1) {
+		close(pc.closeCh)
+	}
+	return nil
+}

lib/enet/paket_test.go

@@ -0,0 +1,78 @@
+package enet
+
+import (
+	"bytes"
+	"github.com/stretchr/testify/assert"
+	"net"
+	"testing"
+)
+
+func TestTcpPacketConn(t *testing.T) {
+	bs := bytes.Repeat([]byte{1}, 100)
+	targetAddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:53")
+	assert.NoError(t, err)
+
+	finish := make(chan struct{}, 0)
+
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	assert.NoError(t, err)
+	go func() {
+		conn, err := ln.Accept()
+		assert.NoError(t, err)
+		b := make([]byte, 1024)
+		n, addr, err := NewTcpPacketConn(conn).ReadFrom(b)
+		assert.NoError(t, err)
+
+		assert.Equal(t, targetAddr, addr)
+		assert.Equal(t, n, 100)
+		finish <- struct{}{}
+	}()
+
+	conn, err := net.Dial("tcp", ln.Addr().String())
+	assert.NoError(t, err)
+
+	_, err = NewTcpPacketConn(conn).WriteTo(bs, targetAddr)
+	assert.NoError(t, err)
+
+	<-finish
+}
+
+func TestPacketConn(t *testing.T) {
+	finish := make(chan struct{}, 0)
+
+	sPacketConn, err := net.ListenPacket("udp", "127.0.0.1:0")
+	assert.NoError(t, err)
+
+	cPacketConn, err := net.ListenPacket("udp", "127.0.0.1:0")
+	assert.NoError(t, err)
+
+	bPacketConn := NewReaderPacketConn(sPacketConn, nil, sPacketConn.LocalAddr())
+
+	sendAddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:53")
+	assert.NoError(t, err)
+
+	go func() {
+		b := make([]byte, 1024)
+		n, addr, err := bPacketConn.ReadFrom(b)
+		assert.NoError(t, err)
+		assert.Equal(t, sendAddr, addr)
+		assert.Equal(t, n, 4)
+
+		_, err = bPacketConn.WriteTo(bytes.Repeat(b[:n], 10), cPacketConn.LocalAddr())
+		assert.NoError(t, err)
+
+		finish <- struct{}{}
+	}()
+
+	err = bPacketConn.SendPacket([]byte{0, 0, 0, 0}, sendAddr)
+	assert.NoError(t, err)
+
+	b := make([]byte, 1024)
+	n, addr, err := cPacketConn.ReadFrom(b)
+	assert.NoError(t, err)
+	assert.Equal(t, n, 40)
+	assert.Equal(t, addr, sPacketConn.LocalAddr())
+
+	<-finish
+
+}

lib/enet/s5_packet.go

@@ -0,0 +1,55 @@
+package enet
+
+import (
+	"ehang.io/nps/lib/common"
+	"ehang.io/nps/lib/pool"
+	"github.com/pkg/errors"
+	"net"
+)
+
+var packetBp = pool.NewBufferPool(1500)
+
+type S5PacketConn struct {
+	net.PacketConn
+	remoteAddr net.Addr
+}
+
+func NewS5PacketConn(pc net.PacketConn, remoteAddr net.Addr) *S5PacketConn {
+	return &S5PacketConn{PacketConn: pc, remoteAddr: remoteAddr}
+}
+
+func (s *S5PacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
+	b := packetBp.Get()
+	defer packetBp.Put(b)
+	n, addr, err = s.PacketConn.ReadFrom(b)
+	if err != nil {
+		return
+	}
+	var targetAddr common.Addr
+	targetAddr, err = common.SplitAddr(b[3:])
+	if err != nil {
+		return
+	}
+	n = copy(p, b[3+len(targetAddr):n])
+	addr, err = net.ResolveUDPAddr("udp", targetAddr.String())
+	return
+}
+
+func (s *S5PacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
+	n = len(p)
+	b := packetBp.Get()
+	defer packetBp.Put(b)
+	var sAddr common.Addr
+	sAddr, err = common.ParseAddr(addr.String())
+	if err != nil {
+		return
+	}
+	copy(b[3:], sAddr)
+	if (3 + len(sAddr) + len(p)) > len(b) {
+		err = errors.Errorf("data too long(%d)", len(p))
+		return
+	}
+	copy(b[3+len(sAddr):], p)
+	_, err = s.PacketConn.WriteTo(b[:3+len(sAddr)+len(p)], s.remoteAddr)
+	return
+}

lib/enet/s5_packet_test.go

@@ -0,0 +1,50 @@
+package enet
+
+import (
+	"ehang.io/nps/lib/common"
+	"github.com/stretchr/testify/assert"
+	"net"
+	"testing"
+)
+
+func TestNewS5PacketConn(t *testing.T) {
+	serverPc, err := net.ListenPacket("udp", "127.0.0.1:0")
+	assert.NoError(t, err)
+	localPc, err := net.ListenPacket("udp", "127.0.0.1:0")
+	assert.NoError(t, err)
+	appAddr, err := net.ResolveUDPAddr("udp", "8.8.8.8:53")
+	assert.NoError(t, err)
+	data := []byte("test")
+	go func() {
+		p := make([]byte, 1500)
+		n, addr, err := serverPc.ReadFrom(p)
+		assert.NoError(t, err)
+		pc := NewReaderPacketConn(serverPc, p[:n], addr)
+		err = pc.SendPacket(p[:n], addr)
+		assert.NoError(t, err)
+
+		_, addr, err = pc.FirstPacket()
+		assert.NoError(t, err)
+		s5Pc := NewS5PacketConn(pc, addr)
+		n, addr, err = s5Pc.ReadFrom(p)
+		assert.NoError(t, err)
+		assert.Equal(t, data, p[:n])
+		assert.Equal(t, addr.String(), "8.8.8.8:53")
+		_, err = s5Pc.WriteTo(data, appAddr)
+		assert.NoError(t, err)
+	}()
+	b := []byte{0, 0, 0}
+	pAddr, err := common.ParseAddr(appAddr.String())
+	assert.NoError(t, err)
+	b = append(b, pAddr...)
+	b = append(b, data...)
+	_, err = localPc.WriteTo(b, serverPc.LocalAddr())
+	assert.NoError(t, err)
+	p := make([]byte, 1500)
+	n, _, err := localPc.ReadFrom(p)
+	assert.NoError(t, err)
+	respAddr, err := common.SplitAddr(p[3:])
+	assert.NoError(t, err)
+	assert.Equal(t, respAddr.String(), appAddr.String())
+	assert.Equal(t, p[3+len(respAddr):n], data)
+}