package main
import (
"net"
"time"
)
// 初始化代理服务
func initProxy() {
Log.Infof("Proxying %s -> %s\n", 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.\n", 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.\n", 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
}
}
}
}
时间: 2024-11-03 21:53:04