github.com/codingeasygo/util@v0.0.0-20231206062002-1ce2f004b7d9/proxy/forward.go (about)

     1  package proxy
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"net"
     7  	"net/url"
     8  	"strings"
     9  	"sync"
    10  
    11  	"github.com/codingeasygo/util/proxy/socks"
    12  	"github.com/codingeasygo/util/xdebug"
    13  	"github.com/codingeasygo/util/xio"
    14  )
    15  
    16  type RouterPiperDialer struct {
    17  	Router string
    18  	Next   xio.PiperDialer
    19  }
    20  
    21  func (r *RouterPiperDialer) DialPiper(uri string, bufferSize int) (raw xio.Piper, err error) {
    22  	raw, err = r.Next.DialPiper(strings.Replace(r.Router, "${HOST}", uri, -1), bufferSize)
    23  	return
    24  }
    25  
    26  type Forward struct {
    27  	Name       string
    28  	BufferSize int
    29  	Dialer     xio.PiperDialer
    30  	forwardLck sync.RWMutex
    31  	forwardAll map[string][]interface{}
    32  }
    33  
    34  func NewForward(name string) (forward *Forward) {
    35  	forward = &Forward{
    36  		Name:       name,
    37  		BufferSize: 8 * 1024,
    38  		Dialer:     xio.PiperDialerF(xio.DialNetPiper),
    39  		forwardLck: sync.RWMutex{},
    40  		forwardAll: map[string][]interface{}{},
    41  	}
    42  	return
    43  }
    44  
    45  // StartForward will forward address to uri
    46  func (f *Forward) StartForward(name string, listen *url.URL, router string) (listener net.Listener, err error) {
    47  	f.forwardLck.Lock()
    48  	defer f.forwardLck.Unlock()
    49  	if f.forwardAll[name] != nil || len(name) < 1 {
    50  		err = fmt.Errorf("the name(%v) is already used", name)
    51  		WarnLog("Forward(%v) start forward by %v fail with %v", f.Name, listen, router, err)
    52  		return
    53  	}
    54  	switch listen.Scheme {
    55  	case "socks":
    56  		sp := socks.NewServer()
    57  		sp.BufferSize = f.BufferSize
    58  		sp.Dialer = &RouterPiperDialer{Router: router, Next: f.Dialer}
    59  		listener, err = sp.Start("tcp", listen.Host)
    60  		if err == nil {
    61  			f.forwardAll[name] = []interface{}{listen.Scheme, listener, listen}
    62  			InfoLog("Forward(%v) start socket forward on %v success by %v->%v", f.Name, listener.Addr(), listen, router)
    63  		}
    64  	case "proxy":
    65  		dialer := &RouterPiperDialer{Router: router, Next: f.Dialer}
    66  		sp := NewServer(dialer)
    67  		sp.SOCKS.BufferSize = f.BufferSize
    68  		sp.HTTP.BufferSize = f.BufferSize
    69  		listener, err = sp.Start("tcp", listen.Host)
    70  		if err == nil {
    71  			f.forwardAll[name] = []interface{}{listen.Scheme, listener, listen, router}
    72  			InfoLog("Forward(%v) start proxy forward on %v success by %v->%v", f.Name, listener.Addr(), listen, router)
    73  		}
    74  	default:
    75  		listener, err = net.Listen(listen.Scheme, listen.Host)
    76  		if err == nil {
    77  			f.forwardAll[name] = []interface{}{listen.Scheme, listener, listen, router}
    78  			go f.loopForward(listener, name, listen, router)
    79  			InfoLog("Forward(%v) start tcp forward on %v success by %v->%v", f.Name, listener.Addr(), listen, router)
    80  		}
    81  	}
    82  	return
    83  }
    84  
    85  // StopForward will forward address to uri
    86  func (f *Forward) StopForward(name string) (err error) {
    87  	InfoLog("Forward(%v) stop forward by name:%v", f.Name, name)
    88  	f.forwardLck.Lock()
    89  	forward := f.forwardAll[name]
    90  	delete(f.forwardAll, name)
    91  	f.forwardLck.Unlock()
    92  	if len(forward) > 0 {
    93  		err = forward[1].(io.Closer).Close()
    94  	}
    95  	return
    96  }
    97  
    98  func (f *Forward) loopForward(l net.Listener, name string, listen *url.URL, uri string) {
    99  	defer func() {
   100  		f.forwardLck.Lock()
   101  		delete(f.forwardAll, name)
   102  		f.forwardLck.Unlock()
   103  	}()
   104  	var err error
   105  	var piper xio.Piper
   106  	var conn net.Conn
   107  	InfoLog("Forward(%v) proxy forward(%v->%v) accept runner is starting", f.Name, l.Addr(), uri)
   108  	for {
   109  		conn, err = l.Accept()
   110  		if err != nil {
   111  			break
   112  		}
   113  		DebugLog("Forward(%v) accepting forward(%v->%v) connection from %v", f.Name, l.Addr(), uri, conn.RemoteAddr())
   114  		piper, err = f.Dialer.DialPiper(uri, f.BufferSize)
   115  		if err == nil {
   116  			DebugLog("Forward(%v) proxy forward(%v->%v) success", f.Name, l.Addr(), uri)
   117  			go f.procForward(l, name, piper, conn, uri)
   118  		} else {
   119  			WarnLog("Forward(%v) proxy forward(%v->%v) fail with %v", f.Name, l.Addr(), uri, err)
   120  			conn.Close()
   121  		}
   122  	}
   123  	l.Close()
   124  	InfoLog("Forward(%v) proxy forward(%v->%v) accept runner is stopped", f.Name, l.Addr(), uri)
   125  }
   126  
   127  func (f *Forward) procForward(l net.Listener, name string, piper xio.Piper, conn net.Conn, uri string) {
   128  	defer func() {
   129  		if perr := recover(); perr != nil {
   130  			ErrorLog("Forward(%v) transfer forward (%v->%v) is pance with %v, callstack is \n%v", f.Name, l.Addr(), uri, perr, xdebug.CallStack())
   131  		}
   132  	}()
   133  	piper.PipeConn(conn, uri)
   134  }
   135  
   136  // Stop will stop all
   137  func (f *Forward) Stop() (err error) {
   138  	InfoLog("Forward(%v) is closing", f.Name)
   139  	f.forwardLck.RLock()
   140  	for key, forward := range f.forwardAll {
   141  		forward[1].(net.Listener).Close()
   142  		InfoLog("Forward(%v) forwad %v is closed", f.Name, key)
   143  	}
   144  	f.forwardAll = map[string][]interface{}{}
   145  	f.forwardLck.RUnlock()
   146  	return
   147  }