P2p first version

lib/mux/bytes.go

@@ -0,0 +1,32 @@
+package mux
+
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+)
+
+//write bytes with int32 length
+func WriteLenBytes(buf []byte, w io.Writer) (int, error) {
+	raw := bytes.NewBuffer([]byte{})
+	if err := binary.Write(raw, binary.LittleEndian, int32(len(buf))); err != nil {
+		return 0, err
+	}
+	if err := binary.Write(raw, binary.LittleEndian, buf); err != nil {
+		return 0, err
+	}
+	return w.Write(raw.Bytes())
+}
+
+//read bytes by length
+func ReadLenBytes(buf []byte, r io.Reader) (int, error) {
+	var l int32
+	var err error
+	if binary.Read(r, binary.LittleEndian, &l) != nil {
+		return 0, err
+	}
+	if _, err = io.ReadFull(r, buf[:l]); err != nil {
+		return 0, err
+	}
+	return int(l), nil
+}

lib/mux/conn.go

@@ -0,0 +1,148 @@
+package mux
+
+import (
+	"errors"
+	"github.com/cnlh/nps/lib/pool"
+	"io"
+	"net"
+	"time"
+)
+
+type conn struct {
+	net.Conn
+	readMsgCh        chan []byte
+	getStatusCh      chan struct{}
+	connStatusOkCh   chan struct{}
+	connStatusFailCh chan struct{}
+	readTimeOut      time.Time
+	writeTimeOut     time.Time
+	sendMsgCh        chan *msg  //mux
+	sendStatusCh     chan int32 //mux
+	connId           int32
+	isClose          bool
+	mux              *Mux
+}
+
+type msg struct {
+	connId  int32
+	content []byte
+}
+
+func NewMsg(connId int32, content []byte) *msg {
+	return &msg{
+		connId:  connId,
+		content: content,
+	}
+}
+
+func NewConn(connId int32, mux *Mux, sendMsgCh chan *msg, sendStatusCh chan int32) *conn {
+	return &conn{
+		readMsgCh:        make(chan []byte),
+		getStatusCh:      make(chan struct{}),
+		connStatusOkCh:   make(chan struct{}),
+		connStatusFailCh: make(chan struct{}),
+		readTimeOut:      time.Time{},
+		writeTimeOut:     time.Time{},
+		sendMsgCh:        sendMsgCh,
+		sendStatusCh:     sendStatusCh,
+		connId:           connId,
+		isClose:          false,
+		mux:              mux,
+	}
+}
+
+func (s *conn) Read(buf []byte) (int, error) {
+	if s.isClose {
+		return 0, errors.New("the conn has closed")
+	}
+	var b []byte
+	if t := s.readTimeOut.Sub(time.Now()); t > 0 {
+		timer := time.NewTimer(t)
+		select {
+		case <-timer.C:
+			s.Close()
+			return 0, errors.New("read timeout")
+		case b = <-s.readMsgCh:
+		}
+	} else {
+		b = <-s.readMsgCh
+	}
+	defer pool.PutBufPoolCopy(b)
+	if s.isClose {
+		return 0, io.EOF
+	}
+	s.sendStatusCh <- s.connId
+	return copy(buf, b), nil
+}
+
+func (s *conn) Write(buf []byte) (int, error) {
+	if s.isClose {
+		return 0, errors.New("the conn has closed")
+	}
+
+	if t := s.writeTimeOut.Sub(time.Now()); t > 0 {
+		timer := time.NewTimer(t)
+		select {
+		case <-timer.C:
+			s.Close()
+			return 0, errors.New("write timeout")
+		case s.sendMsgCh <- NewMsg(s.connId, buf):
+		}
+	} else {
+		s.sendMsgCh <- NewMsg(s.connId, buf)
+	}
+
+	if t := s.writeTimeOut.Sub(time.Now()); t > 0 {
+		timer := time.NewTimer(t)
+		select {
+		case <-timer.C:
+			s.Close()
+			return 0, errors.New("write timeout")
+		case <-s.getStatusCh:
+		}
+	} else {
+		<-s.getStatusCh
+	}
+
+	if s.isClose {
+		return 0, io.EOF
+	}
+	return len(buf), nil
+}
+
+func (s *conn) Close() error {
+	if s.isClose {
+		return errors.New("the conn has closed")
+	}
+	s.isClose = true
+	close(s.getStatusCh)
+	close(s.readMsgCh)
+	close(s.connStatusOkCh)
+	close(s.connStatusFailCh)
+	s.sendMsgCh <- NewMsg(s.connId, nil)
+	return nil
+}
+
+func (s *conn) LocalAddr() net.Addr {
+	return s.mux.conn.LocalAddr()
+}
+
+func (s *conn) RemoteAddr() net.Addr {
+	return s.mux.conn.RemoteAddr()
+}
+
+func (s *conn) SetDeadline(t time.Time) error {
+	s.readTimeOut = t
+	s.writeTimeOut = t
+	return nil
+}
+
+func (s *conn) SetReadDeadline(t time.Time) error {
+	s.readTimeOut = t
+	return nil
+}
+
+func (s *conn) SetWriteDeadline(t time.Time) error {
+	s.writeTimeOut = t
+	return nil
+}

lib/mux/map.go

@@ -0,0 +1,64 @@
+package mux
+
+import (
+	"sync"
+	"time"
+)
+
+type connMap struct {
+	connMap map[int32]*conn
+	closeCh chan struct{}
+	sync.RWMutex
+}
+
+func NewConnMap() *connMap {
+	connMap := &connMap{
+		connMap: make(map[int32]*conn),
+		closeCh: make(chan struct{}),
+	}
+	go connMap.clean()
+	return connMap
+}
+
+func (s *connMap) Get(id int32) (*conn, bool) {
+	s.Lock()
+	defer s.Unlock()
+	if v, ok := s.connMap[id]; ok {
+		return v, true
+	}
+	return nil, false
+}
+
+func (s *connMap) Set(id int32, v *conn) {
+	s.Lock()
+	defer s.Unlock()
+	s.connMap[id] = v
+}
+
+func (s *connMap) Close() {
+	s.Lock()
+	defer s.Unlock()
+	for _, v := range s.connMap {
+		v.isClose = true
+	}
+	s.closeCh <- struct{}{}
+}
+
+func (s *connMap) clean() {
+	ticker := time.NewTimer(time.Minute * 1)
+	for {
+		select {
+		case <-ticker.C:
+			s.Lock()
+			for _, v := range s.connMap {
+				if v.isClose {
+					delete(s.connMap, v.connId)
+				}
+			}
+			s.Unlock()
+		case <-s.closeCh:
+			ticker.Stop()
+			return
+		}
+	}
+}

lib/mux/mux.go

@@ -0,0 +1,226 @@
+package mux
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"github.com/cnlh/nps/lib/pool"
+	"math"
+	"net"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+const (
+	MUX_PING_FLAG int32 = iota
+	MUX_NEW_CONN_OK
+	MUX_NEW_CONN_Fail
+	MUX_NEW_MSG
+	MUX_MSG_SEND_OK
+	MUX_NEW_CONN
+	MUX_PING
+	MUX_CONN_CLOSE
+)
+
+type Mux struct {
+	net.Listener
+	conn         net.Conn
+	connMap      *connMap
+	sendMsgCh    chan *msg  //write msg chan
+	sendStatusCh chan int32 //write read ok chan
+	newConnCh    chan *conn
+	id           int32
+	closeChan    chan struct{}
+	isClose      bool
+	sync.Mutex
+}
+
+func NewMux(c net.Conn) *Mux {
+	m := &Mux{
+		conn:         c,
+		connMap:      NewConnMap(),
+		sendMsgCh:    make(chan *msg),
+		sendStatusCh: make(chan int32),
+		id:           0,
+		closeChan:    make(chan struct{}),
+		newConnCh:    make(chan *conn),
+		isClose:      false,
+	}
+	//read session by flag
+	go m.readSession()
+	//write session
+	go m.writeSession()
+	//ping
+	go m.ping()
+	return m
+}
+
+func (s *Mux) NewConn() (*conn, error) {
+	if s.isClose {
+		return nil, errors.New("the mux has closed")
+	}
+	conn := NewConn(s.getId(), s, s.sendMsgCh, s.sendStatusCh)
+	raw := bytes.NewBuffer([]byte{})
+	if err := binary.Write(raw, binary.LittleEndian, MUX_NEW_CONN); err != nil {
+		return nil, err
+	}
+	if err := binary.Write(raw, binary.LittleEndian, conn.connId); err != nil {
+		return nil, err
+	}
+	//it must be set before send
+	s.connMap.Set(conn.connId, conn)
+	if _, err := s.conn.Write(raw.Bytes()); err != nil {
+		return nil, err
+	}
+	select {
+	case <-conn.connStatusOkCh:
+		return conn, nil
+	case <-conn.connStatusFailCh:
+	}
+	return nil, errors.New("create connection fail，the server refused the connection")
+}
+
+func (s *Mux) Accept() (net.Conn, error) {
+	if s.isClose {
+		return nil, errors.New("accpet error,the conn has closed")
+	}
+	return <-s.newConnCh, nil
+}
+
+func (s *Mux) Addr() net.Addr {
+	return s.conn.LocalAddr()
+}
+
+func (s *Mux) ping() {
+	go func() {
+		ticker := time.NewTicker(time.Second * 5)
+		raw := bytes.NewBuffer([]byte{})
+		for {
+			select {
+			case <-ticker.C:
+			}
+			//Avoid going beyond the scope
+			if (math.MaxInt32 - s.id) < 10000 {
+				s.id = 0
+			}
+			raw.Reset()
+			binary.Write(raw, binary.LittleEndian, MUX_PING_FLAG)
+			binary.Write(raw, binary.LittleEndian, MUX_PING)
+			if _, err := s.conn.Write(raw.Bytes()); err != nil {
+				s.Close()
+				break
+			}
+		}
+	}()
+	select {
+	case <-s.closeChan:
+	}
+}
+
+func (s *Mux) writeSession() {
+	go func() {
+		raw := bytes.NewBuffer([]byte{})
+		for {
+			raw.Reset()
+			select {
+			case msg := <-s.sendMsgCh:
+				if msg.content == nil { //close
+					binary.Write(raw, binary.LittleEndian, MUX_CONN_CLOSE)
+					binary.Write(raw, binary.LittleEndian, msg.connId)
+					break
+				}
+				binary.Write(raw, binary.LittleEndian, MUX_NEW_MSG)
+				binary.Write(raw, binary.LittleEndian, msg.connId)
+				binary.Write(raw, binary.LittleEndian, int32(len(msg.content)))
+				binary.Write(raw, binary.LittleEndian, msg.content)
+			case connId := <-s.sendStatusCh:
+				binary.Write(raw, binary.LittleEndian, MUX_MSG_SEND_OK)
+				binary.Write(raw, binary.LittleEndian, connId)
+			}
+			if _, err := s.conn.Write(raw.Bytes()); err != nil {
+				s.Close()
+				break
+			}
+		}
+	}()
+	select {
+	case <-s.closeChan:
+	}
+}
+
+func (s *Mux) readSession() {
+	go func() {
+		raw := bytes.NewBuffer([]byte{})
+		for {
+			var flag, i int32
+			if binary.Read(s.conn, binary.LittleEndian, &flag) == nil {
+				if binary.Read(s.conn, binary.LittleEndian, &i) != nil {
+					break
+				}
+				switch flag {
+				case MUX_NEW_CONN: //new conn
+					conn := NewConn(i, s, s.sendMsgCh, s.sendStatusCh)
+					s.connMap.Set(i, conn) //it has been set before send ok
+					s.newConnCh <- conn
+					raw.Reset()
+					binary.Write(raw, binary.LittleEndian, MUX_NEW_CONN_OK)
+					binary.Write(raw, binary.LittleEndian, i)
+					s.conn.Write(raw.Bytes())
+					continue
+				case MUX_PING_FLAG: //ping
+					continue
+				}
+				if conn, ok := s.connMap.Get(i); ok {
+					switch flag {
+					case MUX_NEW_MSG: //new msg from remote conn
+						buf := pool.BufPoolCopy.Get().([]byte)
+						if n, err := ReadLenBytes(buf, s.conn); err == nil {
+							if !conn.isClose {
+								conn.readMsgCh <- buf[:n]
+							} else {
+								pool.PutBufPoolCopy(buf)
+							}
+						} else { //read len bytes error,the mux has broken
+							break
+						}
+					case MUX_MSG_SEND_OK: //the remote has read
+						conn.getStatusCh <- struct{}{}
+					case MUX_NEW_CONN_OK: //conn ok
+						conn.connStatusOkCh <- struct{}{}
+					case MUX_NEW_CONN_Fail:
+						conn.connStatusFailCh <- struct{}{}
+					case MUX_CONN_CLOSE: //close the connection
+						conn.Close()
+					}
+				}
+			} else {
+				break
+			}
+		}
+		s.Close()
+	}()
+	select {
+	case <-s.closeChan:
+	}
+}
+
+func (s *Mux) Close() error {
+	if s.isClose {
+		return errors.New("the mux has closed")
+	}
+	s.isClose = true
+	s.connMap.Close()
+	s.closeChan <- struct{}{}
+	s.closeChan <- struct{}{}
+	s.closeChan <- struct{}{}
+	close(s.closeChan)
+	close(s.sendMsgCh)
+	close(s.sendStatusCh)
+	return s.conn.Close()
+}
+
+//get new connId as unique flag
+func (s *Mux) getId() int32 {
+	return atomic.AddInt32(&s.id, 1)
+}

lib/mux/mux_test.go

@@ -0,0 +1,96 @@
+package mux
+
+import (
+	"github.com/cnlh/nps/lib/common"
+	"github.com/cnlh/nps/vender/github.com/astaxie/beego/logs"
+	"log"
+	"net"
+	"net/http"
+	_ "net/http/pprof"
+	"testing"
+	"time"
+)
+
+var conn1 net.Conn
+var conn2 net.Conn
+
+func TestNewMux(t *testing.T) {
+	go func() {
+		http.ListenAndServe("0.0.0.0:8899", nil)
+	}()
+	logs.EnableFuncCallDepth(true)
+	logs.SetLogFuncCallDepth(3)
+	server()
+	client()
+	time.Sleep(time.Second * 3)
+	go func() {
+		m2 := NewMux(conn2)
+		for {
+			c, err := m2.Accept()
+			if err != nil {
+				log.Fatalln(err)
+			}
+			go func(c net.Conn) {
+				c2, err := net.Dial("tcp", "127.0.0.1:8080")
+				if err != nil {
+					log.Fatalln(err)
+				}
+				go common.CopyBuffer(c2, c)
+				common.CopyBuffer(c, c2)
+				c.Close()
+				c2.Close()
+			}(c)
+		}
+	}()
+
+	go func() {
+		m1 := NewMux(conn1)
+		l, err := net.Listen("tcp", "127.0.0.1:7777")
+		if err != nil {
+			log.Fatalln(err)
+		}
+		for {
+			conn, err := l.Accept()
+			if err != nil {
+				log.Fatalln(err)
+			}
+			go func(conn net.Conn) {
+				tmpCpnn, err := m1.NewConn()
+				if err != nil {
+					log.Fatalln(err)
+				}
+				go common.CopyBuffer(tmpCpnn, conn)
+				common.CopyBuffer(conn, tmpCpnn)
+				conn.Close()
+				tmpCpnn.Close()
+			}(conn)
+		}
+	}()
+
+	for {
+		time.Sleep(time.Second * 5)
+	}
+}
+
+func server() {
+	var err error
+	l, err := net.Listen("tcp", "127.0.0.1:9999")
+	if err != nil {
+		log.Fatalln(err)
+	}
+	go func() {
+		conn1, err = l.Accept()
+		if err != nil {
+			log.Fatalln(err)
+		}
+	}()
+	return
+}
+
+func client() {
+	var err error
+	conn2, err = net.Dial("tcp", "127.0.0.1:9999")
+	if err != nil {
+		log.Fatalln(err)
+	}
+}