frp/utils/net/udp.go

259 lines
4.8 KiB
Go
Raw Normal View History

2017-03-09 02:03:47 +08:00
// 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.
// 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 net
import (
"fmt"
"io"
"net"
"sync"
"time"
2018-05-08 02:13:30 +08:00
"github.com/fatedier/golib/pool"
2017-03-09 02:03:47 +08:00
)
2020-05-24 17:48:37 +08:00
type UDPPacket struct {
2017-03-09 02:03:47 +08:00
Buf []byte
LocalAddr net.Addr
RemoteAddr net.Addr
}
2020-05-24 17:48:37 +08:00
type FakeUDPConn struct {
l *UDPListener
2017-03-09 02:03:47 +08:00
localAddr net.Addr
remoteAddr net.Addr
packets chan []byte
closeFlag bool
lastActive time.Time
mu sync.RWMutex
}
2020-05-24 17:48:37 +08:00
func NewFakeUDPConn(l *UDPListener, laddr, raddr net.Addr) *FakeUDPConn {
fc := &FakeUDPConn{
2017-03-09 02:03:47 +08:00
l: l,
localAddr: laddr,
remoteAddr: raddr,
packets: make(chan []byte, 20),
}
go func() {
for {
time.Sleep(5 * time.Second)
fc.mu.RLock()
if time.Now().Sub(fc.lastActive) > 10*time.Second {
fc.mu.RUnlock()
fc.Close()
break
}
fc.mu.RUnlock()
}
}()
return fc
}
2020-05-24 17:48:37 +08:00
func (c *FakeUDPConn) putPacket(content []byte) {
2017-03-09 02:03:47 +08:00
defer func() {
if err := recover(); err != nil {
}
}()
select {
case c.packets <- content:
default:
}
}
2020-05-24 17:48:37 +08:00
func (c *FakeUDPConn) Read(b []byte) (n int, err error) {
2017-03-09 02:03:47 +08:00
content, ok := <-c.packets
if !ok {
return 0, io.EOF
}
c.mu.Lock()
c.lastActive = time.Now()
c.mu.Unlock()
if len(b) < len(content) {
n = len(b)
} else {
n = len(content)
}
copy(b, content)
return n, nil
}
2020-05-24 17:48:37 +08:00
func (c *FakeUDPConn) Write(b []byte) (n int, err error) {
2017-03-09 02:03:47 +08:00
c.mu.RLock()
if c.closeFlag {
c.mu.RUnlock()
return 0, io.ErrClosedPipe
}
c.mu.RUnlock()
2020-05-24 17:48:37 +08:00
packet := &UDPPacket{
2017-03-09 02:03:47 +08:00
Buf: b,
LocalAddr: c.localAddr,
RemoteAddr: c.remoteAddr,
}
2020-05-24 17:48:37 +08:00
c.l.writeUDPPacket(packet)
2017-03-09 02:03:47 +08:00
c.mu.Lock()
c.lastActive = time.Now()
c.mu.Unlock()
return len(b), nil
}
2020-05-24 17:48:37 +08:00
func (c *FakeUDPConn) Close() error {
2017-03-09 02:03:47 +08:00
c.mu.Lock()
defer c.mu.Unlock()
if !c.closeFlag {
c.closeFlag = true
close(c.packets)
}
return nil
}
2020-05-24 17:48:37 +08:00
func (c *FakeUDPConn) IsClosed() bool {
2017-03-09 02:03:47 +08:00
c.mu.RLock()
defer c.mu.RUnlock()
return c.closeFlag
}
2020-05-24 17:48:37 +08:00
func (c *FakeUDPConn) LocalAddr() net.Addr {
2017-03-09 02:03:47 +08:00
return c.localAddr
}
2020-05-24 17:48:37 +08:00
func (c *FakeUDPConn) RemoteAddr() net.Addr {
2017-03-09 02:03:47 +08:00
return c.remoteAddr
}
2020-05-24 17:48:37 +08:00
func (c *FakeUDPConn) SetDeadline(t time.Time) error {
2017-03-09 02:03:47 +08:00
return nil
}
2020-05-24 17:48:37 +08:00
func (c *FakeUDPConn) SetReadDeadline(t time.Time) error {
2017-03-09 02:03:47 +08:00
return nil
}
2020-05-24 17:48:37 +08:00
func (c *FakeUDPConn) SetWriteDeadline(t time.Time) error {
2017-03-09 02:03:47 +08:00
return nil
}
2020-05-24 17:48:37 +08:00
type UDPListener struct {
2019-10-12 20:13:12 +08:00
addr net.Addr
acceptCh chan net.Conn
2020-05-24 17:48:37 +08:00
writeCh chan *UDPPacket
2017-03-09 02:03:47 +08:00
readConn net.Conn
closeFlag bool
2020-05-24 17:48:37 +08:00
fakeConns map[string]*FakeUDPConn
2017-03-09 02:03:47 +08:00
}
2020-05-24 17:48:37 +08:00
func ListenUDP(bindAddr string, bindPort int) (l *UDPListener, err error) {
2017-03-09 02:03:47 +08:00
udpAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", bindAddr, bindPort))
if err != nil {
return l, err
}
readConn, err := net.ListenUDP("udp", udpAddr)
2020-05-24 17:48:37 +08:00
l = &UDPListener{
2019-10-12 20:13:12 +08:00
addr: udpAddr,
acceptCh: make(chan net.Conn),
2020-05-24 17:48:37 +08:00
writeCh: make(chan *UDPPacket, 1000),
fakeConns: make(map[string]*FakeUDPConn),
2017-03-09 02:03:47 +08:00
}
// for reading
go func() {
for {
buf := pool.GetBuf(1450)
n, remoteAddr, err := readConn.ReadFromUDP(buf)
if err != nil {
2019-10-12 20:13:12 +08:00
close(l.acceptCh)
2017-03-09 02:03:47 +08:00
close(l.writeCh)
return
}
fakeConn, exist := l.fakeConns[remoteAddr.String()]
if !exist || fakeConn.IsClosed() {
2020-05-24 17:48:37 +08:00
fakeConn = NewFakeUDPConn(l, l.Addr(), remoteAddr)
2017-03-09 02:03:47 +08:00
l.fakeConns[remoteAddr.String()] = fakeConn
}
fakeConn.putPacket(buf[:n])
2019-10-12 20:13:12 +08:00
l.acceptCh <- fakeConn
2017-03-09 02:03:47 +08:00
}
}()
// for writing
go func() {
for {
packet, ok := <-l.writeCh
if !ok {
return
}
if addr, ok := packet.RemoteAddr.(*net.UDPAddr); ok {
readConn.WriteToUDP(packet.Buf, addr)
}
}
}()
return
}
2020-05-24 17:48:37 +08:00
func (l *UDPListener) writeUDPPacket(packet *UDPPacket) (err error) {
2017-03-09 02:03:47 +08:00
defer func() {
2017-03-13 02:44:47 +08:00
if errRet := recover(); errRet != nil {
err = fmt.Errorf("udp write closed listener")
2017-03-09 02:03:47 +08:00
}
}()
l.writeCh <- packet
2017-03-13 02:44:47 +08:00
return
}
2020-05-24 17:48:37 +08:00
func (l *UDPListener) WriteMsg(buf []byte, remoteAddr *net.UDPAddr) (err error) {
2017-03-13 02:44:47 +08:00
// only set remote addr here
2020-05-24 17:48:37 +08:00
packet := &UDPPacket{
2017-03-13 02:44:47 +08:00
Buf: buf,
RemoteAddr: remoteAddr,
}
2020-05-24 17:48:37 +08:00
err = l.writeUDPPacket(packet)
2017-03-13 02:44:47 +08:00
return
2017-03-09 02:03:47 +08:00
}
2020-05-24 17:48:37 +08:00
func (l *UDPListener) Accept() (net.Conn, error) {
2019-10-12 20:13:12 +08:00
conn, ok := <-l.acceptCh
2017-03-09 02:03:47 +08:00
if !ok {
return conn, fmt.Errorf("channel for udp listener closed")
}
return conn, nil
}
2020-05-24 17:48:37 +08:00
func (l *UDPListener) Close() error {
2017-03-09 02:03:47 +08:00
if !l.closeFlag {
l.closeFlag = true
2020-09-07 14:57:23 +08:00
if l.readConn != nil {
l.readConn.Close()
}
2017-03-09 02:03:47 +08:00
}
return nil
}
2019-10-12 20:13:12 +08:00
2020-05-24 17:48:37 +08:00
func (l *UDPListener) Addr() net.Addr {
2019-10-12 20:13:12 +08:00
return l.addr
}