XiaoMo/frp
1
0
Fork 0
mirror of https://github.com/fatedier/frp.git synced 2025-07-27 15:45:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

support udp

Browse Source
This commit is contained in:
fatedier
2017-03-13 02:44:47 +08:00
parent a1023fdfc2
commit 54bbfe26b0
10 changed files with 378 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
models/msg/msg.go
View File

@@ -14,7 +14,10 @@
package msg
import "reflect"
import (
"net"
"reflect"
)
const (
TypeLogin = 'o'
@@ -26,6 +29,7 @@ const (
TypeStartWorkConn = 's'
TypePing = 'h'
TypePong = '4'
TypeUdpPacket = 'u'
)
var (
@@ -46,6 +50,7 @@ func init() {
TypeMap[TypeStartWorkConn] = reflect.TypeOf(StartWorkConn{})
TypeMap[TypePing] = reflect.TypeOf(Ping{})
TypeMap[TypePong] = reflect.TypeOf(Pong{})
TypeMap[TypeUdpPacket] = reflect.TypeOf(UdpPacket{})
for k, v := range TypeMap {
TypeStringMap[v] = k
@@ -116,3 +121,9 @@ type Ping struct {
type Pong struct {
}
type UdpPacket struct {
Content string `json:"c"`
LocalAddr *net.UDPAddr `json:"l"`
RemoteAddr *net.UDPAddr `json:"r"`
}

3
models/msg/process_test.go
View File

@@ -17,6 +17,9 @@ package msg
import (
"bytes"
"encoding/binary"
"encoding/json"
"fmt"
"net"
"reflect"
"testing"

145
models/proto/udp/udp.go
View File

@@ -1,4 +1,4 @@
// Copyright 2016 fatedier, fatedier@gmail.com
// Copyright 2017 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.
@@ -16,57 +16,120 @@ package udp
import (
"encoding/base64"
"encoding/json"
"net"
"sync"
"time"
"github.com/fatedier/frp/models/msg"
"github.com/fatedier/frp/utils/errors"
"github.com/fatedier/frp/utils/pool"
)
type UdpPacket struct {
Content []byte `json:"-"`
Src *net.UDPAddr `json:"-"`
Dst *net.UDPAddr `json:"-"`
EncodeContent string `json:"content"`
SrcStr string `json:"src"`
DstStr string `json:"dst"`
func NewUdpPacket(buf []byte, laddr, raddr *net.UDPAddr) *msg.UdpPacket {
return &msg.UdpPacket{
Content: base64.StdEncoding.EncodeToString(buf),
LocalAddr: laddr,
RemoteAddr: raddr,
}
}
func NewUdpPacket(content []byte, src, dst *net.UDPAddr) *UdpPacket {
up := &UdpPacket{
Src: src,
Dst: dst,
EncodeContent: base64.StdEncoding.EncodeToString(content),
SrcStr: src.String(),
DstStr: dst.String(),
}
return up
func GetContent(m *msg.UdpPacket) (buf []byte, err error) {
buf, err = base64.StdEncoding.DecodeString(m.Content)
return
}
// parse one udp packet struct to bytes
func (up *UdpPacket) Pack() []byte {
b, _ := json.Marshal(up)
return b
func ForwardUserConn(udpConn *net.UDPConn, readCh <-chan *msg.UdpPacket, sendCh chan<- *msg.UdpPacket) {
// read
go func() {
for udpMsg := range readCh {
buf, err := GetContent(udpMsg)
if err != nil {
continue
}
udpConn.WriteToUDP(buf, udpMsg.RemoteAddr)
}
}()
// write
go func() {
buf := pool.GetBuf(1500)
defer pool.PutBuf(buf)
for {
n, remoteAddr, err := udpConn.ReadFromUDP(buf)
if err != nil {
udpConn.Close()
return
}
udpMsg := NewUdpPacket(buf[:n], nil, remoteAddr)
select {
case sendCh <- udpMsg:
default:
}
}
}()
}
// parse from bytes to UdpPacket struct
func (up *UdpPacket) UnPack(packet []byte) error {
err := json.Unmarshal(packet, &up)
if err != nil {
return err
func Forwarder(dstAddr *net.UDPAddr, readCh <-chan *msg.UdpPacket, sendCh chan<- *msg.UdpPacket) {
var (
mu sync.RWMutex
)
udpConnMap := make(map[string]*net.UDPConn)
// read from dstAddr and write to sendCh
writerFn := func(raddr *net.UDPAddr, udpConn *net.UDPConn) {
addr := raddr.String()
defer func() {
mu.Lock()
delete(udpConnMap, addr)
mu.Unlock()
}()
buf := pool.GetBuf(1500)
for {
udpConn.SetReadDeadline(time.Now().Add(30 * time.Second))
n, _, err := udpConn.ReadFromUDP(buf)
if err != nil {
return
}
udpMsg := NewUdpPacket(buf[:n], nil, raddr)
if err = errors.PanicToError(func() {
select {
case sendCh <- udpMsg:
default:
}
}); err != nil {
return
}
}
}
up.Content, err = base64.StdEncoding.DecodeString(up.EncodeContent)
if err != nil {
return err
}
// read from readCh
go func() {
for udpMsg := range readCh {
buf, err := GetContent(udpMsg)
if err != nil {
continue
}
mu.Lock()
udpConn, ok := udpConnMap[udpMsg.RemoteAddr.String()]
if !ok {
udpConn, err = net.DialUDP("udp", nil, dstAddr)
if err != nil {
continue
}
udpConnMap[udpMsg.RemoteAddr.String()] = udpConn
}
mu.Unlock()
up.Src, err = net.ResolveUDPAddr("udp", up.SrcStr)
if err != nil {
return err
}
_, err = udpConn.Write(buf)
if err != nil {
udpConn.Close()
}
up.Dst, err = net.ResolveUDPAddr("udp", up.DstStr)
if err != nil {
return err
}
return nil
if !ok {
go writerFn(udpMsg.RemoteAddr, udpConn)
}
}
}()
}

50
models/proto/udp/udp_test.go
View File

@@ -1,50 +0,0 @@
// 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 udp
import (
"net"
"testing"
"github.com/stretchr/testify/assert"
)
var (
content string = "udp packet test"
src string = "1.1.1.1:1000"
dst string = "2.2.2.2:2000"
udpMsg *UdpPacket
)
func init() {
srcAddr, _ := net.ResolveUDPAddr("udp", src)
dstAddr, _ := net.ResolveUDPAddr("udp", dst)
udpMsg = NewUdpPacket([]byte(content), srcAddr, dstAddr)
}
func TestPack(t *testing.T) {
assert := assert.New(t)
msg := udpMsg.Pack()
assert.Equal(string(msg), `{"content":"dWRwIHBhY2tldCB0ZXN0","src":"1.1.1.1:1000","dst":"2.2.2.2:2000"}`)
}
func TestUnpack(t *testing.T) {
assert := assert.New(t)
udpMsg.UnPack([]byte(`{"content":"dWRwIHBhY2tldCB0ZXN0","src":"1.1.1.1:1000","dst":"2.2.2.2:2000"}`))
assert.Equal(content, string(udpMsg.Content))
assert.Equal(src, udpMsg.Src.String())
assert.Equal(dst, udpMsg.Dst.String())
}