nps/lib/bridge.go
2019-01-08 00:21:02 +08:00

229 lines
4.6 KiB
Go
Executable File
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package lib
import (
"errors"
"log"
"net"
"sync"
"time"
)
type list struct {
connList chan *Conn
}
func (l *list) Add(c *Conn) {
l.connList <- c
}
func (l *list) Pop() *Conn {
return <-l.connList
}
func (l *list) Len() int {
return len(l.connList)
}
func newList() *list {
l := new(list)
l.connList = make(chan *Conn, 1000)
return l
}
type Tunnel struct {
tunnelPort int //通信隧道端口
listener *net.TCPListener //server端监听
signalList map[string]*list //通信
tunnelList map[string]*list //隧道
lock sync.Mutex
tunnelLock sync.Mutex
}
func newTunnel(tunnelPort int) *Tunnel {
t := new(Tunnel)
t.tunnelPort = tunnelPort
t.signalList = make(map[string]*list)
t.tunnelList = make(map[string]*list)
return t
}
func (s *Tunnel) StartTunnel() error {
var err error
s.listener, err = net.ListenTCP("tcp", &net.TCPAddr{net.ParseIP("0.0.0.0"), s.tunnelPort, ""})
if err != nil {
return err
}
go s.tunnelProcess()
return nil
}
//tcp server
func (s *Tunnel) tunnelProcess() error {
var err error
for {
conn, err := s.listener.Accept()
if err != nil {
log.Println(err)
continue
}
go s.cliProcess(NewConn(conn))
}
return err
}
//验证失败返回错误验证flag并且关闭连接
func (s *Tunnel) verifyError(c *Conn) {
c.conn.Write([]byte(VERIFY_EER))
c.conn.Close()
}
func (s *Tunnel) cliProcess(c *Conn) error {
c.conn.(*net.TCPConn).SetReadDeadline(time.Now().Add(time.Duration(5) * time.Second))
vval := make([]byte, 32)
if _, err := c.conn.Read(vval); err != nil {
log.Println("客户端读超时。客户端地址为::", c.conn.RemoteAddr())
c.conn.Close()
return err
}
if !verify(string(vval)) {
log.Println("当前客户端连接校验错误,关闭此客户端:", c.conn.RemoteAddr())
s.verifyError(c)
return err
}
c.conn.(*net.TCPConn).SetReadDeadline(time.Time{})
//做一个判断 添加到对应的channel里面以供使用
if flag, err := c.ReadFlag(); err != nil {
return err
} else {
return s.typeDeal(flag, c, string(vval))
}
}
//tcp连接类型区分
func (s *Tunnel) typeDeal(typeVal string, c *Conn, cFlag string) error {
switch typeVal {
case WORK_MAIN:
s.addList(s.signalList, c, cFlag)
case WORK_CHAN:
s.addList(s.tunnelList, c, cFlag)
default:
return errors.New("无法识别")
}
c.SetAlive()
return nil
}
//加到对应的list中
func (s *Tunnel) addList(m map[string]*list, c *Conn, cFlag string) {
s.lock.Lock()
if v, ok := m[cFlag]; ok {
v.Add(c)
} else {
l := newList()
l.Add(c)
m[cFlag] = l
}
s.lock.Unlock()
}
//新建隧道
func (s *Tunnel) newChan(cFlag string) error {
if err := s.wait(s.signalList, cFlag); err != nil {
return err
}
retry:
connPass := s.signalList[cFlag].Pop()
_, err := connPass.conn.Write([]byte("chan"))
if err != nil {
log.Println(err)
goto retry
}
s.signalList[cFlag].Add(connPass)
return nil
}
//得到一个tcp隧道
func (s *Tunnel) GetTunnel(cFlag string, en, de int, crypt, mux bool) (c *Conn, err error) {
s.tunnelLock.Lock()
if v, ok := s.tunnelList[cFlag]; !ok || v.Len() < 3 { //新建通道
go s.newChan(cFlag)
}
retry:
if err = s.wait(s.tunnelList, cFlag); err != nil {
return
}
c = s.tunnelList[cFlag].Pop()
if _, err = c.wTest(); err != nil {
c.Close()
goto retry
}
c.WriteConnInfo(en, de, crypt, mux)
s.tunnelLock.Unlock()
return
}
//得到一个通信通道
func (s *Tunnel) GetSignal(cFlag string) (err error, conn *Conn) {
if v, ok := s.signalList[cFlag]; !ok || v.Len() == 0 {
err = errors.New("客户端未连接")
return
}
conn = s.signalList[cFlag].Pop()
return
}
//重回slice 复用
func (s *Tunnel) ReturnSignal(conn *Conn, cFlag string) {
if v, ok := s.signalList[cFlag]; ok {
v.Add(conn)
}
}
//重回slice 复用
func (s *Tunnel) ReturnTunnel(conn *Conn, cFlag string) {
if v, ok := s.tunnelList[cFlag]; ok {
FlushConn(conn.conn)
v.Add(conn)
}
}
//删除通信通道
func (s *Tunnel) DelClientSignal(cFlag string) {
s.delClient(cFlag, s.signalList)
}
//删除隧道
func (s *Tunnel) DelClientTunnel(cFlag string) {
s.delClient(cFlag, s.tunnelList)
}
func (s *Tunnel) delClient(cFlag string, l map[string]*list) {
if t := l[getverifyval(cFlag)]; t != nil {
for {
if t.Len() <= 0 {
break
}
t.Pop().Close()
}
delete(l, getverifyval(cFlag))
}
}
//等待
func (s *Tunnel) wait(m map[string]*list, cFlag string) error {
ticker := time.NewTicker(time.Millisecond * 100)
stop := time.After(time.Second * 10)
loop:
for {
select {
case <-ticker.C:
if _, ok := m[cFlag]; ok {
ticker.Stop()
break loop
}
case <-stop:
return errors.New("client key: " + cFlag + ",err: get client conn timeout")
}
}
return nil
}