support udp type

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
 }
 
-func (p *ProxyClient) GetLocalConn() (c *conn.Conn, err error) {
-	c, err = conn.ConnectServer(fmt.Sprintf("%s:%d", p.LocalIp, p.LocalPort))
+// 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 (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,29 +110,27 @@ 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
 	}
 
@@ -87,12 +138,12 @@ func (p *ProxyClient) GetRemoteConn(addr string, port int64) (c *conn.Conn, err
 	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
 }

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

src/models/client/process_udp.go

@@ -0,0 +1,153 @@
+// Copyright 2016 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package client
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/fatedier/frp/src/models/msg"
+	"github.com/fatedier/frp/src/utils/conn"
+	"github.com/fatedier/frp/src/utils/pool"
+)
+
+type UdpProcesser struct {
+	tcpConn *conn.Conn
+	closeCh chan struct{}
+
+	localAddr string
+
+	// cache local udp connections
+	// key is remoteAddr
+	localUdpConns map[string]*net.UDPConn
+	mutex         sync.RWMutex
+	tcpConnMutex  sync.RWMutex
+}
+
+func NewUdpProcesser(c *conn.Conn, localIp string, localPort int64) *UdpProcesser {
+	return &UdpProcesser{
+		tcpConn:       c,
+		closeCh:       make(chan struct{}),
+		localAddr:     fmt.Sprintf("%s:%d", localIp, localPort),
+		localUdpConns: make(map[string]*net.UDPConn),
+	}
+}
+
+func (up *UdpProcesser) UpdateTcpConn(c *conn.Conn) {
+	up.tcpConnMutex.Lock()
+	defer up.tcpConnMutex.Unlock()
+	up.tcpConn = c
+}
+
+func (up *UdpProcesser) Run() {
+	go up.ReadLoop()
+}
+
+func (up *UdpProcesser) ReadLoop() {
+	var (
+		buf string
+		err error
+	)
+	for {
+		udpPacket := &msg.UdpPacket{}
+
+		// read udp package from frps
+		buf, err = up.tcpConn.ReadLine()
+		if err != nil {
+			if err == io.EOF {
+				return
+			} else {
+				continue
+			}
+		}
+		err = udpPacket.UnPack([]byte(buf))
+		if err != nil {
+			continue
+		}
+
+		// write to local udp port
+		sendConn, ok := up.GetUdpConn(udpPacket.SrcStr)
+		if !ok {
+			dstAddr, err := net.ResolveUDPAddr("udp", up.localAddr)
+			if err != nil {
+				continue
+			}
+			sendConn, err = net.DialUDP("udp", nil, dstAddr)
+			if err != nil {
+				continue
+			}
+
+			up.SetUdpConn(udpPacket.SrcStr, sendConn)
+		}
+
+		_, err = sendConn.Write(udpPacket.Content)
+		if err != nil {
+			sendConn.Close()
+			continue
+		}
+
+		if !ok {
+			go up.Forward(udpPacket, sendConn)
+		}
+	}
+}
+
+func (up *UdpProcesser) Forward(udpPacket *msg.UdpPacket, singleConn *net.UDPConn) {
+	addr := udpPacket.SrcStr
+	defer up.RemoveUdpConn(addr)
+
+	buf := pool.GetBuf(2048)
+	for {
+		singleConn.SetReadDeadline(time.Now().Add(120 * time.Second))
+		n, remoteAddr, err := singleConn.ReadFromUDP(buf)
+		if err != nil {
+			return
+		}
+
+		// forward to frps
+		forwardPacket := msg.NewUdpPacket(buf[0:n], remoteAddr, udpPacket.Src)
+		up.tcpConnMutex.RLock()
+		err = up.tcpConn.WriteString(string(forwardPacket.Pack()) + "\n")
+		up.tcpConnMutex.RUnlock()
+		if err != nil {
+			return
+		}
+	}
+}
+
+func (up *UdpProcesser) GetUdpConn(addr string) (singleConn *net.UDPConn, ok bool) {
+	up.mutex.RLock()
+	defer up.mutex.RUnlock()
+	singleConn, ok = up.localUdpConns[addr]
+	return
+}
+
+func (up *UdpProcesser) SetUdpConn(addr string, conn *net.UDPConn) {
+	up.mutex.Lock()
+	defer up.mutex.Unlock()
+	up.localUdpConns[addr] = conn
+}
+
+func (up *UdpProcesser) RemoveUdpConn(addr string) {
+	up.mutex.Lock()
+	defer up.mutex.Unlock()
+	if c, ok := up.localUdpConns[addr]; ok {
+		c.Close()
+	}
+	delete(up.localUdpConns, addr)
+}