package main
import (
"net"
"time"
)
// 初始化代理服务
func initProxy() { Log.Infof("Proxying %s -> %s
", Config.Bind, Config.Backend) server, err := net.Listen("tcp", Config.Bind) if err != nil { Log.Fatal(err)
}
// 等待的队列长度
waitQueue := make(chan net.Conn, Config.WaitQueueLen)
// 最大并发连接
connPools := make(chan bool, Config.MaxConn)
for i := 0; i < Config.MaxConn; i++ { connPools <- true
}
// 等待连接处理
go waitConn(waitQueue, connPools)
// 接收连接并抛给管道处理
for { conn, err := server.Accept()
if err != nil { Log.Error(err)
continue
}
Log.Infof("Received connection from %s.
", conn.RemoteAddr()) waitQueue <- conn
}
}
// 连接数控制
func waitConn(waitQueue chan net.Conn, connPools chan bool) { for conn := range waitQueue { // 接收一个链接，连接池释放一个
<-connPools
go func(conn net.Conn) { handleConn(conn)
// 链接处理完毕，增加
connPools <- true
Log.Infof("Closed connection from %s.
", conn.RemoteAddr()) }(conn)
}
}
// 处理连接
func handleConn(conn net.Conn) { defer conn.Close()
// 根据链接哈希选择机器
proxySvr, ok := getBackendSvr(conn)
if !ok { return
}
// 链接远程代理服务器
remote, err := net.Dial("tcp", proxySvr.identify) if err != nil { Log.Error(err)
proxySvr.failTimes++
return
}
// 等待双向连接完成
complete := make(chan bool, 2)
oneSwitch := make(chan bool, 1)
otherSwitch := make(chan bool, 1)
// 将当前客户端链接发送的数据发送给远程被代理的服务器
go transaction(conn, remote, complete, oneSwitch, otherSwitch)
// 将远程服务返回的数据返回给客户端
go transaction(remote, conn, complete, otherSwitch, oneSwitch)
<-complete
<-complete
remote.Close()
}
// 数据交换传输（从from读数据，再写入to）
func transaction(from, to net.Conn, complete, oneSwitch, otherSwitch chan bool) { var err error
var read int
bytes := make([]byte, 1024)
for { select { case <-otherSwitch:
complete <- true
return
default:
timeOutSec := time.Duration(Config.Timeout) * time.Second
// 设置超时时间
from.SetReadDeadline(time.Now().Add(timeOutSec))
read, err = from.Read(bytes)
if err != nil { complete <- true
oneSwitch <- true
return
}
// 设置超时时间
to.SetWriteDeadline(time.Now().Add(timeOutSec))
_, err = to.Write(bytes[:read])
if err != nil { complete <- true
oneSwitch <- true
return
}
}
}
}