golang实现简单的TCP代理

TCP Proxy

package mainimport ("io""log""net""sync"
)// go实现简单的TCP代理// TCPProxy 是一个 TCP 代理
type TCPProxy struct {listenAddr stringtargetAddr string
}// NewTCPProxy 创建一个新的 TCP 代理
func NewTCPProxy(listenAddr, targetAddr string) *TCPProxy {return &TCPProxy{listenAddr: listenAddr,targetAddr: targetAddr,}
}// Serve 开始监听并代理 TCP 连接
func (p *TCPProxy) Serve() error {// 监听 TCP 连接listener, err := net.Listen("tcp", p.listenAddr)if err != nil {return err}defer listener.Close()log.Printf("tcp proxy listening at: %s\n", listener.Addr())// 处理 TCP 连接for {conn, err := listener.Accept()if err != nil {log.Printf("listener accept error: %s\n", err)continue}go p.handleConnection(conn)}
}// handleConnection 处理 TCP 连接
func (p *TCPProxy) handleConnection(conn net.Conn) {defer conn.Close() // 关闭代理连接log.Printf("accepted new connection %s\n", conn.RemoteAddr())// 创建目标连接targetConn, err := net.Dial("tcp", p.targetAddr)if err != nil {log.Printf("dial target addr: %s error: %s\n", p.targetAddr, err)return}defer targetConn.Close() // 关闭目标连接log.Printf("tcp proxy client addr: %s\n", targetConn.LocalAddr())// 处理 TCP 流 ，在客户端和目标服务器之间双向转发数据wg := sync.WaitGroup{}wg.Add(2)go func() {defer wg.Done()n, err := io.Copy(conn, targetConn)if err != nil {log.Printf("copy bytes from %s to %s error: %s", conn.RemoteAddr(), targetConn.RemoteAddr(), err)return}log.Printf("copied %d bytes from %s to %s", n, conn.RemoteAddr(), targetConn.RemoteAddr())}()go func() {defer wg.Done()n, err := io.Copy(targetConn, conn)if err != nil {log.Printf("copy bytes from %s to %s error: %s", targetConn.RemoteAddr(), conn.RemoteAddr(), err)return}log.Printf("copied %d bytes from %s to %s", n, targetConn.RemoteAddr(), conn.RemoteAddr())}()// 等待两个复制操作完成wg.Wait()
}func main() {proxy := NewTCPProxy("127.0.0.1:8080", "127.0.0.1:8081")if err := proxy.Serve(); err != nil {log.Fatal(err)}
}

Server

package mainimport ("bufio""log""net"
)// TCP Server端测试
// handle 处理函数
func handle(conn net.Conn) {defer conn.Close() // 关闭连接reader := bufio.NewReader(conn)for {var buf [512]byten, err := reader.Read(buf[:]) // 读取数据if err != nil {log.Printf("read from client err: %s\n", err)break}log.Printf("read %d bytes from client\n", n)recvStr := string(buf[:n])log.Printf("received data from client: %s\n", recvStr)n, err = conn.Write([]byte(recvStr)) // 发送数据if err != nil {log.Printf("write data to client err: %s\n", err)return}log.Printf("write %d bytes to client\n", n)}
}func main() {listener, err := net.Listen("tcp", "127.0.0.1:8081")if err != nil {log.Fatal(err)}log.Printf("server listening at: %s\n", listener.Addr())for {conn, err := listener.Accept() // 监听客户端的连接请求if err != nil {log.Printf("listener accept error: %s\n", err)continue}go handle(conn) // 启动一个goroutine来处理客户端的连接请求}
}

Client

package mainimport ("bufio""log""net""os""strings"
)// TCP 客户端
func main() {conn, err := net.Dial("tcp", "127.0.0.1:8080")if err != nil {log.Fatal(err)}log.Printf("client addr: %s\n", conn.LocalAddr())defer conn.Close() // 关闭TCP连接inputReader := bufio.NewReader(os.Stdin)for {input, _ := inputReader.ReadString('\n') // 读取用户输入inputInfo := strings.Trim(input, "\r\n")if strings.ToUpper(inputInfo) == "Q" { // 如果输入q就退出return}n, err := conn.Write([]byte(inputInfo)) // 发送数据if err != nil {log.Printf("write data to server err: %s\n", err)return}log.Printf("write %d bytes to server\n", n)buf := [512]byte{}n, err = conn.Read(buf[:])if err != nil {log.Printf("read data from server err: %s\n", err)return}log.Printf("read %d bytes from server\n", n)log.Printf("received data from server: %s\n", string(buf[:n]))}
}