求助: tcp 连接转发

12次阅读

共计 928 个字符,预计需要花费 3 分钟才能阅读完成。

目前我做测试想做一个特殊的转发,A 代码运行在公网监听两个端口,一个是用户请求;另一个是 B 代码的请求连接。A 接受到请求后将用户的请求转发到 B,目前表象一直卡在了 B 回传回去的问题上 (没有对 B 请求成功是否验证),请问各位大佬,这个如何操作

A 代码如下

package main

import (
	"fmt"
	"io"
	"net"
	"sync"
)

var (
	bConnMu sync.Mutex
	bConn   net.Conn
)

func handleUserRequest(userConn net.Conn) {defer func() {userConn.Close()
		fmt.Println("User connection closed")
	}()
	fmt.Println("Received user request")

	bConnMu.Lock()
	if bConn == nil {bConnMu.Unlock()
		fmt.Println("No connection to B machine")
		fmt.Fprintln(userConn, "No connection to B machine")
		return
	}
	bConnMu.Unlock()

	fmt.Println("Forwarding request to B machine")

	// 将用户请求转发给 B 机器
	err := forwardRequest(userConn, bConn)
	if err != nil {fmt.Println("Error forwarding request to B machine:", err)
		return
	}

	fmt.Println("User request completed")
}

func forwardRequest(userConn, bConn net.Conn) error {done := make(chan error, 1)

	// 将用户请求转发给 B 机器
	go func() {_, err := io.Copy(bConn, userConn)
		if err != nil {fmt.Println("Error forwarding request to B machine:", err)
		} else {fmt.Println("Finished forwarding request to B machine")
		}
		done <- err}()

	// 将 B 机器的响应转发给用户
	go func() {_, err := io.Copy(userConn, bConn)
		if err != nil {fmt.Println("Error forwarding response to user:", err)
		} else {fmt.Println("Finished forwarding response to user")
		}
		done <- err}()

	err := <-done
	if err != nil {return err}

	err = <-done
	return err
}

func handleBConnection(conn net.Conn) {fmt.Println("B machine connected")

	bConnMu.Lock()
	if bConn != nil {bConn.Close()
		fmt.Println("Closed previous connection to B machine")
	}
	bConn = conn
	bConnMu.Unlock()

	// 监听 B 机器的断开连接
	_, err := io.Copy(io.Discard, conn)
	if err != nil {fmt.Println("B machine disconnected with error:", err)
	} else {fmt.Println("B machine disconnected")
	}

	bConnMu.Lock()
	if bConn == conn {
		bConn = nil
		fmt.Println("Removed connection to B machine")
	}
	bConnMu.Unlock()
	conn.Close()}

func main11() {
	// 启动监听用户代理请求的 goroutine
	fmt.Println("A machine listening for user requests on :12345")
	go func() {listener, err := net.Listen("tcp", ":12345")
		if err != nil {fmt.Println("Failed to listen for user requests:", err)
			return
		}
		defer listener.Close()

		for {conn, err := listener.Accept()
			if err != nil {fmt.Println("Failed to accept user request:", err)
				continue
			}

			go handleUserRequest(conn)
		}
	}()

	// 启动监听 B 机器连接的 goroutine
	fmt.Println("A machine listening for B machine connection on :12346")
	listener, err := net.Listen("tcp", ":12346")
	if err != nil {fmt.Println("Failed to listen for B machine connection:", err)
		return
	}
	defer listener.Close()

	for {conn, err := listener.Accept()
		if err != nil {fmt.Println("Failed to accept B machine connection:", err)
			continue
		}

		go handleBConnection(conn)
	}
}

B 代码如下

package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"time"
)

func handleTunnel(conn net.Conn) {defer func() {conn.Close()
		fmt.Println("Tunnel connection closed")
	}()

	fmt.Println("Handling new tunnel connection")

	reader := bufio.NewReader(conn)
	for {conn.SetReadDeadline(time.Now().Add(30 * time.Second))

		request, err := http.ReadRequest(reader)
		if err != nil {
			if err == io.EOF {fmt.Println("A machine closed the connection")
				return
			}
			fmt.Println("Error reading request:", err)
			return
		}

		fmt.Printf("Received request from A machine: %s %sn", request.Method, request.URL)

		if request.Method == http.MethodConnect {fmt.Printf("Processing CONNECT request: %sn", request.URL.Host)

			targetConn, err := net.DialTimeout("tcp", request.URL.Host, 10*time.Second)
			if err != nil {log.Printf("Error connecting to target: %vn", err)
				// 修改这里来手动发送 HTTP 响应
				conn.Write([]byte("HTTP/1.1 503 Service Unavailablernrn"))
				return
			}

			conn.Write([]byte("HTTP/1.1 200 Connection Establishedrnrn"))

			// 设置通道同步 goroutines
			done := make(chan struct{})

			// 开启 goroutine 处理从 client 到目标主机的流量
			go func() {io.Copy(targetConn, conn) // 假设 conn 是从 net.Listen 获取的连接
				close(done)
			}()

			// 开启 goroutine 处理从目标主机到 client 的流量
			go func() {io.Copy(conn, targetConn)
				close(done)
			}()

			// 等待至少一个方向的流完成
			<-done

			targetConn.Close()
			conn.Close()} else {fmt.Printf("Processing normal request: %sn", request.URL)

			// 处理普通请求
			targetConn, err := net.DialTimeout("tcp", request.Host, 10*time.Second)
			if err != nil {fmt.Println("Error connecting to target host:", err)
				conn.Write([]byte("HTTP/1.1 503 Service Unavailablernrn"))
				continue
			}
			defer func() {targetConn.Close()
				fmt.Printf("Target connection to %s closedn", request.Host)
			}()

			fmt.Println("Connected to target host:", request.Host)

			// 将请求转发给目标主机
			err = request.Write(targetConn)
			if err != nil {fmt.Println("Error forwarding request to target host:", err)
				conn.Write([]byte("HTTP/1.1 503 Service Unavailablernrn"))
				continue
			}
			fmt.Println("Forwarded request to target host")

			// 将目标主机的响应转发给 A 机器
			response, err := http.ReadResponse(bufio.NewReader(targetConn), request)
			if err != nil {fmt.Println("Error reading response from target host:", err)
				conn.Write([]byte("HTTP/1.1 503 Service Unavailablernrn"))
				continue
			}
			fmt.Println("Read response from target host")

			err = response.Write(conn)
			if err != nil {fmt.Println("Error forwarding response to A machine:", err)
				continue
			}
			fmt.Println("Forwarded response to A machine")
		}
	}
}

func connectToA() {
	for {fmt.Println("Attempting to connect to A machine")
		conn, err := net.Dial("tcp", "A 机器地址:12346")
		if err != nil {fmt.Println("Failed to connect to A machine:", err)
			time.Sleep(5 * time.Second)
			continue
		}
		fmt.Println("Connected to A machine")

		// 发送心跳包
		go func() {
			for {_, err := conn.Write([]byte("ping"))
				if err != nil {fmt.Println("Failed to send heartbeat:", err)
					conn.Close()
					return
				}
				time.Sleep(5 * time.Second)
			}
		}()

		handleTunnel(conn)
	}
}

func main() {go connectToA()

	select {} // 阻塞主线程}
正文完
 0