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
            }
        }
    }
}