快捷導(dǎo)航

用GO實現(xiàn)IP門禁優(yōu)化網(wǎng)絡(luò)流量管理

更新時間：2023年12月25日 10:31:13 作者：南丞?PFinalClub

這篇文章主要為大家介紹了用GO實現(xiàn)IP門禁優(yōu)化網(wǎng)絡(luò)流量管理,有需要的朋友可以借鑒參考下,希望能夠有所幫助,祝大家多多進步,早日升職加薪

背景

最近在處理數(shù)據(jù)時，發(fā)現(xiàn)線上某個接口的訪問量異常高。在初次設(shè)計時，并未對流量和訪問量進行限制，因此對具體情況并不太清楚。為了解決這一問題，搞了一個簡單的IP過濾腳本，并進行了記錄。

之前搞IP過濾的時候, 使用了 PHP 的 workerman 來搞的其實思路差不多一樣:

$worker = new Worker('tcp:0.0.0.0:8080');
// 監(jiān)聽一個端口
$worker->count = 2;
// 設(shè)置多進程
$worker->onConnect = function (TcpConnection $connection) {
    $list_ip = []; // ip 白名單 
    $remote_ip = $connection->getRemoteIp();
    // 攔截IP
    if (!in_array($remote_ip, $list_ip)) {
        $connection->close();
    }
    // 放行連接，連接內(nèi)部目標(biāo)端口
    $to_connection = new AsyncTcpConnection('tcp:127.0.0.1:80');
    // 互相轉(zhuǎn)發(fā)流量
    $connection->pipe($to_connection);
    $to_connection->pipe($connection);
    $to_connection->connect();
}

以上PHP代碼創(chuàng)建了一個TCP監(jiān)聽器，在連接時進行IP過濾，并將流量轉(zhuǎn)發(fā)到目標(biāo)端口。

Go實現(xiàn)客戶端

func main() {
	listener, err := net.Listen("tcp", "127.0.0.1:8080")
	if err != nil {
		fmt.Println("Listen() failed, err: ", err)
		return
	}
	defer func(listener net.Listener) {
		_ = listener.Close()
	}(listener)
	for {
		conn, err := listener.Accept()
		if err != nil {
			fmt.Println("Accept() failed, err: ", err)
			continue
		}
		ip, _, _ := net.SplitHostPort(conn.RemoteAddr().String())
		if !isIPAllowed(ip) {
			fmt.Println("IP: ", ip, " is not allowed")
			_ = conn.Close()
			// TODO 對非白名單IP的連接，應(yīng)該進行一些處理，比如記錄日志
			continue
		}
		fmt.Println("IP: ", ip, " is allowed")
		go handleConn(conn)  // 轉(zhuǎn)發(fā)到固定的目標(biāo)端口
	}
}
func handleConn(conn net.Conn) {
	defer func(conn net.Conn) {
		_ = conn.Close()
	}(conn)
	target, err := net.Dial("tcp", "127.0.0.1:80")
	if err != nil {
		fmt.Println("Dial() failed, err: ", err)
		return
	}
	defer func(target net.Conn) {
		_ = target.Close()
	}(target)
	// 使用 copyAndCount 的并發(fā)執(zhí)行版本（同時處理流向）
	go copyAndCount(target, conn, &inTraffic)
	copyAndCount(conn, target, &outTraffic) // Block until this copy is done
}