当前位置：首页 > article >正文

高效 TCP 代理服务器的实战解析：Go 语言编写的高性能代理实例20241028

article 2024/11/1 17:40:21

高效 TCP 代理服务器的实战解析：Go 语言编写的高性能代理实例

引言

Go 语言天然支持并发处理，因此在高性能网络服务和代理服务开发中颇具优势。本文将从配置、并发控制、日志记录和优雅关机等方面，深入解读一个基于 Go 的 TCP 代理服务器项目。希望通过这个实战项目，帮助读者更好地理解 Go 语言在构建高性能应用时的应用场景与技巧。

目录结构

在项目目录结构设计上，我们保持文件的简洁性，并通过 Makefile 和 .gitignore 管理构建过程和版本控制：

go-tcp-proxy/
├── go_tcp_proxy.go                # 主程序文件
├── config.json                    # 配置文件
├── Makefile                       # 构建和运行自动化
├── README.md                      # 项目说明文档
└── .gitignore                     # Git忽略规则

通过 .gitignore 文件，我们可以忽略编译后的二进制文件和配置文件：

# 忽略二进制文件
go-tcp-proxy
*.exe

# 配置文件（如需私有配置）
config.json

Makefile 文件可以帮助我们简化编译和运行过程，内容如下：

APP_NAME := go-tcp-proxy

.PHONY: all build clean run

all: build run

build:
	@echo "==> Building $(APP_NAME)..."
	@go build -o $(APP_NAME) go_tcp_proxy.go

clean:
	@echo "==> Cleaning up..."
	@rm -f $(APP_NAME)

run:
	@echo "==> Running $(APP_NAME)..."
	./$(APP_NAME)

代码实现解析

以下是 go_tcp_proxy.go 的完整代码，包含详细注释：

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"
)

// Config 配置结构体定义
type Config struct {
	ListenPort    int    `json:"listen_port"`
	TargetHost    string `json:"target_host"`
	TargetPort    int    `json:"target_port"`
	RequestDelay  int    `json:"request_delay"`
	ResponseDelay int    `json:"response_delay"`
	LogLevel      string `json:"log_level"`
}

// 全局变量
var (
	maxConcurrentConnections = 20                        // 最大并发连接数
	semaphore                = make(chan struct{}, maxConcurrentConnections) // 控制并发连接的信号量
	configMutex              sync.Mutex                  // 保护配置的锁
	currentConfig            Config                      // 当前加载的配置
)

// 配置文件路径
func getConfigPath() string {
	homeDir, err := os.UserHomeDir()
	if err != nil {
		log.Fatal(err)
	}
	return filepath.Join(homeDir, "etc", "config.json")
}

// 加载配置
func loadConfig() Config {
	configPath := getConfigPath()
	file, err := os.Open(configPath)
	if err != nil {
		logError("加载配置失败: %v", err)
		return currentConfig
	}
	defer file.Close()

	var config Config
	decoder := json.NewDecoder(file)
	err = decoder.Decode(&config)
	if err != nil {
		logError("配置解析失败: %v", err)
		return currentConfig
	}

	// 默认日志级别为 "info"
	if config.LogLevel == "" {
		config.LogLevel = "info"
	}

	logInfo("成功加载配置: %+v", config)
	return config
}

// 日志控制函数
func logInfo(format string, v ...interface{}) {
	if currentConfig.LogLevel == "info" || currentConfig.LogLevel == "debug" {
		log.Printf("[INFO] "+format, v...)
	}
}

func logDebug(format string, v ...interface{}) {
	if currentConfig.LogLevel == "debug" {
		log.Printf("[DEBUG] "+format, v...)
	}
}

func logError(format string, v ...interface{}) {
	if currentConfig.LogLevel == "error" || currentConfig.LogLevel == "info" || currentConfig.LogLevel == "debug" {
		log.Printf("[ERROR] "+format, v...)
	}
}

func logFatal(format string, v ...interface{}) {
	log.Fatalf("[FATAL] "+format, v...)
}

// 处理客户端连接
func handleConnection(clientConn net.Conn) {
	defer clientConn.Close()
	logInfo("客户端 %v 已连接", clientConn.RemoteAddr())

	// 控制并发连接数量
	semaphore <- struct{}{}
	defer func() { <-semaphore }()

	// 加载延迟设置
	configMutex.Lock()
	requestDelay := time.Duration(currentConfig.RequestDelay) * time.Second
	responseDelay := time.Duration(currentConfig.ResponseDelay) * time.Second
	configMutex.Unlock()

	// 请求延迟处理
	if requestDelay > 0 {
		logInfo("请求延迟: %v 秒", requestDelay.Seconds())
		time.Sleep(requestDelay)
	}

	// 连接目标服务器
	targetConn, err := net.Dial("tcp", fmt.Sprintf("%s:%d", currentConfig.TargetHost, currentConfig.TargetPort))
	if err != nil {
		logError("连接目标服务器失败: %v", err)
		return
	}
	defer func() {
		logInfo("关闭与目标服务器 %v 的连接", targetConn.RemoteAddr())
		targetConn.Close()
	}()

	logInfo("成功连接到目标服务器 %s:%d", currentConfig.TargetHost, currentConfig.TargetPort)

	// 双向数据转发
	go forwardData(targetConn, clientConn, responseDelay) // 从目标到客户端
	forwardData(clientConn, targetConn, 0)                // 从客户端到目标
	logInfo("客户端 %v 的连接已关闭", clientConn.RemoteAddr())
}

// 数据转发函数
func forwardData(source, destination net.Conn, delay time.Duration) {
	buffer := make([]byte, 4096)
	for {
		n, err := source.Read(buffer)
		if n > 0 {
			if delay > 0 {
				logDebug("响应延迟: %v", delay)
				time.Sleep(delay)
			}

			_, writeErr := destination.Write(buffer[:n])
			if writeErr != nil {
				logError("数据写入目标连接失败: %v", writeErr)
				break
			}
			logDebug("成功转发数据 %d 字节", n)
		}
		if err != nil {
			if err != io.EOF && !strings.Contains(err.Error(), "use of closed network connection") {
				logError("数据转发异常: %v", err)
			}
			break
		}
	}
}

// 优雅关机函数
func gracefulShutdown(listener net.Listener, done chan bool, exit chan bool) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
	<-c

	logInfo("正在关闭服务器...")
	listener.Close()
	done <- true
	exit <- true
}

// 主函数
func main() {
	currentConfig = loadConfig()
	listenAddr := fmt.Sprintf(":%d", currentConfig.ListenPort)

	// 启动监听器
	listener, err := net.Listen("tcp", listenAddr)
	if err != nil {
		logFatal("监听失败: %v", err)
	}
	defer listener.Close()
	logInfo("代理服务器监听端口 %d，最大并发连接数为 %d", currentConfig.ListenPort, maxConcurrentConnections)

	// 优雅退出通道
	done := make(chan bool)
	exit := make(chan bool)

	go gracefulShutdown(listener, done, exit)

	for {
		select {
		case <-done:
			logInfo("服务器已成功关闭.")
			return
		default:
			clientConn, err := listener.Accept()
			if err != nil {
				select {
				case <-exit:
					return
				default:
					logError("接受连接失败: %v", err)
				}
				continue
			}
			go handleConnection(clientConn)
		}
	}
}