natpass/code/client/pool/conn.go

package pool

import (
	"fmt"
	"natpass/code/client/global"
	"natpass/code/network"
	"strings"
	"sync"
	"time"

	"github.com/lwch/logging"
)

type Conn struct {
	sync.RWMutex
	ID          string
	parent      *Pool
	conn        *network.Conn
	read        map[string]chan *network.Msg // link id => channel
	unknownRead chan *network.Msg            // read message without link
	write       chan *network.Msg            // link id => channel
}

func newConn(parent *Pool, conn *network.Conn) *Conn {
	ret := &Conn{
		ID:          fmt.Sprintf("%d", time.Now().UnixNano()),
		parent:      parent,
		conn:        conn,
		read:        make(map[string]chan *network.Msg),
		unknownRead: make(chan *network.Msg),
		write:       make(chan *network.Msg),
	}
	logging.Info("new connection: %s", ret.ID)
	go ret.loopRead()
	go ret.loopWrite()
	return ret
}

func (conn *Conn) hasLink(id string) bool {
	conn.RLock()
	defer conn.RUnlock()
	_, ok := conn.read[id]
	return ok
}

// AddLink attach read message
func (conn *Conn) AddLink(id string) {
	conn.Lock()
	conn.read[id] = make(chan *network.Msg)
	conn.Unlock()
}

// RemoveLink detach read message
func (conn *Conn) RemoveLink(id string) {
	conn.Lock()
	ch := conn.read[id]
	if ch != nil {
		close(ch)
	}
	delete(conn.read, id)
	conn.Unlock()
}

// Close close connection
func (conn *Conn) Close() {
	conn.conn.Close()
	conn.Lock()
	for id, ch := range conn.read {
		close(ch)
		delete(conn.read, id)
	}
	conn.Unlock()
	if conn.unknownRead != nil {
		close(conn.unknownRead)
		conn.unknownRead = nil
	}
	if conn.write != nil {
		close(conn.write)
		conn.write = nil
	}
	conn.parent.onClose(conn.ID)
	logging.Error("connection %s closed", conn.ID)
}

func (conn *Conn) loopRead() {
	defer conn.Close()
	for {
		msg, err := conn.conn.ReadMessage(global.ReadTimeout)
		if err != nil {
			if strings.Contains(err.Error(), "i/o timeout") {
				continue
			}
			logging.Error("read message: %v", err)
			return
		}
		var linkID string
		switch msg.GetXType() {
		case network.Msg_connect_req:
			linkID = msg.GetCreq().GetId()
		case network.Msg_connect_rep:
			linkID = msg.GetCrep().GetId()
		case network.Msg_disconnect:
			linkID = msg.GetXDisconnect().GetId()
		case network.Msg_forward:
			linkID = msg.GetXData().GetLid()
		}
		conn.RLock()
		ch := conn.read[linkID]
		conn.RUnlock()
		if ch == nil {
			ch = conn.unknownRead
		}
		select {
		case ch <- msg:
		case <-time.After(global.ReadTimeout):
			logging.Error("write read channel for link %s timeouted", linkID)
			if ch == conn.unknownRead {
				continue
			}
			close(ch)
		}
	}
}

func (conn *Conn) loopWrite() {
	defer conn.Close()
	for {
		msg := <-conn.write
		if msg == nil {
			return
		}
		msg.From = conn.parent.cfg.ID
		err := conn.conn.WriteMessage(msg, global.WriteTimeout)
		if err == nil {
			continue
		}
		logging.Error("write message error on %d: %v", conn.ID, err)
		return
	}
}

// ChanRead get read channel from link id
func (conn *Conn) ChanRead(id string) <-chan *network.Msg {
	conn.RLock()
	defer conn.RUnlock()
	return conn.read[id]
}

func (conn *Conn) ChanUnknown() <-chan *network.Msg {
	return conn.unknownRead
}