support udp type

fatedier · Dec 18, 2016 · f2999e3 · f2999e3
1 parent adcb2c1
commit f2999e3
Show file tree

Hide file tree

Showing 18 changed files with 556 additions and 76 deletions.
diff --git a/.travis.yml b/.travis.yml
@@ -3,8 +3,8 @@ language: go
 
 go:
     - 1.5.4
-    - 1.6.3
-    - 1.7
+    - 1.6.4
+    - 1.7.4
 
 install:
     - make

diff --git a/conf/frpc.ini b/conf/frpc.ini
@@ -29,6 +29,13 @@ use_gzip = false
 # connections will be established in advance, default value is zero
 pool_count = 10
 
+[dns]
+type = udp
+local_ip = 127.0.0.1
+local_port = 53
+use_encryption = true
+use_gzip = true
+
 # Resolve your domain names to [server_addr] so you can use http://web01.yourdomain.com to browse web01 and http://web02.yourdomain.com to browse web02, the domains are set in frps.ini
 [web01]
 type = http

diff --git a/conf/frps.ini b/conf/frps.ini
@@ -34,6 +34,12 @@ auth_token = 123
 bind_addr = 0.0.0.0
 listen_port = 6000
 
+[dns]
+type = udp
+auth_token = 123
+bind_addr = 0.0.0.0
+listen_port = 53
+
 [web01]
 # if type equals http, vhost_http_port must be set
 type = http

diff --git a/src/cmd/frpc/control.go b/src/cmd/frpc/control.go
@@ -120,7 +120,7 @@ func msgSender(cli *client.ProxyClient, c *conn.Conn, msgSendChan chan interface
 		}
 
 		buf, _ := json.Marshal(msg)
-		err := c.Write(string(buf) + "\n")
+		err := c.WriteString(string(buf) + "\n")
 		if err != nil {
 			log.Warn("ProxyName [%s], write to server error, proxy exit", cli.Name)
 			c.Close()
@@ -165,7 +165,7 @@ func loginToServer(cli *client.ProxyClient) (c *conn.Conn, err error) {
 	}
 
 	buf, _ := json.Marshal(req)
-	err = c.Write(string(buf) + "\n")
+	err = c.WriteString(string(buf) + "\n")
 	if err != nil {
 		log.Error("ProxyName [%s], write to server error, %v", cli.Name, err)
 		return
@@ -190,6 +190,12 @@ func loginToServer(cli *client.ProxyClient) (c *conn.Conn, err error) {
 	}
 
 	log.Info("ProxyName [%s], connect to server [%s:%d] success!", cli.Name, client.ServerAddr, client.ServerPort)
+
+	if cli.Type == "udp" {
+		// we only need one udp work connection
+		// all udp messages will be forwarded throngh this connection
+		go cli.StartUdpTunnelOnce(client.ServerAddr, client.ServerPort)
+	}
 	return
 }
 

diff --git a/src/cmd/frps/control.go b/src/cmd/frps/control.go
@@ -71,14 +71,14 @@ func controlWorker(c *conn.Conn) {
 	// login when type is NewCtlConn or NewWorkConn
 	ret, info := doLogin(cliReq, c)
 	// if login type is NewWorkConn, nothing will be send to frpc
-	if cliReq.Type != consts.NewWorkConn {
+	if cliReq.Type == consts.NewCtlConn {
 		cliRes := &msg.ControlRes{
 			Type: consts.NewCtlConnRes,
 			Code: ret,
 			Msg:  info,
 		}
 		byteBuf, _ := json.Marshal(cliRes)
-		err = c.Write(string(byteBuf) + "\n")
+		err = c.WriteString(string(byteBuf) + "\n")
 		if err != nil {
 			log.Warn("ProxyName [%s], write to client error, proxy exit", cliReq.ProxyName)
 			return
@@ -144,9 +144,11 @@ func msgReader(s *server.ProxyServer, c *conn.Conn, msgSendChan chan interface{}
 		if err != nil {
 			if err == io.EOF {
 				log.Warn("ProxyName [%s], client is dead!", s.Name)
+				s.Close()
 				return err
 			} else if c == nil || c.IsClosed() {
 				log.Warn("ProxyName [%s], client connection is closed", s.Name)
+				s.Close()
 				return err
 			}
 			log.Warn("ProxyName [%s], read error: %v", s.Name, err)
@@ -183,7 +185,7 @@ func msgSender(s *server.ProxyServer, c *conn.Conn, msgSendChan chan interface{}
 		}
 
 		buf, _ := json.Marshal(msg)
-		err := c.Write(string(buf) + "\n")
+		err := c.WriteString(string(buf) + "\n")
 		if err != nil {
 			log.Warn("ProxyName [%s], write to client error, proxy exit", s.Name)
 			s.Close()
@@ -193,6 +195,9 @@ func msgSender(s *server.ProxyServer, c *conn.Conn, msgSendChan chan interface{}
 }
 
 // if success, ret equals 0, otherwise greater than 0
+// NewCtlConn
+// NewWorkConn
+// NewWorkConnUdp
 func doLogin(req *msg.ControlReq, c *conn.Conn) (ret int64, info string) {
 	ret = 1
 	// check if PrivilegeMode is enabled
@@ -325,6 +330,13 @@ func doLogin(req *msg.ControlReq, c *conn.Conn) (ret int64, info string) {
 		}
 		// the connection will close after join over
 		s.RegisterNewWorkConn(c)
+	} else if req.Type == consts.NewWorkConnUdp {
+		// work conn for udp
+		if s.Status != consts.Working {
+			log.Warn("ProxyName [%s], is not working when it gets one new work connnection for udp", req.ProxyName)
+			return
+		}
+		s.RegisterNewWorkConnUdp(c)
 	} else {
 		info = fmt.Sprintf("Unsupport login message type [%d]", req.Type)
 		log.Warn("Unsupport login message type [%d]", req.Type)

diff --git a/src/models/client/client.go b/src/models/client/client.go
@@ -17,6 +17,7 @@ package client
 import (
 	"encoding/json"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/fatedier/frp/src/models/config"
@@ -34,19 +35,71 @@ type ProxyClient struct {
 
 	RemotePort    int64
 	CustomDomains []string
+
+	udpTunnel *conn.Conn
+	once      sync.Once
+}
+
+// if proxy type is udp, keep a tcp connection for transferring udp packages
+func (pc *ProxyClient) StartUdpTunnelOnce(addr string, port int64) {
+	pc.once.Do(func() {
+		var err error
+		var c *conn.Conn
+		udpProcessor := NewUdpProcesser(nil, pc.LocalIp, pc.LocalPort)
+		for {
+			if pc.udpTunnel == nil || pc.udpTunnel.IsClosed() {
+				if HttpProxy == "" {
+					c, err = conn.ConnectServer(fmt.Sprintf("%s:%d", addr, port))
+				} else {
+					c, err = conn.ConnectServerByHttpProxy(HttpProxy, fmt.Sprintf("%s:%d", addr, port))
+				}
+				if err != nil {
+					log.Error("ProxyName [%s], udp tunnel connect to server [%s:%d] error, %v", pc.Name, addr, port, err)
+					time.Sleep(5 * time.Second)
+					continue
+				}
+
+				nowTime := time.Now().Unix()
+				req := &msg.ControlReq{
+					Type:          consts.NewWorkConnUdp,
+					ProxyName:     pc.Name,
+					PrivilegeMode: pc.PrivilegeMode,
+					Timestamp:     nowTime,
+				}
+				if pc.PrivilegeMode == true {
+					req.PrivilegeKey = pcrypto.GetAuthKey(pc.Name + PrivilegeToken + fmt.Sprintf("%d", nowTime))
+				} else {
+					req.AuthKey = pcrypto.GetAuthKey(pc.Name + pc.AuthToken + fmt.Sprintf("%d", nowTime))
+				}
+
+				buf, _ := json.Marshal(req)
+				err = c.WriteString(string(buf) + "\n")
+				if err != nil {
+					log.Error("ProxyName [%s], udp tunnel write to server error, %v", pc.Name, err)
+					c.Close()
+					time.Sleep(1 * time.Second)
+					continue
+				}
+				pc.udpTunnel = c
+				udpProcessor.UpdateTcpConn(pc.udpTunnel)
+				udpProcessor.Run()
+			}
+			time.Sleep(1 * time.Second)
+		}
+	})
 }
 
-func (p *ProxyClient) GetLocalConn() (c *conn.Conn, err error) {
-	c, err = conn.ConnectServer(fmt.Sprintf("%s:%d", p.LocalIp, p.LocalPort))
+func (pc *ProxyClient) GetLocalConn() (c *conn.Conn, err error) {
+	c, err = conn.ConnectServer(fmt.Sprintf("%s:%d", pc.LocalIp, pc.LocalPort))
 	if err != nil {
-		log.Error("ProxyName [%s], connect to local port error, %v", p.Name, err)
+		log.Error("ProxyName [%s], connect to local port error, %v", pc.Name, err)
 	}
 	return
 }
 
-func (p *ProxyClient) GetRemoteConn(addr string, port int64) (c *conn.Conn, err error) {
+func (pc *ProxyClient) GetRemoteConn(addr string, port int64) (c *conn.Conn, err error) {
 	defer func() {
-		if err != nil {
+		if err != nil && c != nil {
 			c.Close()
 		}
 	}()
@@ -57,42 +110,40 @@ func (p *ProxyClient) GetRemoteConn(addr string, port int64) (c *conn.Conn, err
 		c, err = conn.ConnectServerByHttpProxy(HttpProxy, fmt.Sprintf("%s:%d", addr, port))
 	}
 	if err != nil {
-		log.Error("ProxyName [%s], connect to server [%s:%d] error, %v", p.Name, addr, port, err)
+		log.Error("ProxyName [%s], connect to server [%s:%d] error, %v", pc.Name, addr, port, err)
 		return
 	}
 
 	nowTime := time.Now().Unix()
 	req := &msg.ControlReq{
 		Type:          consts.NewWorkConn,
-		ProxyName:     p.Name,
-		PrivilegeMode: p.PrivilegeMode,
+		ProxyName:     pc.Name,
+		PrivilegeMode: pc.PrivilegeMode,
 		Timestamp:     nowTime,
 	}
-	if p.PrivilegeMode == true {
-		privilegeKey := pcrypto.GetAuthKey(p.Name + PrivilegeToken + fmt.Sprintf("%d", nowTime))
-		req.PrivilegeKey = privilegeKey
+	if pc.PrivilegeMode == true {
+		req.PrivilegeKey = pcrypto.GetAuthKey(pc.Name + PrivilegeToken + fmt.Sprintf("%d", nowTime))
 	} else {
-		authKey := pcrypto.GetAuthKey(p.Name + p.AuthToken + fmt.Sprintf("%d", nowTime))
-		req.AuthKey = authKey
+		req.AuthKey = pcrypto.GetAuthKey(pc.Name + pc.AuthToken + fmt.Sprintf("%d", nowTime))
 	}
 
 	buf, _ := json.Marshal(req)
-	err = c.Write(string(buf) + "\n")
+	err = c.WriteString(string(buf) + "\n")
 	if err != nil {
-		log.Error("ProxyName [%s], write to server error, %v", p.Name, err)
+		log.Error("ProxyName [%s], write to server error, %v", pc.Name, err)
 		return
 	}
 
 	err = nil
 	return
 }
 
-func (p *ProxyClient) StartTunnel(serverAddr string, serverPort int64) (err error) {
-	localConn, err := p.GetLocalConn()
+func (pc *ProxyClient) StartTunnel(serverAddr string, serverPort int64) (err error) {
+	localConn, err := pc.GetLocalConn()
 	if err != nil {
 		return
 	}
-	remoteConn, err := p.GetRemoteConn(serverAddr, serverPort)
+	remoteConn, err := pc.GetRemoteConn(serverAddr, serverPort)
 	if err != nil {
 		return
 	}
@@ -101,7 +152,7 @@ func (p *ProxyClient) StartTunnel(serverAddr string, serverPort int64) (err erro
 	log.Debug("Join two connections, (l[%s] r[%s]) (l[%s] r[%s])", localConn.GetLocalAddr(), localConn.GetRemoteAddr(),
 		remoteConn.GetLocalAddr(), remoteConn.GetRemoteAddr())
 	needRecord := false
-	go msg.JoinMore(localConn, remoteConn, p.BaseConf, needRecord)
+	go msg.JoinMore(localConn, remoteConn, pc.BaseConf, needRecord)
 
 	return nil
 }
diff --git a/src/models/client/config.go b/src/models/client/config.go
@@ -130,7 +130,7 @@ func LoadConf(confFile string) (err error) {
 			proxyClient.Type = "tcp"
 			tmpStr, ok = section["type"]
 			if ok {
-				if tmpStr != "tcp" && tmpStr != "http" && tmpStr != "https" {
+				if tmpStr != "tcp" && tmpStr != "http" && tmpStr != "https" && tmpStr != "udp" {
 					return fmt.Errorf("Parse conf error: proxy [%s] type error", proxyClient.Name)
 				}
 				proxyClient.Type = tmpStr
-Original file line number
+Diff line change
@@ Expand Up / @@ -3,8 +3,8 @@ language: go @@
     go:
         - 1.5.4
-        - 1.6.3
-        - 1.7
+        - 1.6.4
+        - 1.7.4
     install:
         - make
@@ Expand Down @@