一个简单的tcp代理

 1 // TcpProxy project main.go
 2 package main
 3 
 4 import (
 5     "fmt"
 6     "net"
 7     "os"
 8     "regexp"
 9     "strconv"
10 )
11 
12 func main() {
13     if len(os.Args) < 4 {
14         fmt.Println("missing message!")
15         return
16     }
17     ip := os.Args[1]
18     port, err := strconv.Atoi(os.Args[2])
19     if err != nil {
20         fmt.Println("error happened ,exit")
21         return
22     }
23     addr := os.Args[3]
24     host := "Host: " + addr
25 
26     Service(ip, port, addr, host)
27 }
28 
29 func Service(ip string, port int, dstaddr string, dsthost string) {
30     // listen and accept
31     listen, err := net.ListenTCP("tcp", &net.TCPAddr{net.ParseIP(ip), port, ""})
32     if err != nil {
33         fmt.Println("listen error: ", err.Error())
34         return
35     }
36     fmt.Println("init done...")
37 
38     for {
39         client, err := listen.AcceptTCP()
40         if err != nil {
41             fmt.Println("accept error: ", err.Error())
42             continue
43         }
44         go Channal(client, dstaddr, dsthost)
45     }
46 }
47 
48 func Channal(client *net.TCPConn, addr string, rhost string) {
49 
50     tcpAddr, _ := net.ResolveTCPAddr("tcp4", addr)
51     conn, err := net.DialTCP("tcp", nil, tcpAddr)
52     if err != nil {
53         fmt.Println("connection error: ", err.Error())
54         client.Close()
55         return
56     }
57 
58     go ReadRequest(client, conn, rhost)
59     ReadResponse(conn, client)
60 }
61 
62 func ReadRequest(lconn *net.TCPConn, rconn *net.TCPConn, dsthost string) {
63     for {
64         buf := make([]byte, 10240)
65         n, err := lconn.Read(buf)
66         if err != nil {
67             break
68         }
69 
70         mesg := changeHost(string(buf[:n]), dsthost)
71         // print request
72         fmt.Println(mesg)
73         rconn.Write([]byte(mesg))
74     }
75     lconn.Close()
76 }
77 
78 func ReadResponse(lconn *net.TCPConn, rconn *net.TCPConn) {
79     for {
80         buf := make([]byte, 10240)
81         n, err := lconn.Read(buf)
82         if err != nil {
83             break
84         }
85 
86         fmt.Println(string(buf[:n]))
87         rconn.Write(buf[:n])
88     }
89     lconn.Close()
90 }
91 
92 func changeHost(request string, newhost string) string {
93     reg := regexp.MustCompile(`Host[^\r\n]+`)
94     return reg.ReplaceAllString(request, newhost)
95 }