mirror of
https://github.com/fatedier/frp.git
synced 2026-01-11 22:23:12 +00:00
start refactoring
This commit is contained in:
fatedier
32
utils/net/conn.go
Normal file
32
utils/net/conn.go
Normal file
@@ -0,0 +1,32 @@
|
||||
// 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 net
|
||||
|
||||
import (
|
||||
"net"
|
||||
|
||||
"github.com/fatedier/frp/utils/log"
|
||||
)
|
||||
|
||||
// Conn is the interface of connections used in frp.
|
||||
type Conn interface {
|
||||
net.Conn
|
||||
log.Logger
|
||||
}
|
||||
|
||||
type Listener interface {
|
||||
Accept() (Conn, error)
|
||||
Close() error
|
||||
}
|
||||
162
utils/net/tcp.go
Normal file
162
utils/net/tcp.go
Normal file
@@ -0,0 +1,162 @@
|
||||
// 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 net
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
||||
"github.com/fatedier/frp/utils/log"
|
||||
)
|
||||
|
||||
type TcpListener struct {
|
||||
net.Addr
|
||||
listener net.Listener
|
||||
accept chan Conn
|
||||
closeFlag bool
|
||||
}
|
||||
|
||||
func ListenTcp(bindAddr string, bindPort int64) (l *TcpListener, err error) {
|
||||
tcpAddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%d", bindAddr, bindPort))
|
||||
if err != nil {
|
||||
return l, err
|
||||
}
|
||||
listener, err := net.ListenTCP("tcp", tcpAddr)
|
||||
if err != nil {
|
||||
return l, err
|
||||
}
|
||||
|
||||
l = &TcpListener{
|
||||
Addr: listener.Addr(),
|
||||
listener: listener,
|
||||
accept: make(chan Conn),
|
||||
closeFlag: false,
|
||||
}
|
||||
|
||||
go func() {
|
||||
for {
|
||||
conn, err := listener.AcceptTCP()
|
||||
if err != nil {
|
||||
if l.closeFlag {
|
||||
close(l.accept)
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
c := NewTcpConn(conn)
|
||||
l.accept <- c
|
||||
}
|
||||
}()
|
||||
return l, err
|
||||
}
|
||||
|
||||
// Wait util get one new connection or listener is closed
|
||||
// if listener is closed, err returned.
|
||||
func (l *TcpListener) Accept() (Conn, error) {
|
||||
conn, ok := <-l.accept
|
||||
if !ok {
|
||||
return conn, fmt.Errorf("channel for tcp listener closed")
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func (l *TcpListener) Close() error {
|
||||
if !l.closeFlag {
|
||||
l.closeFlag = true
|
||||
l.listener.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Wrap for TCPConn.
|
||||
type TcpConn struct {
|
||||
net.Conn
|
||||
log.Logger
|
||||
}
|
||||
|
||||
func NewTcpConn(conn *net.TCPConn) (c *TcpConn) {
|
||||
c = &TcpConn{
|
||||
Conn: conn,
|
||||
Logger: log.NewPrefixLogger(""),
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func ConnectTcpServer(addr string) (c Conn, err error) {
|
||||
servertAddr, err := net.ResolveTCPAddr("tcp", addr)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
conn, err := net.DialTCP("tcp", nil, servertAddr)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
c = NewTcpConn(conn)
|
||||
return
|
||||
}
|
||||
|
||||
// ConnectTcpServerByHttpProxy try to connect remote server by http proxy.
|
||||
// If httpProxy is empty, it will connect server directly.
|
||||
func ConnectTcpServerByHttpProxy(httpProxy string, serverAddr string) (c Conn, err error) {
|
||||
if httpProxy == "" {
|
||||
return ConnectTcpServer(serverAddr)
|
||||
}
|
||||
|
||||
var proxyUrl *url.URL
|
||||
if proxyUrl, err = url.Parse(httpProxy); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
var proxyAuth string
|
||||
if proxyUrl.User != nil {
|
||||
proxyAuth = "Basic " + base64.StdEncoding.EncodeToString([]byte(proxyUrl.User.String()))
|
||||
}
|
||||
|
||||
if proxyUrl.Scheme != "http" {
|
||||
err = fmt.Errorf("Proxy URL scheme must be http, not [%s]", proxyUrl.Scheme)
|
||||
return
|
||||
}
|
||||
|
||||
if c, err = ConnectTcpServer(proxyUrl.Host); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
req, err := http.NewRequest("CONNECT", "http://"+serverAddr, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if proxyAuth != "" {
|
||||
req.Header.Set("Proxy-Authorization", proxyAuth)
|
||||
}
|
||||
req.Header.Set("User-Agent", "Mozilla/5.0")
|
||||
req.Write(c)
|
||||
|
||||
resp, err := http.ReadResponse(bufio.NewReader(c), req)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode != 200 {
|
||||
err = fmt.Errorf("ConnectTcpServer using proxy error, StatusCode [%d]", resp.StatusCode)
|
||||
return
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
243
utils/net/udp.go
Normal file
243
utils/net/udp.go
Normal file
@@ -0,0 +1,243 @@
|
||||
// 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"
|
||||
|
||||
flog "github.com/fatedier/frp/utils/log"
|
||||
"github.com/fatedier/frp/utils/pool"
|
||||
)
|
||||
|
||||
type UdpPacket struct {
|
||||
Buf []byte
|
||||
LocalAddr net.Addr
|
||||
RemoteAddr net.Addr
|
||||
}
|
||||
|
||||
type FakeUdpConn struct {
|
||||
flog.Logger
|
||||
l *UdpListener
|
||||
|
||||
localAddr net.Addr
|
||||
remoteAddr net.Addr
|
||||
packets chan []byte
|
||||
closeFlag bool
|
||||
|
||||
lastActive time.Time
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func NewFakeUdpConn(l *UdpListener, laddr, raddr net.Addr) *FakeUdpConn {
|
||||
fc := &FakeUdpConn{
|
||||
Logger: flog.NewPrefixLogger(""),
|
||||
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
|
||||
}
|
||||
|
||||
func (c *FakeUdpConn) putPacket(content []byte) {
|
||||
defer func() {
|
||||
if err := recover(); err != nil {
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case c.packets <- content:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (c *FakeUdpConn) Read(b []byte) (n int, err error) {
|
||||
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
|
||||
}
|
||||
|
||||
func (c *FakeUdpConn) Write(b []byte) (n int, err error) {
|
||||
c.mu.RLock()
|
||||
if c.closeFlag {
|
||||
c.mu.RUnlock()
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
c.mu.RUnlock()
|
||||
|
||||
packet := &UdpPacket{
|
||||
Buf: b,
|
||||
LocalAddr: c.localAddr,
|
||||
RemoteAddr: c.remoteAddr,
|
||||
}
|
||||
c.l.writeUdpPacket(packet)
|
||||
|
||||
c.mu.Lock()
|
||||
c.lastActive = time.Now()
|
||||
c.mu.Unlock()
|
||||
return len(b), nil
|
||||
}
|
||||
|
||||
func (c *FakeUdpConn) Close() error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if !c.closeFlag {
|
||||
c.closeFlag = true
|
||||
close(c.packets)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *FakeUdpConn) IsClosed() bool {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
return c.closeFlag
|
||||
}
|
||||
|
||||
func (c *FakeUdpConn) LocalAddr() net.Addr {
|
||||
return c.localAddr
|
||||
}
|
||||
|
||||
func (c *FakeUdpConn) RemoteAddr() net.Addr {
|
||||
return c.remoteAddr
|
||||
}
|
||||
|
||||
func (c *FakeUdpConn) SetDeadline(t time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *FakeUdpConn) SetReadDeadline(t time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *FakeUdpConn) SetWriteDeadline(t time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type UdpListener struct {
|
||||
net.Addr
|
||||
accept chan Conn
|
||||
writeCh chan *UdpPacket
|
||||
readConn net.Conn
|
||||
closeFlag bool
|
||||
|
||||
fakeConns map[string]*FakeUdpConn
|
||||
}
|
||||
|
||||
func ListenUDP(bindAddr string, bindPort int64) (l *UdpListener, err error) {
|
||||
udpAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", bindAddr, bindPort))
|
||||
if err != nil {
|
||||
return l, err
|
||||
}
|
||||
readConn, err := net.ListenUDP("udp", udpAddr)
|
||||
|
||||
l = &UdpListener{
|
||||
Addr: udpAddr,
|
||||
accept: make(chan Conn),
|
||||
writeCh: make(chan *UdpPacket, 1000),
|
||||
fakeConns: make(map[string]*FakeUdpConn),
|
||||
}
|
||||
|
||||
// for reading
|
||||
go func() {
|
||||
for {
|
||||
buf := pool.GetBuf(1450)
|
||||
n, remoteAddr, err := readConn.ReadFromUDP(buf)
|
||||
if err != nil {
|
||||
close(l.accept)
|
||||
close(l.writeCh)
|
||||
return
|
||||
}
|
||||
|
||||
fakeConn, exist := l.fakeConns[remoteAddr.String()]
|
||||
if !exist || fakeConn.IsClosed() {
|
||||
fakeConn = NewFakeUdpConn(l, l.Addr, remoteAddr)
|
||||
l.fakeConns[remoteAddr.String()] = fakeConn
|
||||
}
|
||||
fakeConn.putPacket(buf[:n])
|
||||
|
||||
l.accept <- fakeConn
|
||||
}
|
||||
}()
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
func (l *UdpListener) writeUdpPacket(packet *UdpPacket) {
|
||||
defer func() {
|
||||
if err := recover(); err != nil {
|
||||
}
|
||||
}()
|
||||
l.writeCh <- packet
|
||||
}
|
||||
|
||||
func (l *UdpListener) Accept() (Conn, error) {
|
||||
conn, ok := <-l.accept
|
||||
if !ok {
|
||||
return conn, fmt.Errorf("channel for udp listener closed")
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func (l *UdpListener) Close() error {
|
||||
if !l.closeFlag {
|
||||
l.closeFlag = true
|
||||
l.readConn.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
27
utils/net/udp_test.go
Normal file
27
utils/net/udp_test.go
Normal file
@@ -0,0 +1,27 @@
|
||||
package net
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestA(t *testing.T) {
|
||||
l, err := ListenUDP("0.0.0.0", 9000)
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
}
|
||||
for {
|
||||
c, _ := l.Accept()
|
||||
go func() {
|
||||
for {
|
||||
buf := make([]byte, 1450)
|
||||
n, err := c.Read(buf)
|
||||
if err != nil {
|
||||
fmt.Println(buf[:n])
|
||||
}
|
||||
|
||||
c.Write(buf[:n])
|
||||
}
|
||||
}()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user